"""
Follow-up Report Agent Graph — LangGraph for appending sections to existing reports.

     ┌── direct_answer ──→ END
load_context ──→ think ──┤
                         └── execute ──→ reflect ──→ write_section ──→ END
                                          ↑            │
                                          └── loop ────┘ (need more data)

Generates "Phụ lục A, B, C..." sections appended after the main report.
Used by api/report_html_route.py via `run_followup_agent()`.
"""

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

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

from agent.report_agent.core import (ThinkingStreamer, _extract_report_outline, call_llm, call_llm_streaming,
                                     execute_tools_parallel, extract_html, load_report_context, make_result_preview,
                                     parse_json, summarize_results)
from agent.report_agent.prompts.follow_up_prompt import FOLLOWUP_AGENT_PROMPT, FOLLOWUP_WRITER_PROMPT

try:
    from agent.report_agent.error_recovery import classify_and_log_error
    from agent.report_agent.context_manager import ReportContextManager
except ImportError:
    classify_and_log_error = None
    ReportContextManager = None

logger = logging.getLogger(__name__)

MAX_FOLLOWUP_CYCLES = 2  # Follow-ups are focused, fewer cycles needed

# Appendix labels
APPENDIX_LABELS = [
    "Phụ lục A", "Phụ lục B", "Phụ lục C", "Phụ lục D",
    "Phụ lục E", "Phụ lục F", "Phụ lục G", "Phụ lục H",
]


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

class FollowupState(TypedDict):
    # Input
    question: str
    model: str
    codex_token: str | None
    openai_key: str | None
    parent_report_id: int
    appendix_index: int       # 0=A, 1=B, 2=C...

    # Context loaded from parent report
    report_outline: str
    report_topic: str
    tools_already_used: str

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

    # Output — events for SSE streaming
    events: Annotated[list[dict], add]

    # Final
    html_section: str
    tools_used: list[str]


# ─── Nodes ───────────────────────────────────────────────────────────

async def load_context_node(state: FollowupState, writer: StreamWriter) -> dict:
    """Load the parent report outline and metadata."""
    parent_report_id = state["parent_report_id"]

    writer({"type": "thinking", "step": "📖 Đang đọc báo cáo trước..."})

    row = await load_report_context(parent_report_id)
    if not row:
        writer({"type": "error", "message": f"Không tìm thấy báo cáo #{parent_report_id}"})
        return {
            "report_outline": "",
            "report_topic": "",
            "tools_already_used": "",
            "events": [],
        }

    # Extract outline from HTML
    html_content = row.get("html_content", "")
    outline = _extract_report_outline(html_content) if html_content else "(no outline)"

    writer({"type": "thinking", "step": f"📋 Đã đọc outline: {row['prompt'][:60]}..."})

    return {
        "report_outline": outline,
        "report_topic": row.get("prompt", ""),
        "tools_already_used": str(row.get("tools_used", "")),
        "events": [],
    }


async def think_node(state: FollowupState, writer: StreamWriter) -> dict:
    """Analyze the follow-up question and decide what to query."""
    question = state["question"]
    model = state["model"]
    outline = state.get("report_outline", "")
    topic = state.get("report_topic", "")
    tools_used = state.get("tools_already_used", "")

    writer({"type": "thinking", "step": "🤔 Đang phân tích câu hỏi follow-up..."})

    think_input = (
        f"## PREVIOUS REPORT:\n"
        f"Topic: {topic}\n"
        f"Tools already used: {tools_used}\n\n"
        f"## REPORT OUTLINE:\n{outline}\n\n"
        f"## FOLLOW-UP QUESTION:\n{question}\n\n"
        f"Analyze this follow-up. What new data do we need?\n"
        f"RESPOND WITH RAW JSON ONLY."
    )

    agent_raw = ""
    streamer = ThinkingStreamer(writer)
    async for token in call_llm_streaming(
        FOLLOWUP_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)

    # Direct answer (no SQL needed)
    if agent_response.get("action") == "direct_answer":
        answer = agent_response.get("answer", "")
        writer({"type": "direct_answer", "message": answer})

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


