"""
AI Report Agent — SSE Streaming Endpoint with ReAct Reflect Loop.

POST /api/report-agent
  Input:  { question: str, history: list[{role, content}] }
  Output: SSE stream of events

ReAct Agent Loop:
  1. THINK → decide what tools to call
  2. ACT  → execute tools (parallel)
  3. OBSERVE → collect results
  4. REFLECT → is data sufficient? quality ok?
  5. If sufficient → WRITE report
  6. If not → go back to step 1 with additional queries
  Max 3 reflect cycles to prevent infinite loops.
"""

import asyncio
import json
import logging
import re
import traceback
from typing import Any

from fastapi import APIRouter
from pydantic import BaseModel
from starlette.responses import StreamingResponse

from common.starrocks_connection import StarRocksConnection
from common.postgres_readonly import PostgresReadonly
from prompts.report_prompt import AGENT_PROMPT, WRITER_PROMPT

logger = logging.getLogger(__name__)
router = APIRouter()

# ─── Request Model ───────────────────────────────────────────────────

class ChatMessage(BaseModel):
    role: str
    content: str

class ReportAgentRequest(BaseModel):
    question: str
    history: list[ChatMessage] = []
    model: str = "codex/gpt-5.3-codex"


# ─── SQL Safety ──────────────────────────────────────────────────────

_FORBIDDEN = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|REPLACE|MERGE|GRANT|REVOKE)\b",
    re.IGNORECASE,
)

def _validate_sql(sql: str) -> str | None:
    sql_clean = sql.strip().rstrip(";")
    if not sql_clean.upper().startswith("SELECT"):
        return "Only SELECT queries allowed"
    if _FORBIDDEN.search(sql_clean):
        return "Forbidden SQL keyword"
    check = sql_clean.lower().replace("_text_embedding", "").replace("_embedding", "")
    if re.search(r'\b(embedding|vector)\b', check):
        return "Cannot query embedding/vector column"
    return None


# ─── LLM Call (Codex) ────────────────────────────────────────────────

async def _call_llm(system_prompt: str, user_content: str, model_id: str = "codex/gpt-5.3-codex") -> str:
    """Call Codex LLM and return raw text response."""
    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 = get_codex_access_token()
    if not access_token:
        raise RuntimeError("Codex auth not available. Run `codex login` first.")

    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


