"""Polleo AI — conversational analytics service (agentic).

Instead of one-shot text-to-SQL, Claude runs an AGENTIC LOOP with two tools:

  run_sql(sql)        — execute one read-only SELECT and get rows back. Claude
                        calls this as many times as it needs to gather/verify
                        data and compute an answer.
  respond(answer,     — Claude calls this exactly once to deliver the final
          chart)        answer (+ optional chart over its last query's columns).

The loop: POST messages+tools -> if Claude emits run_sql, execute it read-only
and feed the rows back as a tool_result, then POST again -> repeat until Claude
calls respond (or we hit MAX_TURNS). This is what makes it an analyst rather
than a single query: it can break a question into steps, check assumptions, and
reason over the results.

Safety is unchanged and layered: validate_sql (SELECT-only) + the read-only
Postgres role + a rolled-back read-only transaction + statement_timeout (see
repositories/polleo_ai_repo.py).

Sync on purpose — the whole backend is sync; FastAPI runs this endpoint in a
threadpool so the blocking httpx calls don't stall the event loop. The static
system prompt is sent with cache_control so prompt caching amortises it across
the (now multiple) calls per question.
"""
from __future__ import annotations

import json
import logging
import re
from collections.abc import Iterator
from typing import Any

import httpx

from backend.config import settings
from backend.repositories.polleo_ai_repo import (
    SqlExecutionError,
    SqlTimeoutError,
    run_readonly,
)
from backend.schemas.polleo_ai import ChartConfig, ChatResponse, ChatTurn
from backend.services.polleo_ai_system_prompt import SYSTEM_PROMPT

logger = logging.getLogger(__name__)

ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"
MAX_TOKENS = 4096
# Keep at most this many prior turns (≈ 10 user/assistant pairs) for token thrift.
MAX_HISTORY_TURNS = 20
# Hard cap on Claude round-trips per question (bounds cost/latency on the loop).
# On the final turn we FORCE the respond tool so the user always gets a real
# answer with whatever was gathered, instead of an "out of steps" fallback.
MAX_TURNS = 8
# Rows handed back to Claude per query (full set, up to 500, is kept for the UI).
TOOL_ROW_CAP = 60

FORBIDDEN_KEYWORDS = {
    "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "CREATE", "TRUNCATE",
    "GRANT", "REVOKE", "EXECUTE", "EXEC", "COPY", "MERGE", "CALL", "VACUUM",
    "REINDEX", "REFRESH", "COMMENT",
}
FORBIDDEN_FUNCTIONS = (
    "PG_SLEEP", "PG_READ_FILE", "PG_READ_BINARY_FILE", "PG_LS_DIR",
    "LO_IMPORT", "LO_EXPORT", "DBLINK", "PG_TERMINATE_BACKEND",
    "PG_CANCEL_BACKEND", "PG_STAT_FILE",
)

_WORD_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")

# ── Tool definitions sent to the Anthropic API ────────────────────────────
RUN_SQL_TOOL = {
    "name": "run_sql",
    "description": (
        "Run ONE read-only SELECT / WITH...SELECT query against the Postgres "
        "database and get the rows back as JSON. Call it as many times as you "
        "need to gather and verify data before answering."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "sql": {
                "type": "string",
                "description": "A single SELECT or WITH...SELECT statement (no writes, no ';' chaining).",
            }
        },
        "required": ["sql"],
    },
}
RESPOND_TOOL = {
    "name": "respond",
    "description": (
        "Deliver your FINAL answer to the user. Call this exactly once, after "
        "you have gathered everything you need."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "answer": {
                "type": "string",
                "description": "The full text answer, in the same language the user used.",
            },
            "chart": {
                "type": ["object", "null"],
                "description": "Optional chart over the columns of your MOST RECENT run_sql result, else null.",
                "properties": {
                    "type": {"type": "string", "enum": ["bar", "line", "area", "pie"]},
                    "x_key": {"type": "string"},
                    "y_keys": {"type": "array", "items": {"type": "string"}},
                    "title": {"type": "string"},
                },
            },
        },
        "required": ["answer"],
    },
}


class PolleoAiNotConfigured(Exception):
    """ANTHROPIC_API_KEY is not set."""


