"""
Inline Edit Agent Graph — LangGraph StateGraph for report section editing.

     ┌── simple_edit ──→ rewrite → END    (rewrite/shorten/fix — no SQL needed)
think ──┤
     └── agent_edit ──→ query_data → rewrite_with_data → END  (enrich with real data)

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

import json
import logging
import re
from typing import Any, TypedDict

from langgraph.graph import StateGraph, START, END

from agent.report_agent.core import (
    call_llm,
    parse_json,
    execute_tools_parallel,
    summarize_results,
)
from agent.report_agent.prompts.inline_prompt import (
    INLINE_EDIT_PROMPT,
    AGENT_SECTION_PROMPT,
    AGENT_WRITER_PROMPT,
)

logger = logging.getLogger(__name__)


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

class InlineState(TypedDict):
    # Input
    selected_text: str
    action: str          # rewrite | enrich | shorten | fix | agent_rewrite
    context: str
    model: str
    codex_token: str | None
    openai_key: str | None

    # Internal
    needs_data: bool
    tools_to_run: list[dict]
    data_summary: str
    thinking: str

    # Output
    new_text: str
    explanation: str
    error: str | None


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

async def think_node(state: InlineState) -> dict:
    """Analyze the selected text and decide: simple edit or agent-powered rewrite."""
    action = state["action"]

    # Simple edits → no SQL needed
    if action != "agent_rewrite":
        return {"needs_data": False, "tools_to_run": [], "thinking": ""}

    # Agent rewrite → analyze what data is needed
    think_input = (
        f"Section text: \"{state['selected_text']}\"\n"
        f"Surrounding context: {state['context'][:500]}\n\n"
        f"Generate SQL queries to fetch data for enriching this section.\n"
        f"Return JSON only."
    )

    think_raw = await call_llm(
        AGENT_SECTION_PROMPT, think_input, state["model"],
        codex_token=state.get("codex_token"),
        openai_key=state.get("openai_key"),
        json_mode=True,
    )
    think_response = parse_json(think_raw)

    tools = think_response.get("tools", [])
    skip = think_response.get("action") == "skip"

    return {
        "needs_data": bool(tools) and not skip,
        "tools_to_run": tools,
        "thinking": think_response.get("thinking", ""),
    }


async def query_node(state: InlineState) -> dict:
    """Execute SQL tools to fetch real data for enriching the section."""
    tools_to_run = state.get("tools_to_run", [])
    if not tools_to_run:
        return {"data_summary": ""}

    results = await execute_tools_parallel(tools_to_run)

    all_results: dict[str, Any] = {}
    for i, (tool_spec, result) in enumerate(zip(tools_to_run, results)):
        if isinstance(result, Exception):
            result = {"error": str(result)[:200], "data": []}
        all_results[f"{tool_spec.get('name', 'q')}_{i}"] = result

    data_summary = summarize_results(all_results)

    return {"data_summary": data_summary}


async def simple_rewrite_node(state: InlineState) -> dict:
    """Simple rewrite without data: rewrite/shorten/fix."""
    user_input = (
        f"Selected text: \"{state['selected_text']}\"\n"
        f"Action: {state['action']}\n"
        f"Surrounding context: {state['context'][:500]}\n\n"
        f"Return JSON only."
    )

    raw = await call_llm(
        INLINE_EDIT_PROMPT, user_input, state["model"],
        codex_token=state.get("codex_token"),
        openai_key=state.get("openai_key"),
    )

    json_match = re.search(r'\{[\s\S]*\}', raw)
    if json_match:
        parsed = json.loads(json_match.group())
        return {
            "new_text": parsed.get("new_text", raw.strip()),
            "explanation": parsed.get("explanation", "AI đã chỉnh sửa văn bản"),
        }
    return {"new_text": raw.strip(), "explanation": "AI đã chỉnh sửa văn bản"}


async def rewrite_with_data_node(state: InlineState) -> dict:
    """Rewrite the section using real data from SQL queries."""
    data_summary = state.get("data_summary", "")

    if not data_summary.strip() or "no data" in data_summary.lower():
        return {
            "new_text": state["selected_text"],
            "explanation": "Không có dữ liệu mới để bổ sung",
        }

    write_input = (
        f"Original section:\n\"{state['selected_text']}\"\n\n"
        f"New data from queries:\n{data_summary}\n\n"
        f"Rewrite this section incorporating the new data. Return JSON only."
    )

    write_raw = await call_llm(
        AGENT_WRITER_PROMPT, write_input, state["model"],
        codex_token=state.get("codex_token"),
        openai_key=state.get("openai_key"),
    )

    json_match = re.search(r'\{[\s\S]*\}', write_raw)
    if json_match:
        parsed = json.loads(json_match.group())
        return {
            "new_text": parsed.get("new_text", write_raw.strip()),
            "explanation": parsed.get("explanation", "AI đã bổ sung dữ liệu mới"),
        }
    return {"new_text": write_raw.strip(), "explanation": "AI đã bổ sung dữ liệu mới"}


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

def route_after_think(state: InlineState) -> str:
    """Think decides: simple edit → simple_rewrite, data needed → query."""
    if state.get("needs_data"):
        return "query"
    return "simple_rewrite"


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

def build_inline_graph() -> StateGraph:
    """Build and compile the inline edit agent graph."""
    graph = StateGraph(InlineState)

    # Add nodes
    graph.add_node("think", think_node)
    graph.add_node("query", query_node)
    graph.add_node("simple_rewrite", simple_rewrite_node)
    graph.add_node("rewrite_with_data", rewrite_with_data_node)

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

    # Think → either simple_rewrite or query
    graph.add_conditional_edges("think", route_after_think, ["query", "simple_rewrite"])

    # query → rewrite_with_data → END
    graph.add_edge("query", "rewrite_with_data")
    graph.add_edge("rewrite_with_data", END)

    # simple_rewrite → END
    graph.add_edge("simple_rewrite", END)

    return graph.compile()


# Compiled graph instance
inline_graph = build_inline_graph()


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

async def run_inline_agent(
    *,
    selected_text: str,
    action: str = "rewrite",
    context: str = "",
    model: str = "codex/gpt-5.3-codex",
    codex_token: str | None = None,
    openai_key: str | None = None,
) -> dict:
    """
    Run the inline edit agent and return the result.

    Returns: {"new_text": str, "explanation": str} or {"error": str}
    """
    initial_state: InlineState = {
        "selected_text": selected_text,
        "action": action,
        "context": context,
        "model": model,
        "codex_token": codex_token,
        "openai_key": openai_key,
        "needs_data": False,
        "tools_to_run": [],
        "data_summary": "",
        "thinking": "",
        "new_text": "",
        "explanation": "",
        "error": None,
    }

    try:
        result = await inline_graph.ainvoke(initial_state)
        return {
            "new_text": result.get("new_text", selected_text),
            "explanation": result.get("explanation", ""),
        }
    except Exception as e:
        logger.error("Inline agent error: %s", e)
        return {"error": str(e)}