async def execute_node(state: FollowupState) -> dict:
    """Execute tools to fetch new data."""
    agent_response = state["agent_response"]
    tool_counter = state.get("tool_counter", 0)
    cycle = state.get("cycle", 0)

    # Get tools from either initial "execute" or reflect's "next_tools"
    action = agent_response.get("action", "execute")
    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}

    events: list[dict] = []

    # 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
    try:
        results = await execute_tools_parallel(tools_to_run)
        
        if ReportContextManager:
            ctx_mgr = ReportContextManager(max_tokens=60000)
            results = [ctx_mgr.truncate_result(res) if isinstance(res, (dict, list)) else res 
                       for res in results]
    except Exception as e:
        if classify_and_log_error:
            classify_and_log_error(e, context={"node": "execute_follow_up", "tools": tools_to_run})
        results = [{"error": str(e)[:300]} for _ in tools_to_run]

    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: FollowupState) -> dict:
    """Check if we have enough data for the follow-up answer."""
    question = state["question"]
    model = state["model"]
    cycle = state.get("cycle", 1)
    all_tool_results = state.get("all_tool_results", {})

    data_summary = summarize_results(all_tool_results)
    reflect_input = (
        f"Follow-up question: {question}\n\n"
        f"## DATA COLLECTED (cycle {cycle}/{MAX_FOLLOWUP_CYCLES}):\n{data_summary}\n\n"
        f"Is this enough to answer the follow-up? Respond with action='reflect'.\n"
        f"If data_sufficient=true → write. If not, provide next_tools.\n"
        f"RESPOND WITH RAW JSON ONLY."
    )

    reflect_raw = await call_llm(
        FOLLOWUP_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)

    events: list[dict] = [{
        "type": "reflect", "cycle": cycle,
        "thinking": agent_response.get("thinking", ""),
        "data_sufficient": agent_response.get("data_sufficient", False),
    }]

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


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

def route_after_think(state: FollowupState) -> str:
    """Think decides: direct answer → END, needs data → execute."""
    action = state["agent_response"].get("action", "execute")
    if action == "direct_answer":
        return END
    return "execute"


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

    if cycle >= MAX_FOLLOWUP_CYCLES:
        return END

    data_sufficient = agent_response.get("data_sufficient", False)
    next_tools = agent_response.get("next_tools", [])

    if data_sufficient:
        return END
    if next_tools:
        return "execute"
    return END


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

def build_followup_graph() -> StateGraph:
    """Build and compile the follow-up agent graph."""
    graph = StateGraph(FollowupState)

    graph.add_node("load_context", load_context_node)
    graph.add_node("think", think_node)
    graph.add_node("execute", execute_node)
    graph.add_node("reflect", reflect_node)

    # Flow
    graph.add_edge(START, "load_context")
    graph.add_edge("load_context", "think")
    graph.add_conditional_edges("think", route_after_think, ["execute", END])
    graph.add_edge("execute", "reflect")
    graph.add_conditional_edges("reflect", route_after_reflect, ["execute", END])

    return graph.compile()


followup_graph = build_followup_graph()


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

async def run_followup_agent(
    *,
    question: str,
    parent_report_id: int,
    appendix_index: int = 0,
    model: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
):
    """
    Run the follow-up agent and yield SSE events.
    Generates a single follow-up section ("Phụ lục A/B/C...").
    """
    initial_state: FollowupState = {
        "question": question,
        "model": model,
        "codex_token": codex_token,
        "openai_key": openai_key,
        "parent_report_id": parent_report_id,
        "appendix_index": appendix_index,
        "report_outline": "",
        "report_topic": "",
        "tools_already_used": "",
        "agent_response": {},
        "all_tool_results": {},
        "cycle": 0,
        "tool_counter": 0,
        "events": [],
        "html_section": "",
        "tools_used": [],
    }

    tool_counter = 0
    cycle = 0
    all_tool_results = {}
    direct_answer = False
    report_outline = ""

    async for stream_mode, chunk_data in followup_graph.astream(
        initial_state,
        stream_mode=["updates", "custom"],
    ):
        if stream_mode == "custom":
            data = chunk_data
            if isinstance(data, dict) and data.get("type") == "direct_answer":
                direct_answer = True
            yield data
            continue

        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 "report_outline" in updates: report_outline = updates["report_outline"]

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

    if direct_answer:
        yield {"type": "done"}
        return

    # WRITE PHASE — generate the follow-up section HTML
    appendix_label = APPENDIX_LABELS[min(appendix_index, len(APPENDIX_LABELS) - 1)]
    yield {"type": "thinking", "step": f"✍️ Đang viết {appendix_label}..."}

    data_context = summarize_results(all_tool_results)
    writer_input = (
        f"## APPENDIX LABEL: {appendix_label}\n"
        f"## USER FOLLOW-UP QUESTION: {question}\n\n"
        f"## PREVIOUS REPORT OUTLINE:\n{report_outline}\n\n"
        f"## NEW DATA FROM QUERIES ({tool_counter} queries):\n{data_context}\n\n"
        f"Generate a follow-up section as a `<div class='followup-section'>` "
        f"with header '{appendix_label}' and the analysis. "
        f"Use existing report CSS classes. Output RAW HTML ONLY."
    )

    html_section = ""
    async for token in call_llm_streaming(
        FOLLOWUP_WRITER_PROMPT, writer_input, model,
        codex_token=codex_token,
        openai_key=openai_key,
    ):
        html_section += token
        yield {"type": "html_token", "token": token}

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

    yield {
        "type": "followup_complete",
        "html": html_section,
        "appendix_label": appendix_label,
        "tools_used": tools_used,
        "cycles_count": cycle,
    }
    yield {"type": "done"}

    logger.info(
        "✅ Follow-up %s complete: %d chars, %d tools",
        appendix_label, len(html_section), tool_counter,
    )
