"""
AI Report HTML Generator — SSE Streaming Endpoint.

POST /api/report-html
  Input:  { question: str }
  Output: SSE stream → final event contains body HTML

Same ReAct data gathering loop as report_agent_route.py,
but Writer outputs raw HTML body instead of JSON.
"""

import asyncio
import json
import logging
import re
import traceback
import time
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_html_prompt import HTML_AGENT_PROMPT, HTML_WRITER_PROMPT

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


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

class ReportHTMLRequest(BaseModel):
    question: str
    model: str = "codex/gpt-5.3-codex"
    parent_report_id: int | None = None


# ─── SQL Safety (same as report_agent_route) ─────────────────────────

_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:
    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:
    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)


def _extract_html(raw: str) -> str:
    """Extract HTML from LLM response, stripping markdown fences if present."""
    html_match = re.search(r'```(?:html)?\s*([\s\S]*?)```', raw)
    if html_match:
        return html_match.group(1).strip()
    return raw.strip()


# ─── Tool Executors (same as report_agent_route) ────────────────────

async def _exec_tool(tool_name: str, params: dict) -> dict[str, Any]:
    if tool_name in ("sql_langfuse", "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", "expression": expr}
        except Exception as e:
            return {"error": str(e), "expression": expr}

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


def _serialize_row(row: dict) -> dict:
    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:
            row[key] = 0
        elif hasattr(val, "as_tuple"):
            row[key] = float(val)
    return row


async def _execute_tools_parallel(tools: list[dict]) -> list[dict]:
    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:
    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 []
    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:
    """Compress tool results into concise format optimized for token efficiency.
    
    Strategy inspired by DeepAgents argument truncation:
    - Small results (≤20 rows): tab-separated table
    - Large results (>20 rows): auto-compute summary statistics
    - Numeric columns: min, max, avg, sum
    - Text columns: unique count + top values
    """
    parts = []
    for key, result in all_results.items():
        if isinstance(result, Exception):
            parts.append(f"\n## {key}\nERROR: {str(result)[:200]}")
            continue
        
        result_dict = dict(result) if isinstance(result, dict) else {"error": str(result)}
        
        if "error" in result_dict:
            parts.append(f"\n## {key}\nERROR: {result_dict['error']}")
            continue
        
        if "result" in result_dict:
            parts.append(f"\n## {key}\nResult: {result_dict['result']}")
            continue
        
        data = result_dict.get("data", [])
        if not isinstance(data, list) or not data:
            parts.append(f"\n## {key}\n(no data)")
            continue
        
        columns = result_dict.get("columns", list(data[0].keys()) if data else [])
        row_count = len(data)
        
        if row_count <= 20:
            # Small result: show full tab-separated table (perfect for charts)
            lines = [f"\n## {key} ({row_count} rows)"]
            lines.append("\t".join(str(c) for c in columns))
            for row in data:
                cells = []
                for col in columns:
                    val = row.get(col, "")
                    s = str(val) if val is not None else ""
                    if len(s) > 60:
                        s = s[:57] + "..."
                    cells.append(s)
                lines.append("\t".join(cells))
            parts.append("\n".join(lines))
        else:
            # Large result: auto-aggregate to save tokens
            lines = [f"\n## {key} ({row_count} rows → auto-summarized)"]
            
            # Show first 10 rows for chart data
            lines.append("### Sample (first 10):")
            lines.append("\t".join(str(c) for c in columns))
            for row in data[:10]:
                cells = []
                for col in columns:
                    val = row.get(col, "")
                    s = str(val) if val is not None else ""
                    if len(s) > 40:
                        s = s[:37] + "..."
                    cells.append(s)
                lines.append("\t".join(cells))
            
            # Compute stats for numeric columns
            lines.append("### Statistics:")
            for col in columns:
                values = [row.get(col) for row in data if row.get(col) is not None]
                if not values:
                    continue
                # Check if numeric
                try:
                    nums = [float(v) for v in values if v != "" and v is not None]
                    if nums and len(nums) > len(values) * 0.5:
                        lines.append(
                            f"  {col}: min={min(nums):.2f}, max={max(nums):.2f}, "
                            f"avg={sum(nums)/len(nums):.2f}, sum={sum(nums):.2f}, count={len(nums)}"
                        )
                except (ValueError, TypeError):
                    # Text column: show unique count + top 5
                    unique = list(set(str(v) for v in values[:100]))
                    top = unique[:5]
                    lines.append(f"  {col}: {len(unique)} unique values, top: {', '.join(top)}")
            
            parts.append("\n".join(lines))
    
    return "\n".join(parts)


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

