"""Optional Claude reviews of a draft run."""

from __future__ import annotations

import anthropic
from fastapi import APIRouter, Depends, HTTPException
from psycopg import Connection
from pydantic import BaseModel, Field

from .. import keystore, server_settings
from ..auth import User, current_user, require_center
from ..configs import config_by_id
from ..db import Jsonb, fetchall, fetchone, get_conn
from ..retrieval import find_analogs
from ..services import claude
from ..settings import get_settings
from ..style import check as style_check
from .forecasts import _query_from_inputs
from .runs import _load_run

router = APIRouter(tags=["ai"])


class ReviewIn(BaseModel):
    run_id: int
    mode: str = Field(pattern="^(review_draft|draft_bottom_line|sanity_check)$")
    model: str | None = None


def _with_style(row: dict | None) -> dict | None:
    """Attach writing-style findings for Claude's suggested bottom line (computed on read)."""
    if row is not None:
        text = (row.get("output") or {}).get("suggested_bottom_line")
        row["style"] = style_check(text, "bottom_line") if text else []
    return row


class FeedbackIn(BaseModel):
    feedback: int = Field(ge=-1, le=1)
    note: str | None = None


def _month_spend(conn: Connection, center_id: int) -> float:
    row = fetchone(
        conn,
        "SELECT coalesce(sum(cost_usd), 0) AS s FROM ai_reviews WHERE center_id = %s AND created_at >= date_trunc('month', now())",
        (center_id,),
    )
    return float(row["s"])


@router.get("/ai/status")
def status(center_id: int | None = None, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    s = get_settings()
    out = {"enabled": claude.enabled(keystore.anthropic_key(conn)), "models": s.ai_allowed_models, "default_model": s.ai_default_model,
           "monthly_cap_usd": server_settings.ai_cap(conn)}
    if center_id is not None:
        require_center(user, center_id)
        out["month_spend_usd"] = _month_spend(conn, center_id)
    return out


@router.post("/ai/review")
def review(body: ReviewIn, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    s = get_settings()
    api_key = keystore.anthropic_key(conn)
    if not claude.enabled(api_key):
        raise HTTPException(503, "AI review is not configured: a site admin can add the API key on the Admin page.")
    model = body.model or s.ai_default_model
    if model not in s.ai_allowed_models:
        raise HTTPException(422, f"Model must be one of {s.ai_allowed_models}")
    run = _load_run(conn, body.run_id, user)
    if _month_spend(conn, run["center_id"]) >= server_settings.ai_cap(conn):
        raise HTTPException(429, "This center has reached its monthly AI review budget.")
    config = config_by_id(conn, run["config_id"])["config"]
    analogs = find_analogs(conn, run["center_id"], _query_from_inputs(run["inputs"], config), run["valid_date"], run["zone_id"], k=5)["items"]
    params = claude.build_request(run, config, run["center_name"], analogs, body.mode, model)
    try:
        result = claude.run_review(params, api_key)
    except claude.AIUnavailable as e:
        raise HTTPException(503, str(e)) from e
    except anthropic.RateLimitError as e:
        raise HTTPException(429, "Claude is rate limited right now; try again shortly.") from e
    except anthropic.APIStatusError as e:
        raise HTTPException(502, f"Claude API error ({e.status_code}).") from e
    except anthropic.APIConnectionError as e:
        raise HTTPException(502, "Could not reach the Claude API.") from e
    return _with_style(fetchone(
        conn,
        """INSERT INTO ai_reviews (center_id, tool_run_id, user_id, mode, model, prompt_version, retrieved_forecast_ids,
               output, stop_reason, fallback_model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
               cost_usd, error)
           VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING *""",
        (run["center_id"], run["id"], user.id, body.mode, model, claude.PROMPT_VERSION, [a["id"] for a in analogs],
         Jsonb(result["output"]) if result["output"] is not None else None, result["stop_reason"], result["fallback_model"],
         result["input_tokens"], result["output_tokens"], result["cache_read_tokens"], result["cache_write_tokens"],
         result["cost_usd"], result["error"]),
    ))


@router.get("/runs/{run_id}/reviews")
def list_reviews(run_id: int, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    _load_run(conn, run_id, user)
    return [_with_style(r) for r in fetchall(conn, "SELECT * FROM ai_reviews WHERE tool_run_id = %s ORDER BY created_at DESC", (run_id,))]


@router.post("/ai/reviews/{review_id}/feedback")
def feedback(review_id: int, body: FeedbackIn, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    row = fetchone(conn, "SELECT center_id FROM ai_reviews WHERE id = %s", (review_id,))
    if not row:
        raise HTTPException(404, "Review not found")
    require_center(user, row["center_id"])
    conn.execute("UPDATE ai_reviews SET feedback = %s, feedback_note = %s WHERE id = %s", (body.feedback, body.note, review_id))
    return {"ok": True}