class ClaudeTimeoutError(Exception):
    """The Anthropic API call timed out."""


class ClaudeApiError(Exception):
    """The Anthropic API returned an error or unexpected payload."""


def is_configured() -> bool:
    return bool(settings.ANTHROPIC_API_KEY)


def validate_sql(sql: str) -> bool:
    """Allow only a single SELECT / WITH...SELECT statement."""
    if not sql or not sql.strip():
        return False
    s = sql.strip().rstrip(";").strip()
    if not s or ";" in s:  # no multiple statements
        return False
    upper = s.upper()
    first = upper.split(None, 1)[0] if upper.split() else ""
    if first not in ("SELECT", "WITH"):
        return False
    tokens = {m.group(0).upper() for m in _WORD_RE.finditer(s)}
    if tokens & FORBIDDEN_KEYWORDS:
        return False
    if any(fn in upper for fn in FORBIDDEN_FUNCTIONS):
        return False
    return True


def _build_messages(message: str, history: list[ChatTurn]) -> list[dict[str, Any]]:
    """Prior text turns (trimmed, leading-assistant dropped) + the new user message."""
    trimmed = history[-MAX_HISTORY_TURNS:] if history else []
    msgs: list[dict[str, Any]] = [{"role": t.role, "content": t.content} for t in trimmed]
    while msgs and msgs[0]["role"] != "user":
        msgs.pop(0)
    msgs.append({"role": "user", "content": message})
    return msgs


def _call_claude(messages: list[dict[str, Any]], *, force_respond: bool = False) -> dict[str, Any]:
    """POST to the Anthropic API (with tools) and return the parsed JSON response.

    force_respond=True pins tool_choice to `respond` so Claude must deliver a
    final answer this turn (used on the last allowed turn)."""
    headers = {
        "x-api-key": settings.ANTHROPIC_API_KEY,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }
    payload: dict[str, Any] = {
        "model": settings.ANTHROPIC_MODEL,
        "max_tokens": MAX_TOKENS,
        "system": [
            {"type": "text", "text": SYSTEM_PROMPT, "cache_control": {"type": "ephemeral"}}
        ],
        "tools": [RUN_SQL_TOOL, RESPOND_TOOL],
        "messages": messages,
    }
    if force_respond:
        payload["tool_choice"] = {"type": "tool", "name": "respond"}
    try:
        # 60s read budget per call — a complex agentic turn (large system prompt
        # + tool reasoning) can run past 30s. The loop makes several of these.
        with httpx.Client(timeout=httpx.Timeout(60.0, connect=10.0)) as client:
            resp = client.post(ANTHROPIC_URL, headers=headers, json=payload)
            resp.raise_for_status()
            return resp.json()
    except httpx.TimeoutException as exc:
        raise ClaudeTimeoutError(str(exc)) from exc
    except httpx.HTTPStatusError as exc:
        detail = exc.response.text[:500] if exc.response is not None else str(exc)
        logger.warning("Anthropic API error %s: %s", exc.response.status_code, detail)
        raise ClaudeApiError(f"Anthropic API returned {exc.response.status_code}") from exc
    except httpx.HTTPError as exc:
        raise ClaudeApiError(str(exc)) from exc


def _coerce_chart(chart_raw: Any) -> ChartConfig | None:
    if not isinstance(chart_raw, dict):
        return None
    try:
        return ChartConfig(**chart_raw)
    except Exception:  # noqa: BLE001 — a bad chart hint must never break the answer
        return None


def _text_of(blocks: list[dict[str, Any]]) -> str:
    return "".join(
        b.get("text", "") for b in blocks if isinstance(b, dict) and b.get("type") == "text"
    ).strip()