def _parse_json(raw: str) -> dict:
    """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)


# ─── Tool Executors ──────────────────────────────────────────────────

async def _exec_tool(tool_name: str, params: dict) -> dict[str, Any]:
    """Execute a single tool and return results."""

    if tool_name == "sql_langfuse":
        sql = params.get("sql", "")
        err = _validate_sql(sql)
        if err:
            return {"error": err, "data": []}
        db = StarRocksConnection()
        try:
            rows = await db.execute_query_async(sql)
            data = [_serialize_row(dict(r)) for r in rows]
            return {"data": data, "columns": list(rows[0].keys()) if rows else [], "row_count": len(data)}
        except Exception as e:
            return {"error": str(e)[:300], "data": []}

    elif tool_name == "sql_starrocks":
        sql = params.get("sql", "")
        err = _validate_sql(sql)
        if err:
            return {"error": err, "data": []}
        db = StarRocksConnection()
        try:
            rows = await db.execute_query_async(sql)
            data = [_serialize_row(dict(r)) for r in rows]
            return {"data": data, "columns": list(rows[0].keys()) if rows else [], "row_count": len(data)}
        except Exception as e:
            return {"error": str(e)[:300], "data": []}

    elif tool_name == "sql_postgres":
        sql = params.get("sql", "")
        err = _validate_sql(sql)
        if err:
            return {"error": err, "data": []}
        try:
            rows = await PostgresReadonly.execute_query_async(sql)
            data = [_serialize_row(dict(r)) for r in rows]
            return {"data": data, "columns": list(rows[0].keys()) if rows else [], "row_count": len(data)}
        except Exception as e:
            return {"error": str(e)[:300], "data": []}

    elif tool_name == "calculator":
        expr = params.get("expression", "")
        description = params.get("description", "")
        try:
            allowed = set("0123456789.+-*/() ,")
            if all(c in allowed for c in expr.replace(" ", "")):
                result = eval(expr)  # noqa: S307
                return {"expression": expr, "result": result, "description": description}
            else:
                return {"error": "Invalid expression (only math allowed)", "expression": expr}
        except Exception as e:
            return {"error": str(e), "expression": expr}

    else:
        return {"error": f"Unknown tool: {tool_name}"}


def _serialize_row(row: dict) -> dict:
    """Serialize row for JSON (handle Decimal, bytes, datetime, NaN)."""
    for key, val in row.items():
        if isinstance(val, (bytes, bytearray)):
            row[key] = val.hex()
        elif hasattr(val, "isoformat"):
            row[key] = val.isoformat()
        elif isinstance(val, float) and val != val:  # NaN
            row[key] = 0
        elif hasattr(val, "as_tuple"):  # Decimal
            row[key] = float(val)
    return row


async def _execute_tools_parallel(tools: list[dict]) -> list[dict]:
    """Execute multiple tools in parallel, return results."""
    tasks = [_exec_tool(t.get("name", ""), t.get("params", {})) for t in tools]
    return await asyncio.gather(*tasks, return_exceptions=True)


def _make_result_preview(result: dict) -> dict:
    """Create a compact preview of a tool result for SSE streaming."""
    preview = {}
    if isinstance(result, Exception):
        return {"error": str(result)[:200]}
    if "data" in result:
        data = result["data"]
        if isinstance(data, list):
            preview["row_count"] = len(data)
            preview["preview"] = data[:3] if data else []
        elif isinstance(data, dict):
            preview["keys"] = list(data.keys())[:10]
    if "summary" in result:
        preview["summary"] = result["summary"]
    if "error" in result:
        preview["error"] = result["error"]
    if "result" in result:
        preview["result"] = result["result"]
    return preview


def _summarize_results(all_results: dict[str, Any]) -> str:
    """Build a text summary of all collected tool results for the LLM."""
    parts = []
    for key, result in all_results.items():
        result_copy = dict(result) if isinstance(result, dict) else {"error": str(result)}
        if "data" in result_copy and isinstance(result_copy["data"], list):
            result_copy["data"] = result_copy["data"][:50]  # max 50 rows per tool
        parts.append(
            f"\n## Tool Result: {key}\n```json\n{json.dumps(result_copy, ensure_ascii=False, default=str, indent=2)}\n```"
        )
    return "\n".join(parts)


# ─── SSE Streaming Endpoint ─────────────────────────────────────────

MAX_REFLECT_CYCLES = 3

def _sse(data: dict) -> str:
    """Format SSE event."""
    return f"data: {json.dumps(data, ensure_ascii=False, default=str)}\n\n"


@router.post("/api/report-agent", summary="AI Report Agent (SSE Stream)")
async def report_agent_stream(req: ReportAgentRequest):
    """
    Generate an analytical report via SSE — ReAct agent loop with reflection:
      1. THINK → decide what data to gather
      2. ACT  → execute tools (parallel)
      3. OBSERVE → collect results
      4. REFLECT → is data sufficient?
      5. WRITE → generate report from real data
    """

    async def event_stream():
        try:
            logger.info("📝 Report Agent: %s", req.question[:100])
            all_tool_results: dict[str, Any] = {}
            tool_counter = 0
            cycle = 0

            # ── Initial THINK: Agent decides first batch of tools ────
            yield _sse({"type": "thinking", "step": "🧠 Đang phân tích yêu cầu..."})

            agent_input = f"User request: {req.question}"
            agent_raw = await _call_llm(AGENT_PROMPT, agent_input, req.model)
            agent_response = _parse_json(agent_raw)

            # ── ReAct Loop ───────────────────────────────────────────
            while cycle < MAX_REFLECT_CYCLES:
                action = agent_response.get("action", "execute")
                thinking = agent_response.get("thinking", "")

                # Get tools to execute from either "execute" or "reflect with next_tools"
                tools_to_run = []
                if action == "execute":
                    tools_to_run = agent_response.get("tools", [])
                elif action == "reflect":
                    data_sufficient = agent_response.get("data_sufficient", False)
                    missing = agent_response.get("missing", [])
                    quality_issues = agent_response.get("quality_issues", [])

                    yield _sse({
                        "type": "reflect",
                        "cycle": cycle + 1,
                        "thinking": thinking,
                        "data_sufficient": data_sufficient,
                        "missing": missing,
                        "quality_issues": quality_issues,
                    })

                    if data_sufficient and not missing:
                        # Data is sufficient — proceed to write
                        logger.info("✅ Reflect cycle %d: Data sufficient, proceeding to write", cycle + 1)
                        break

                    # Not sufficient — get additional tools
                    tools_to_run = agent_response.get("next_tools", [])
                    if not tools_to_run:
                        logger.info("⚠️ Reflect: not sufficient but no next_tools, forcing write")
                        break

                elif action == "write_report":
                    break

                if not tools_to_run:
                    break

                # ── ACT: Execute tools in parallel ────────────────────
                yield _sse({
                    "type": "plan_ready",
                    "thinking": thinking,
                    "tools_count": len(tools_to_run),
                    "cycle": cycle + 1,
                })

                logger.info("🔄 Cycle %d: executing %d tools", cycle + 1, len(tools_to_run))

                # Stream tool_call events
                for i, tool_spec in enumerate(tools_to_run):
                    yield _sse({
                        "type": "tool_call",
                        "index": tool_counter + i,
                        "tool": tool_spec.get("name", ""),
                        "params": tool_spec.get("params", {}),
                        "purpose": tool_spec.get("purpose", ""),
                    })

                # Execute all tools in parallel
                results = await _execute_tools_parallel(tools_to_run)

                # ── OBSERVE: Collect and stream results ───────────────
                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

                    yield _sse({
                        "type": "tool_result",
                        "index": tool_counter + i,
                        "tool": tool_name,
                        "success": "error" not in result,
                        "preview": _make_result_preview(result),
                    })

                tool_counter += len(tools_to_run)
                cycle += 1

                # ── REFLECT: Ask LLM to assess data sufficiency ──────
                yield _sse({"type": "thinking", "step": f"🔍 Đánh giá dữ liệu (vòng {cycle}/{MAX_REFLECT_CYCLES})..."})

                results_summary = _summarize_results(all_tool_results)
                reflect_input = (
                    f"User request: {req.question}\n\n"
                    f"## DATA COLLECTED SO FAR (cycle {cycle}):\n{results_summary}\n\n"
                    f"## INSTRUCTION:\n"
                    f"Review the data above. Respond with action='reflect'.\n"
                    f"Assess: Is this data SUFFICIENT for a comprehensive report? "
                    f"Any GAPS or QUALITY ISSUES? Need additional queries?\n"
                    f"If data is sufficient, set data_sufficient=true and missing=[].\n"
                    f"If not, provide next_tools with additional queries.\n"
                    f"RESPOND WITH RAW JSON ONLY."
                )

                reflect_raw = await _call_llm(AGENT_PROMPT, reflect_input, req.model)
                agent_response = _parse_json(reflect_raw)

                logger.info("🔍 Reflect cycle %d: action=%s", cycle, agent_response.get("action"))

            # ── WRITE: Generate final report ─────────────────────────
            logger.info("🔧 Total tools executed: %d across %d cycles", tool_counter, cycle)

            yield _sse({"type": "thinking", "step": "✍️ Đang viết báo cáo từ dữ liệu thật..."})

            data_context = _summarize_results(all_tool_results)
            writer_input = (
                f"User request: {req.question}\n\n"
                f"## REAL DATA FROM TOOLS ({tool_counter} queries, {cycle} analysis cycles):\n{data_context}\n\n"
                f"Write the full report using ONLY the data above. RESPOND WITH RAW JSON ONLY."
            )

            report_raw = await _call_llm(WRITER_PROMPT, writer_input, req.model)
            report = _parse_json(report_raw)

            # Stream sections progressively
            sections = report.get("sections", [])
            for i, section in enumerate(sections):
                yield _sse({
                    "type": "section_ready",
                    "index": i,
                    "section": section,
                })
                await asyncio.sleep(0.1)

            # Full report
            yield _sse({"type": "report_complete", "report": report})

            num_sec = len(sections)
            num_chart = sum(1 for s in sections if s.get("chart"))
            num_table = sum(1 for s in sections if s.get("table"))
            logger.info(
                "✅ Report complete: '%s' — %d sections, %d charts, %d tables, %d cycles",
                report.get("title", "?"), num_sec, num_chart, num_table, cycle,
            )

            yield _sse({"type": "done"})

        except json.JSONDecodeError as e:
            logger.error("Report Agent JSON error: %s", e)
            yield _sse({"type": "error", "message": f"AI response was not valid JSON: {str(e)[:200]}"})
        except RuntimeError as e:
            logger.error("Report Agent runtime error: %s", e)
            yield _sse({"type": "error", "message": str(e)[:300]})
        except Exception as e:
            logger.error("Report Agent error: %s\n%s", e, traceback.format_exc())
            yield _sse({"type": "error", "message": str(e)[:300]})

    return StreamingResponse(event_stream(), media_type="text/event-stream")


# ─── Health Check ────────────────────────────────────────────────────

@router.get("/api/report-agent/status", summary="Check Report Agent availability")
async def report_agent_status():
    """Check if Codex is available."""
    from common.codex_auth import get_codex_access_token

    return {
        "codex_available": bool(get_codex_access_token()),
        "data_source": "StarRocks (analytic.chatbot_rsa_trace_event_detail)",
        "agent_type": "ReAct with Reflect Loop (max 3 cycles)",
    }
