from pydantic_ai import Agent
from pydantic_ai.format_prompt import format_as_xml
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from pydantic_graph.beta.join import reduce_list_append

from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import SearchAnswer, resolve_citations
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies
from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.graph.deep_qa.prompts import DECISION_PROMPT, SYNTHESIS_PROMPT
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
from haiku.rag.store.models import SearchResult


def build_deep_qa_graph(
    config: AppConfig = Config,
) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
    """Build the Deep QA graph.

    Args:
        config: AppConfig object (uses config.qa for provider, model, and graph parameters)

    Returns:
        Configured Deep QA graph
    """
    model_config = config.qa.model
    g = GraphBuilder(
        state_type=DeepQAState,
        deps_type=DeepQADeps,
        output_type=DeepQAAnswer,
    )

    # Create and register the plan node using the factory
    plan = g.step(
        create_plan_node(
            model_config=model_config,
            deps_type=DeepQADependencies,  # type: ignore[arg-type]
            activity_message="Planning approach",
            output_retries=None,  # Deep QA doesn't use output_retries
            config=config,
        )
    )  # type: ignore[arg-type]

    # Create and register the search_one node using the factory
    search_one = g.step(
        create_search_node(
            model_config=model_config,
            deps_type=DeepQADependencies,  # type: ignore[arg-type]
            with_step_wrapper=False,  # Deep QA doesn't wrap with agui_emitter step
            success_message_format="Answered: {sub_q}",
            handle_exceptions=True,
            config=config,
        )
    )  # type: ignore[arg-type]

    @g.step
    async def get_batch(
        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
    ) -> list[str] | None:
        """Get all remaining questions for this iteration."""
        state = ctx.state

        if not state.context.sub_questions:
            return None

        # Take ALL remaining questions - max_concurrency controls parallel execution within .map()
        batch = list(state.context.sub_questions)
        state.context.sub_questions.clear()
        return batch

    @g.step
    async def decide(
        ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
    ) -> bool:
        state = ctx.state
        deps = ctx.deps

        if deps.agui_emitter:
            deps.agui_emitter.start_step("decide")
            deps.agui_emitter.update_activity(
                "evaluating", {"message": "Evaluating information sufficiency"}
            )

        try:
            agent = Agent(
                model=get_model(model_config, config),
                output_type=DeepQAEvaluation,
                instructions=DECISION_PROMPT,
                retries=3,
                deps_type=DeepQADependencies,
            )

            context_data = {
                "original_question": state.context.original_question,
                "gathered_answers": [
                    {
                        "question": qa.query,
                        "answer": qa.answer,
                        "confidence": qa.confidence,
                    }
                    for qa in state.context.qa_responses
                ],
            }
            context_xml = format_as_xml(context_data, root_tag="gathered_information")

            prompt = (
                "Evaluate whether we have sufficient information to answer the question.\n\n"
                f"{context_xml}"
            )

            agent_deps = DeepQADependencies(
                client=deps.client,
                context=state.context,
            )
            result = await agent.run(prompt, deps=agent_deps)
            evaluation = result.output

            state.iterations += 1

            for new_q in evaluation.new_questions:
                if new_q not in state.context.sub_questions:
                    state.context.sub_questions.append(new_q)

            if deps.agui_emitter:
                deps.agui_emitter.update_state(state)
                status = "sufficient" if evaluation.is_sufficient else "insufficient"
                deps.agui_emitter.update_activity(
                    "evaluating",
                    {
                        "stepName": "decide",
                        "message": f"Information {status} after {state.iterations} iteration(s)",
                        "is_sufficient": evaluation.is_sufficient,
                        "iterations": state.iterations,
                    },
                )

            should_continue = (
                not evaluation.is_sufficient and state.iterations < state.max_iterations
            )

            return should_continue
        finally:
            if deps.agui_emitter:
                deps.agui_emitter.finish_step()

    @g.step
    async def synthesize(
        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
    ) -> DeepQAAnswer:
        state = ctx.state
        deps = ctx.deps

        if deps.agui_emitter:
            deps.agui_emitter.start_step("synthesize")
            deps.agui_emitter.update_activity(
                "synthesizing", {"message": "Synthesizing final answer"}
            )

        try:
            agent = Agent(
                model=get_model(model_config, config),
                output_type=SearchAnswer,
                instructions=SYNTHESIS_PROMPT,
                retries=3,
                deps_type=DeepQADependencies,
            )

            context_data = {
                "original_question": state.context.original_question,
                "sub_answers": [
                    {
                        "question": qa.query,
                        "answer": qa.answer,
                        "confidence": qa.confidence,
                        "cited_chunks": qa.cited_chunks,
                    }
                    for qa in state.context.qa_responses
                ],
            }
            context_xml = format_as_xml(context_data, root_tag="gathered_information")

            prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"

            agent_deps = DeepQADependencies(
                client=deps.client,
                context=state.context,
            )
            result = await agent.run(prompt, deps=agent_deps)
            llm_answer = result.output

            # Resolve citations by fetching chunks by ID
            search_results = []
            for chunk_id in llm_answer.cited_chunks:
                chunk = await deps.client.chunk_repository.get_by_id(chunk_id)
                if chunk:
                    search_results.append(SearchResult.from_chunk(chunk, score=1.0))
            citations = resolve_citations(llm_answer.cited_chunks, search_results)

            if deps.agui_emitter:
                deps.agui_emitter.update_activity(
                    "synthesizing", {"message": "Answer complete"}
                )

            return DeepQAAnswer(answer=llm_answer.answer, citations=citations)
        finally:
            if deps.agui_emitter:
                deps.agui_emitter.finish_step()

    # Build the graph structure
    collect_answers = g.join(
        reduce_list_append,
        initial_factory=list[SearchAnswer],
    )

    g.add(
        g.edge_from(g.start_node).to(plan),
        g.edge_from(plan).to(get_batch),
    )

    # Branch based on whether we have questions
    g.add(
        g.edge_from(get_batch).to(
            g.decision()
            .branch(g.match(list).label("Has questions").map().to(search_one))
            .branch(g.match(type(None)).label("No questions").to(synthesize))
        ),
        g.edge_from(search_one).to(collect_answers),
        g.edge_from(collect_answers).to(decide),
    )

    # Branch based on decision
    g.add(
        g.edge_from(decide).to(
            g.decision()
            .branch(
                g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
            )
            .branch(
                g.match(bool, matches=lambda x: not x)
                .label("Done with QA")
                .to(synthesize)
            )
        ),
        g.edge_from(synthesize).to(g.end_node),
    )

    return g.build()