MAX_REFLECT_CYCLES = 5

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


@router.post("/api/report-html", summary="AI Report HTML Generator (SSE Stream)")
async def report_html_stream(req: ReportHTMLRequest):
    """
    Generate a report as raw HTML body content via SSE.
    Same ReAct loop for data gathering, but Writer outputs HTML instead of JSON.
    """

    async def event_stream():
        start_time = time.time()
        try:
            logger.info("📝 HTML Report Agent: %s", req.question[:100])
            all_tool_results: dict[str, Any] = {}
            cycle = 0
            MAX_REFLECT_CYCLES = 4
            tool_counter = 0
            
            # --- Handle Follow-up Context ---
            parent_context = ""
            is_followup = bool(req.parent_report_id)
            if is_followup:
                sql = "SELECT prompt, tools_used FROM public.ai_reports WHERE id = $1"
                pool = await PostgresReadonly._get_pool()
                async with pool.acquire() as conn:
                    row = await conn.fetchrow(sql, req.parent_report_id)
                if row:
                    parent_context = f"THIS IS A FOLLOW-UP QUERY to a previous report.\nPrevious Report Topic: {row['prompt']}\nTools already used: {row['tools_used']}\n\nYou must focus ONLY on answering the follow-up question by querying new data if needed.\n\n"
                    logger.info("Context loaded for Follow-up to Report #%s", req.parent_report_id)

            # ── THINK ──
            yield _sse({"type": "thinking", "step": "🧠 Đang phân tích yêu cầu..."})
            
            think_input = (
                f"{parent_context}"
                f"User request: {req.question}\n\n"
                f"## INSTRUCTION:\nBreak down this request into sub_tasks and provide the FIRST set of tools to run.\n"
                f"RESPOND WITH RAW JSON ONLY (no markdown blocks)."
            )
            agent_raw = await _call_llm(HTML_AGENT_PROMPT, think_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", "")
                current_task = agent_response.get("current_task", "")
                sub_tasks = agent_response.get("sub_tasks", [])

                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", [])
                    completed = agent_response.get("completed_sub_tasks", [])
                    drill_downs = agent_response.get("drill_down_opportunities", [])
                    next_task = agent_response.get("next_task", "")
                    next_subs = agent_response.get("next_sub_tasks", [])

                    yield _sse({
                        "type": "reflect", "cycle": cycle + 1,
                        "thinking": thinking, "data_sufficient": data_sufficient,
                        "missing": missing,
                        "current_task": current_task,
                        "completed_sub_tasks": completed,
                        "drill_down_opportunities": drill_downs,
                        "next_task": next_task,
                        "next_sub_tasks": next_subs,
                    })

                    if data_sufficient and not missing:
                        break

                    tools_to_run = agent_response.get("next_tools", [])
                    if not tools_to_run:
                        break
                elif action == "write_report":
                    break

                if not tools_to_run:
                    break

                # ── ACT ──
                yield _sse({
                    "type": "plan_ready", "thinking": thinking,
                    "tools_count": len(tools_to_run), "cycle": cycle + 1,
                    "current_task": current_task,
                    "sub_tasks": sub_tasks,
                })

                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", ""),
                    })

                results = await _execute_tools_parallel(tools_to_run)

                # ── OBSERVE ──
                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 ──
                yield _sse({"type": "thinking", "step": f"🔍 Đánh giá dữ liệu (vòng {cycle}/{MAX_REFLECT_CYCLES})..."})

                # Only send NEW results from this cycle + compact summary of previous
                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)
                reflect_input = (
                    f"User request: {req.question}\n\n"
                    f"## DATA COLLECTED (cycle {cycle}/{MAX_REFLECT_CYCLES}):\n"
                    f"{prev_summary}"
                    f"### New results this cycle:\n{new_summary}\n\n"
                    f"## INSTRUCTION:\nReview the data. Respond with action='reflect'.\n"
                    f"If data_sufficient=true AND missing=[] → proceed to write.\n"
                    f"If not, provide next_tools.\nRESPOND WITH RAW JSON ONLY."
                )

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

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

            data_context = _summarize_results(all_tool_results)
            writer_input = (
                f"{parent_context}"
                f"User request: {req.question}\n\n"
                f"## REAL DATA FROM TOOLS ({tool_counter} queries, {cycle} cycles):\n{data_context}\n\n"
                f"## DATA FORMAT:\n"
                f"Data is provided as tab-separated tables. To create Chart.js charts:\n"
                f"- First column = labels[] array (x-axis: dates, categories, names)\n"
                f"- Other columns = data[] arrays (y-axis: numbers, counts, amounts)\n"
                f"- Example: if data shows 'ngay\\tdoanh_thu', use labels=['2024-01','2024-02',...] and data=[120M, 150M,...]\n\n"
                f"Generate the FULL report HTML body using the data above.\n"
                f"CRITICAL: You MUST generate AT LEAST 6 pages (6 `.page` divs). Include: Cover page, TOC + Executive Summary + KPIs, "
                f"2-3 analysis pages with charts and tables, and a Conclusion + Recommendations page. "
                f"Do NOT skip pages or abbreviate. Each content page MUST have .rh header and .pf footer."
            )
            
            if is_followup:
                writer_input += " OVERRIDE: Since this is a follow-up, DO NOT generate a full HTML page. ONLY generate a `<div class='followup-section'>...</div>` containing the answer to the follow-up question, styled using the existing report_template.css classes (e.g. .section, .kpi-grid). Output RAW HTML ONLY."
            else:
                writer_input += " Output RAW HTML ONLY."

            html_raw = await _call_llm(HTML_WRITER_PROMPT, writer_input, req.model)
            html_body = _extract_html(html_raw)

            # ── SAVE TO DB ──
            generation_time_ms = int((time.time() - start_time) * 1000)
            tools_used = list(all_tool_results.keys())
            
            # Simple title generation from the prompt
            title = req.question[:100] + ("..." if len(req.question) > 100 else "")
            
            insert_sql = """
                INSERT INTO public.ai_reports
                (title, prompt, html_content, tools_used, cycles_count, generation_time_ms)
                VALUES ($1, $2, $3, $4, $5, $6)
                RETURNING id
            """
            
            report_id = None
            try:
                report_id = await PostgresReadonly.execute_insert_async(
                    insert_sql,
                    title,
                    req.question,
                    html_body,
                    tools_used,
                    cycle,
                    generation_time_ms
                )
                logger.info("💾 Saved report to DB with ID: %s", report_id)
            except Exception as e:
                logger.error("Failed to save report to DB: %s", e)
            
            yield _sse({"type": "html_complete", "html": html_body, "report_id": report_id})

            logger.info("✅ HTML Report complete: %d chars, %d tools, %d cycles",
                        len(html_body), tool_counter, cycle)

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

        except json.JSONDecodeError as e:
            logger.error("Report HTML 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 HTML runtime error: %s", e)
            yield _sse({"type": "error", "message": str(e)[:300]})
        except Exception as e:
            logger.error("Report HTML 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")

