"""
SQL Trace Agent — Human-in-the-Loop ReAct loop for AI SQL analysis.

Flow: User asks → AI generates SQL → User approves → SQL runs → Results shown
If SQL errors → AI fixes → User approves again → retry
"""

import asyncio
import json
import logging
import re
import time
import traceback
from typing import Any, AsyncGenerator

import httpx

from agent.sql_agent.session_manager import (
    create_session,
    ensure_cycle,
    public_session,
    sse,
    utc_now_iso,
    create_approval_event,
    get_approval_action,
    cleanup_approval,
)
from agent.sql_agent.tools import (
    execute_tools_parallel,
    make_result_preview,
    summarize_results,
)
from agent.sql_agent.persistence import persist_session_to_db
from agent.sql_agent.prompts.sql_trace_prompt import SQL_TRACE_AGENT_PROMPT

logger = logging.getLogger(__name__)


# ─── LLM Call ────────────────────────────────────────────────────────

async def call_llm(
    system_prompt: str,
    user_content: str,
    model_id: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
) -> str:
    """Call LLM (Codex or OpenAI) and return raw text response."""
    if model_id.startswith("openai/"):
        actual_model = model_id.replace("openai/", "")
        from common.llm_factory import LLMFactory

        factory = LLMFactory()

        if not openai_key:
            raise RuntimeError(
                "OpenAI API key khong ton tai. Them Token trong Cai Dat (Sidebar)."
            )

        model = factory.get_model(
            model_name=actual_model, streaming=False, api_key=openai_key
        )
        from langchain_core.messages import HumanMessage, SystemMessage

        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_content),
        ]
        result = await model._agenerate(messages)
        return result.generations[0].message.content

    else:
        from common.codex_auth import (
            get_codex_access_token,
            get_codex_account_id,
            get_refresh_token,
        )
        from common.codex_chat_model import CodexChatModel
        from langchain_core.messages import HumanMessage, SystemMessage

        access_token = codex_token or get_codex_access_token()
        if not access_token:
            raise RuntimeError(
                "Codex auth not available. Them Token trong Cai Dat (Sidebar)."
            )

        actual_model = (
            model_id.replace("codex/", "") if model_id.startswith("codex/") else model_id
        )

        model = CodexChatModel(
            model_name=actual_model,
            access_token=access_token,
            account_id=get_codex_account_id() or "",
            refresh_token=get_refresh_token(),
            streaming=False,
        )

        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_content),
        ]
        result = await model._agenerate(messages)
        return result.generations[0].message.content


# ─── JSON Parsing ────────────────────────────────────────────────────

def parse_json(raw: str) -> dict[str, Any]:
    """Parse JSON from LLM response, stripping markdown fences."""
    json_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", raw)
    json_str = json_match.group(1).strip() if json_match else raw.strip()
    return json.loads(json_str)


# ─── Error Humanizer ─────────────────────────────────────────────────

def humanize_llm_error(exc: Exception) -> str:
    """Convert LLM exceptions to human-readable Vietnamese messages."""
    if isinstance(exc, httpx.HTTPStatusError):
        response = exc.response
        status = response.status_code if response else "unknown"
        body = ""
        if response is not None:
            try:
                body = response.text[:300]
            except Exception:
                body = str(exc)

        lowered = body.lower()
        if "no route to host" in lowered:
            return (
                f"Codex backend is unreachable right now (HTTP {status}: No route to host). "
                "The SQL trace service cannot connect to the upstream model."
            )
        if status == 503:
            return "Codex backend is temporarily unavailable (HTTP 503). Please retry in a moment."
        return f"Codex API error (HTTP {status}). {body[:180]}".strip()

    if isinstance(exc, httpx.RequestError):
        return f"Network error while calling Codex backend: {exc}"

    return str(exc)


# ─── Main Agent Loop (Human-in-the-Loop) ─────────────────────────────

