"""
Report Agent Graph — LangGraph StateGraph for report generation.

     ┌── direct_response ──→ END
router ──┤
     └── analyze ──→ execute → reflect ──┐
                       ↑                  │
                       └── need_more ─────┘
                            │
                            └── sufficient ──→ write → END

Used by api/report_html_route.py via `run_report_agent()`.
"""

import logging
from operator import add
from typing import Annotated, Any, TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.types import StreamWriter

from agent.report_agent.core import (
    call_llm,
    call_llm_streaming,
    parse_json,
    extract_html,
    execute_tools_parallel,
    make_result_preview,
    summarize_results,
    load_report_context,
    ThinkingStreamer,
)
from agent.report_agent.prompts.agent_prompt import HTML_AGENT_PROMPT
from agent.report_agent.prompts.writer_prompt import HTML_WRITER_PROMPT

logger = logging.getLogger(__name__)

MAX_REFLECT_CYCLES = 4


# ─── State ───────────────────────────────────────────────────────────

class ReportState(TypedDict):
    # Input
    question: str
    model: str
    codex_token: str | None
    openai_key: str | None
    parent_report_id: int | None
    current_html: str | None

    # Internal
    parent_context: str
    is_followup: bool
    agent_response: dict
    all_tool_results: Annotated[dict, lambda a, b: {**a, **b}]
    cycle: int
    tool_counter: int

    # Output — events accumulate across nodes
    events: Annotated[list[dict], add]

    # Final output
    html_body: str
    tools_used: list[str]
    cycles_count: int


# ─── Nodes ───────────────────────────────────────────────────────────
async def router_node(state: ReportState, writer: StreamWriter) -> dict:
    """Classify the question: direct_response (greeting) or analyze (needs data).

    Uses StreamWriter to push thinking tokens in real-time to the frontend.
    """
    question = state["question"]
    model = state["model"]
    parent_report_id = state.get("parent_report_id")

    # Handle follow-up context
    parent_context = ""
    is_followup = bool(parent_report_id)
    if is_followup:
        row = await load_report_context(parent_report_id)
        if row:
            parent_context = (
                f"THIS IS A FOLLOW-UP QUERY to a previous report.\n"
                f"Previous Report ID: {parent_report_id}\n"
                f"Previous Report Topic: {row['prompt']}\n"
                f"Tools already used: {row['tools_used']}\n\n"
                f"TIP: Use read_report_outline(report_id={parent_report_id}) to see what sections "
                f"the previous report covers, so you can avoid repeating analysis.\n"
                f"Use read_report_section(report_id={parent_report_id}, section='...') to read specific sections.\n\n"
                f"Focus ONLY on answering the follow-up question by querying new data.\n\n"
            )
            logger.info("Context loaded for Follow-up to Report #%s", parent_report_id)

    think_input = (
        f"{parent_context}"
        f"User request: {question}\n\n"
        f"## INSTRUCTION:\n"
        f"Break 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)."
    )

    # Stream router tokens in real-time via StreamWriter
    writer({"type": "thinking", "step": "Đang phân tích yêu cầu..."})

    agent_raw = ""
    streamer = ThinkingStreamer(writer)
    async for token in call_llm_streaming(
        HTML_AGENT_PROMPT, think_input, model,
        codex_token=state.get("codex_token"),
        openai_key=state.get("openai_key"),
    ):
        agent_raw += token
        streamer.feed(token)

    agent_response = parse_json(agent_raw)

    # If direct response, push immediately via writer (not events list)
    if agent_response.get("action") == "direct_response":
        message = agent_response.get("message", "Xin chào! Tôi là AI Report Agent của Canifa.")
        writer({"type": "direct_response", "message": message})
        writer({"type": "done"})

    return {
        "agent_response": agent_response,
        "parent_context": parent_context,
        "is_followup": is_followup,
        "events": [],
        "cycle": 0,
        "tool_counter": 0,
        "all_tool_results": {},
    }