def _run_sql_for_tool(sql: str) -> tuple[str, list[dict[str, Any]] | None]:
    """Execute a run_sql tool call. Returns (tool_result_content, rows_or_None).
    rows is None on validation/exec failure (so it doesn't become the shown data)."""
    sql = (sql or "").strip()
    if not validate_sql(sql):
        return ("ERROR: only a single read-only SELECT/WITH query is allowed. Rewrite and retry.", None)
    try:
        rows = run_readonly(sql)
    except SqlTimeoutError:
        return ("ERROR: query timed out (>5s). Narrow the data range and retry.", None)
    except SqlExecutionError as exc:
        return (f"ERROR: {exc}", None)

    shown = rows[:TOOL_ROW_CAP]
    content = json.dumps(
        {"row_count": len(rows), "truncated": len(rows) > len(shown), "rows": shown},
        default=str,           # Decimals/dates -> str for Claude's reasoning
        ensure_ascii=False,
    )
    return (content, rows)


def run_agentic(message: str, history: list[ChatTurn]) -> Iterator[dict[str, Any]]:
    """Agentic generator. Yields progress events as it works, then a final event:

      {"type": "step", "tool": "run_sql", "sql": str, "ok": bool,
       "row_count": int (if ok), "error": str (if not ok)}
      {"type": "final", "response": ChatResponse}

    Raises PolleoAiNotConfigured / ClaudeTimeoutError / ClaudeApiError mid-iteration;
    callers (the SSE endpoint) catch these and emit an error event.
    """
    if not is_configured():
        raise PolleoAiNotConfigured()

    messages = _build_messages(message, history)
    executed_sql: list[str] = []   # all successful queries (for transparency)
    last_rows: list[dict[str, Any]] | None = None
    answer: str | None = None
    chart: ChartConfig | None = None

    for turn in range(MAX_TURNS):
        # On the last allowed turn, force a final answer with what we have.
        resp = _call_claude(messages, force_respond=(turn == MAX_TURNS - 1))
        blocks = resp.get("content") or []
        tool_uses = [b for b in blocks if isinstance(b, dict) and b.get("type") == "tool_use"]

        if not tool_uses:
            # Claude answered in plain text without using respond() — accept it.
            answer = _text_of(blocks)
            break

        # Echo the assistant's tool_use turn back, then answer each tool call.
        messages.append({"role": "assistant", "content": blocks})
        tool_results: list[dict[str, Any]] = []
        done = False
        for tu in tool_uses:
            name, tid, inp = tu.get("name"), tu.get("id"), (tu.get("input") or {})
            if name == "run_sql":
                sql = inp.get("sql", "").strip()
                content, rows = _run_sql_for_tool(sql)
                if rows is not None:
                    executed_sql.append(sql)
                    last_rows = rows
                    yield {"type": "step", "tool": "run_sql", "sql": sql, "ok": True, "row_count": len(rows)}
                else:
                    yield {"type": "step", "tool": "run_sql", "sql": sql, "ok": False, "error": content}
                tool_results.append({"type": "tool_result", "tool_use_id": tid, "content": content})
            elif name == "respond":
                answer = str(inp.get("answer") or "").strip()
                chart = _coerce_chart(inp.get("chart"))
                done = True
                tool_results.append({"type": "tool_result", "tool_use_id": tid, "content": "OK"})
            else:
                tool_results.append(
                    {"type": "tool_result", "tool_use_id": tid, "content": "ERROR: unknown tool", "is_error": True}
                )

        messages.append({"role": "user", "content": tool_results})
        if done:
            break

    if not answer:
        answer = "Nisam uspio doći do odgovora u zadanom broju koraka. Pokušaj suziti ili preformulirati pitanje."

    # Show the last query's rows as the data table; chart only makes sense with data.
    data = last_rows or None
    if data is None:
        chart = None
    sql_executed = "\n\n".join(executed_sql) if executed_sql else None

    yield {
        "type": "final",
        "response": ChatResponse(message=answer, data=data, chart=chart, sql_executed=sql_executed),
    }


def chat(message: str, history: list[ChatTurn]) -> ChatResponse:
    """Non-streaming entry point: drain the agentic generator, return the result.

    Raises PolleoAiNotConfigured / ClaudeTimeoutError / ClaudeApiError for the
    router to translate into HTTP responses.
    """
    final: ChatResponse | None = None
    for ev in run_agentic(message, history):
        if ev.get("type") == "final":
            final = ev["response"]
    if final is None:  # defensive — generator always yields a final
        return ChatResponse(message="Nisam uspio doći do odgovora.", data=None, chart=None, sql_executed=None)
    return final