async def run_trace_agent(
    question: str,
    model: str,
    max_cycles: int,
    conversation_id: str | None = None,
    codex_token: str | None = None,
    openai_key: str | None = None,
) -> AsyncGenerator[str, None]:
    """
    Run the SQL Trace agent with Human-in-the-Loop approval.

    Flow:
    1. AI generates SQL from user question
    2. SSE: pending_approval (with SQL) → frontend shows Accept/Reject
    3. Wait for user approval
    4. If accepted → run SQL → show results
    5. If SQL error → AI fixes SQL → goto step 2
    6. If data sufficient → DONE
    """
    start_time = time.time()
    session = create_session(question, model, max_cycles, conversation_id)
    session_id = session["id"]
    tool_counter = 0
    all_tool_results: dict[str, Any] = {}

    try:
        yield sse(
            session,
            {
                "type": "session_started",
                "question": question,
                "model": model,
                "created_at": session["created_at"],
                "max_cycles": max_cycles,
            },
        )

        yield sse(
            session,
            {"type": "thinking", "step": "Dang phan tich cau hoi..."},
        )

        # ── THINK: Generate SQL ──
        think_input = (
            f"User question: {question}\n\n"
            "Generate the SQL queries to answer this question.\n"
            "Respond with raw JSON only."
        )
        agent_raw = await call_llm(
            SQL_TRACE_AGENT_PROMPT,
            think_input,
            model,
            codex_token=codex_token,
            openai_key=openai_key,
        )
        agent_response = parse_json(agent_raw)

        # ── Main HITL Loop ──
        cycle_no = 0
        while cycle_no < max_cycles:
            action = agent_response.get("action", "execute")
            thinking = agent_response.get("thinking", "")
            tools_to_run = agent_response.get("tools", [])

            # Reflect with no more tools → done
            if action == "reflect":
                data_sufficient = agent_response.get("data_sufficient", False)
                missing = agent_response.get("missing", [])
                tools_to_run = agent_response.get("next_tools", [])

                if data_sufficient or not tools_to_run:
                    summary_text = agent_response.get("summary", "")
                    yield sse(session, {
                        "type": "reflect",
                        "thinking": thinking,
                        "data_sufficient": data_sufficient,
                        "missing": missing,
                        "summary": summary_text,
                    })
                    session["ai_summary"] = summary_text
                    break

            if not tools_to_run:
                yield sse(session, {
                    "type": "reflect",
                    "thinking": thinking,
                    "data_sufficient": True,
                    "missing": [],
                })
                break

            cycle_no += 1
            cycle_obj = ensure_cycle(session, cycle_no)
            cycle_obj["thinking"] = thinking
            cycle_obj["status"] = "pending_approval"

            # ── Prepare tool info for approval ──
            tool_info = []
            for t in tools_to_run:
                tool_info.append({
                    "name": t.get("name", ""),
                    "params": t.get("params", {}),
                    "purpose": t.get("purpose", ""),
                })
                cycle_obj["tool_calls"].append({
                    "index": tool_counter + len(tool_info) - 1,
                    "tool": t.get("name", ""),
                    "params": t.get("params", {}),
                    "purpose": t.get("purpose", ""),
                })

            # ── SSE: pending_approval → Wait for user ──
            approval_event = create_approval_event(session_id)

            yield sse(session, {
                "type": "pending_approval",
                "cycle": cycle_no,
                "thinking": thinking,
                "tools": tool_info,
            })

            # Wait for user to click Accept or Reject (timeout 5 min)
            try:
                await asyncio.wait_for(approval_event.wait(), timeout=300)
            except asyncio.TimeoutError:
                yield sse(session, {
                    "type": "error",
                    "message": "Approval timeout (5 min). Session closed.",
                })
                session["status"] = "timeout"
                break

            user_action = get_approval_action(session_id)
            cleanup_approval(session_id)

            if user_action == "reject":
                cycle_obj["status"] = "rejected"
                yield sse(session, {
                    "type": "rejected",
                    "cycle": cycle_no,
                })
                break

            # ── User accepted → Execute tools ──
            cycle_obj["status"] = "running"
            yield sse(session, {
                "type": "executing",
                "cycle": cycle_no,
                "tools_count": len(tools_to_run),
            })

            results = await execute_tools_parallel(tools_to_run)

            # ── Process results ──
            has_error = False
            for i, (tool_spec, result) in enumerate(zip(tools_to_run, results)):
                tool_name = tool_spec.get("name", "unknown")
                if isinstance(result, Exception):
                    result = {"error": str(result)[:300], "data": []}

                key = f"{tool_name}_{tool_counter + i}"
                all_tool_results[key] = result

                success = "error" not in result
                session["total_queries"] += 1
                if success:
                    session["successful_queries"] += 1
                else:
                    session["failed_queries"] += 1
                    has_error = True

                preview = make_result_preview(result)
                stored_result = {
                    "index": tool_counter + i,
                    "tool": tool_name,
                    "success": success,
                    "purpose": tool_spec.get("purpose", ""),
                    "columns": preview.get("columns", []),
                    "row_count": preview.get("row_count", 0),
                    "preview_rows": preview.get("preview_rows", []),
                    "error": preview.get("error"),
                    "result": preview.get("result"),
                }
                cycle_obj["tool_results"].append(stored_result)

                yield sse(session, {
                    "type": "tool_result",
                    "cycle": cycle_no,
                    "index": tool_counter + i,
                    "tool": tool_name,
                    "success": success,
                    "preview": preview,
                })

            tool_counter += len(tools_to_run)
            cycle_obj["status"] = "done"

            # ── REFLECT: Ask LLM if more queries needed ──
            new_keys = [
                f"{t.get('name', 'unknown')}_{tool_counter - len(tools_to_run) + i}"
                for i, t in enumerate(tools_to_run)
            ]
            new_results = {k: v for k, v in all_tool_results.items() if k in new_keys}
            prev_keys = [k for k in all_tool_results if k not in new_keys]

            prev_summary = ""
            if prev_keys:
                prev_summary = f"### Previously collected ({len(prev_keys)} queries): {', '.join(prev_keys)}\n"

            new_summary = summarize_results(new_results)

            error_hint = ""
            if has_error:
                error_hint = "\nSome queries had errors. Fix the SQL and provide corrected tools.\n"

            reflect_input = (
                f"User question: {question}\n\n"
                f"## DATA COLLECTED (cycle {cycle_no}/{max_cycles}):\n"
                f"{prev_summary}"
                f"### New results this cycle:\n{new_summary}\n"
                f"{error_hint}\n"
                "## INSTRUCTION:\n"
                "Assess whether the data is sufficient to answer the question.\n"
                "If SQL had errors, fix them and provide corrected tools.\n"
                "If more data needed, provide next_tools.\n"
                "Respond with raw JSON only."
            )

            reflect_raw = await call_llm(
                SQL_TRACE_AGENT_PROMPT,
                reflect_input,
                model,
                codex_token=codex_token,
                openai_key=openai_key,
            )
            agent_response = parse_json(reflect_raw)

        # ── COMPLETE ──
        generation_time_ms = int((time.time() - start_time) * 1000)
        session["status"] = "completed"
        session["completed_at"] = utc_now_iso()
        session["summary"] = {
            "generation_time_ms": generation_time_ms,
            "cycles_completed": cycle_no,
            "total_queries": session["total_queries"],
            "successful_queries": session["successful_queries"],
            "failed_queries": session["failed_queries"],
        }

        persist_session_to_db(session)

        yield sse(session, {
            "type": "complete",
            "conversation_id": session["conversation_id"],
            "ai_summary": session.get("ai_summary", ""),
            **session["summary"],
        })
        yield sse(session, {"type": "done"})

    except Exception as e:
        session["status"] = "error"
        session["completed_at"] = utc_now_iso()
        session["summary"] = {
            "generation_time_ms": int((time.time() - start_time) * 1000),
            "cycles_completed": cycle_no if "cycle_no" in dir() else 0,
            "total_queries": session["total_queries"],
            "successful_queries": session["successful_queries"],
            "failed_queries": session["failed_queries"],
        }
        persist_session_to_db(session)
        friendly_error = humanize_llm_error(e)
        logger.error("AI SQL Trace error: %s\n%s", e, traceback.format_exc())
        yield sse(session, {"type": "error", "message": friendly_error[:500]})
    finally:
        cleanup_approval(session_id)