async def execute_node(state: ReportState) -> dict:
    """ACT: Execute tools in parallel and OBSERVE results."""
    agent_response = state["agent_response"]
    tool_counter = state.get("tool_counter", 0)
    cycle = state.get("cycle", 0)
    action = agent_response.get("action", "execute")

    # Determine which tools to run
    if action == "execute":
        tools_to_run = agent_response.get("tools", [])
    elif action == "reflect":
        tools_to_run = agent_response.get("next_tools", [])
    else:
        tools_to_run = []

    if not tools_to_run:
        return {"events": [], "cycle": cycle}

    thinking = agent_response.get("thinking", "")
    current_task = agent_response.get("current_task", "")
    sub_tasks = agent_response.get("sub_tasks", [])

    events: list[dict] = []

    # Plan ready event
    events.append({
        "type": "plan_ready", "thinking": thinking,
        "tools_count": len(tools_to_run), "cycle": cycle + 1,
        "current_task": current_task, "sub_tasks": sub_tasks,
    })

    # Tool call events
    for i, tool_spec in enumerate(tools_to_run):
        events.append({
            "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 results
    new_tool_results: dict[str, Any] = {}
    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}"
        new_tool_results[key] = result

        events.append({
            "type": "tool_result", "index": tool_counter + i,
            "tool": tool_name,
            "success": "error" not in result,
            "preview": make_result_preview(result),
        })

    return {
        "all_tool_results": new_tool_results,
        "tool_counter": tool_counter + len(tools_to_run),
        "cycle": cycle + 1,
        "events": events,
    }


async def reflect_node(state: ReportState) -> dict:
    """REFLECT: Ask LLM to assess data sufficiency."""
    question = state["question"]
    model = state["model"]
    cycle = state.get("cycle", 1)
    all_tool_results = state.get("all_tool_results", {})

    events: list[dict] = [
        {"type": "thinking", "step": f"🔍 Đánh giá dữ liệu (vòng {cycle}/{MAX_REFLECT_CYCLES})..."}
    ]

    data_summary = summarize_results(all_tool_results)
    reflect_input = (
        f"User request: {question}\n\n"
        f"## DATA COLLECTED (cycle {cycle}/{MAX_REFLECT_CYCLES}):\n{data_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, model,
        codex_token=state.get("codex_token"),
        openai_key=state.get("openai_key"),
        json_mode=True,
    )
    agent_response = parse_json(reflect_raw)

    # Emit reflect event
    thinking = agent_response.get("thinking", "")
    data_sufficient = agent_response.get("data_sufficient", False)
    missing = agent_response.get("missing", [])

    events.append({
        "type": "reflect", "cycle": cycle,
        "thinking": thinking, "data_sufficient": data_sufficient,
        "missing": missing,
        "current_task": agent_response.get("current_task", ""),
        "completed_sub_tasks": agent_response.get("completed_sub_tasks", []),
        "drill_down_opportunities": agent_response.get("drill_down_opportunities", []),
        "next_task": agent_response.get("next_task", ""),
        "next_sub_tasks": agent_response.get("next_sub_tasks", []),
    })

    return {
        "agent_response": agent_response,
        "events": events,
    }





# ─── Routing Functions ──────────────────────────────────────────────

def route_after_router(state: ReportState) -> str:
    """Router decides: greeting → END, analysis → execute."""
    action = state["agent_response"].get("action", "execute")
    if action == "direct_response":
        return END
    return "execute"


def route_after_reflect(state: ReportState) -> str:
    """Reflect decides: sufficient → END, need more → execute, max cycles → END."""
    agent_response = state["agent_response"]
    action = agent_response.get("action", "reflect")
    cycle = state.get("cycle", 0)

    # Max cycles reached — force write
    if cycle >= MAX_REFLECT_CYCLES:
        logger.info("⚠️ Max reflect cycles (%d) reached, forcing write", MAX_REFLECT_CYCLES)
        return END

    if action == "write_report":
        return END

    if action == "reflect":
        data_sufficient = agent_response.get("data_sufficient", False)
        missing = agent_response.get("missing", [])
        next_tools = agent_response.get("next_tools", [])

        if data_sufficient and not missing:
            return END

        if next_tools:
            return "execute"

        # No next tools but data not sufficient — just write with what we have
        logger.info("⚠️ Reflect: not sufficient but no next_tools, forcing write")
        return END

    return END


# ─── Build Graph ─────────────────────────────────────────────────────

def build_report_graph() -> StateGraph:
    """Build and compile the report agent graph."""
    graph = StateGraph(ReportState)

    # Add nodes
    graph.add_node("router", router_node)
    graph.add_node("execute", execute_node)
    graph.add_node("reflect", reflect_node)

    # Entry point
    graph.add_edge(START, "router")

    # Router → either END (greeting) or execute (analysis)
    graph.add_conditional_edges("router", route_after_router, ["execute", END])

    # Execute → always goes to reflect
    graph.add_edge("execute", "reflect")

    # Reflect → either back to execute (need more data) or end (sufficient)
    graph.add_conditional_edges("reflect", route_after_reflect, ["execute", END])

    return graph.compile()


# Compiled graph instance
report_graph = build_report_graph()


# ─── Public API ──────────────────────────────────────────────────────

async def run_report_agent(
    *,
    question: str,
    model: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
    parent_report_id: int | None = None,
    current_html: str | None = None,
):
    """
    Run the report agent graph and yield SSE event dicts.

    This is the main entry point used by api/report_html_route.py.
    Uses LangGraph's astream to execute nodes and emit events progressively.
    """
    initial_state: ReportState = {
        "question": question,
        "model": model,
        "codex_token": codex_token,
        "openai_key": openai_key,
        "parent_report_id": parent_report_id,
        "current_html": current_html,
        "parent_context": "",
        "is_followup": False,
        "agent_response": {},
        "all_tool_results": {},
        "cycle": 0,
        "tool_counter": 0,
        "events": [],
        "html_body": "",
        "tools_used": [],
        "cycles_count": 0,
    }

    # Track state for the final write phase
    tool_counter = 0
    cycle = 0
    all_tool_results = {}
    parent_context = ""
    is_followup = False
    direct_response = False

    # Stream with combined modes: "updates" for node state, "custom" for real-time tokens
    async for stream_mode, chunk_data in report_graph.astream(
        initial_state,
        stream_mode=["updates", "custom"],
    ):
        # Custom events (real-time tokens from StreamWriter inside nodes)
        if stream_mode == "custom":
            data = chunk_data
            # Detect direct_response pushed via writer()
            if isinstance(data, dict) and data.get("type") == "direct_response":
                direct_response = True
            yield data  # e.g. {"type": "thinking_token", "token": "..."}
            continue

        # Node state updates
        if stream_mode == "updates":
            for node_name, updates in chunk_data.items():
                if "tool_counter" in updates: tool_counter = updates["tool_counter"]
                if "cycle" in updates: cycle = updates["cycle"]
                if "all_tool_results" in updates: all_tool_results = updates["all_tool_results"]
                if "parent_context" in updates: parent_context = updates["parent_context"]
                if "is_followup" in updates: is_followup = updates["is_followup"]

                new_events = updates.get("events", [])
                for event in new_events:
                    yield event

    if direct_response:
        return

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

    logger.info("🔧 Total tools executed: %d across %d cycles", tool_counter, cycle)

    data_context = summarize_results(all_tool_results)
    writer_input = (
        f"{parent_context}"
        f"User request: {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_body = ""
    async for token in call_llm_streaming(
        HTML_WRITER_PROMPT, writer_input, model,
        codex_token=codex_token,
        openai_key=openai_key,
    ):
        html_body += token
        yield {"type": "html_token", "token": token}

    html_body = extract_html(html_body)
    tools_used = list(all_tool_results.keys())

    yield {
        "type": "html_complete",
        "html": html_body,
        "tools_used": tools_used,
        "cycles_count": cycle,
    }
    yield {"type": "done"}

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