# ─── Report History API ──────────────────────────────────────────────

@router.get("/reports")
async def get_reports():
    """List all saved reports (metadata only, no full HTML)."""
    sql = """
        SELECT id, title, prompt, tools_used, cycles_count, generation_time_ms, report_pages, status, created_by, created_at
        FROM public.ai_reports
        ORDER BY created_at DESC
        LIMIT 50
    """
    try:
        rows = await PostgresReadonly.execute_query_async(sql)
        for r in rows:
            if "created_at" in r and r["created_at"]:
                r["created_at"] = r["created_at"].isoformat()
        return {"reports": rows}
    except Exception as e:
        logger.error("Error fetching reports: %s", e)
        return {"error": str(e), "reports": []}

@router.get("/reports/{report_id}")
async def get_report_by_id(report_id: int):
    """Get a single report with full HTML content."""
    sql = "SELECT * FROM public.ai_reports WHERE id = $1"
    try:
        pool = await PostgresReadonly._get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(sql, report_id)
        
        if not row:
            return {"error": "Report not found"}
            
        r = dict(row)
        if "created_at" in r and r["created_at"]:
            r["created_at"] = r["created_at"].isoformat()
        if "updated_at" in r and r["updated_at"]:
            r["updated_at"] = r["updated_at"].isoformat()
            
        return {"report": r}
    except Exception as e:
        logger.error("Error fetching report %d: %s", report_id, e)
        return {"error": str(e)}

@router.delete("/reports/{report_id}")
async def delete_report(report_id: int):
    """Delete a generated report."""
    sql = "DELETE FROM public.ai_reports WHERE id = $1"
    try:
        pool = await PostgresReadonly._get_pool()
        async with pool.acquire() as conn:
            status = await conn.execute(sql, report_id)
        return {"success": True, "status": status}
    except Exception as e:
        logger.error("Error deleting report %d: %s", report_id, e)
        return {"error": str(e)}
