"""Common node implementations for graph workflows."""

import asyncio
from collections.abc import Awaitable, Callable
from typing import Any, Protocol

from pydantic_ai import Agent, RunContext
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import StepContext

from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import RawSearchAnswer, ResearchPlan, SearchAnswer
from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.store.models import SearchResult


class GraphContext(Protocol):
    """Protocol for graph context objects."""

    original_question: str
    sub_questions: list[str]

    def add_qa_response(self, qa: SearchAnswer) -> None:
        """Add a QA response to context."""
        ...


class GraphState(Protocol):
    """Protocol for graph state objects."""

    context: GraphContext
    max_concurrency: int


class GraphDeps(Protocol):
    """Protocol for graph dependencies."""

    client: HaikuRAG
    agui_emitter: AGUIEmitter[Any, Any] | None
    semaphore: asyncio.Semaphore | None


class GraphAgentDeps(Protocol):
    """Protocol for agent dependencies."""

    client: HaikuRAG
    context: GraphContext
    search_results: list[SearchResult]


def create_plan_node[AgentDepsT: GraphAgentDeps](
    model_config: ModelConfig,
    deps_type: type[AgentDepsT],
    activity_message: str = "Creating plan",
    output_retries: int | None = None,
    config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
    """Create a plan node for any graph.

    Args:
        model_config: ModelConfig with provider, model, and settings
        deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
        activity_message: Message to show during planning activity
        output_retries: Number of output retries for the agent (optional)
        config: AppConfig object (defaults to global Config)

    Returns:
        Async function that can be used as a graph step
    """

    async def plan(ctx: StepContext[Any, Any, None], /) -> None:
        state: GraphState = ctx.state  # type: ignore[assignment]
        deps: GraphDeps = ctx.deps  # type: ignore[assignment]

        if deps.agui_emitter:
            deps.agui_emitter.start_step("plan")
            deps.agui_emitter.update_activity(
                "planning", {"stepName": "plan", "message": activity_message}
            )

        try:
            # Build agent configuration
            agent_config = {
                "model": get_model(model_config, config),
                "output_type": ResearchPlan,
                "instructions": (
                    PLAN_PROMPT
                    + "\n\nUse the gather_context tool once on the main question before planning."
                ),
                "retries": 3,
                "deps_type": deps_type,
            }
            if output_retries is not None:
                agent_config["output_retries"] = output_retries

            plan_agent = Agent(**agent_config)

            @plan_agent.tool
            async def gather_context(
                ctx2: RunContext[AgentDepsT], query: str, limit: int = 6
            ) -> str:
                results = await ctx2.deps.client.search(query, limit=limit)
                results = await ctx2.deps.client.expand_context(results)
                return "\n\n".join(r.content for r in results)

            # Tool is registered via decorator above
            _ = gather_context

            prompt = (
                "Plan a focused approach for the main question.\n\n"
                f"Main question: {state.context.original_question}"
            )

            # Create agent dependencies
            agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]
            plan_result = await plan_agent.run(prompt, deps=agent_deps)
            state.context.sub_questions = list(plan_result.output.sub_questions)

            # State now contains the plan - emit state update and narrate
            if deps.agui_emitter:
                deps.agui_emitter.update_state(state)
                count = len(state.context.sub_questions)
                deps.agui_emitter.update_activity(
                    "planning",
                    {
                        "stepName": "plan",
                        "message": f"Created plan with {count} sub-questions",
                        "sub_questions": list(state.context.sub_questions),
                    },
                )
        finally:
            if deps.agui_emitter:
                deps.agui_emitter.finish_step()

    return plan


def create_search_node[AgentDepsT: GraphAgentDeps](
    model_config: ModelConfig,
    deps_type: type[AgentDepsT],
    with_step_wrapper: bool = True,
    success_message_format: str = "Answered: {sub_q}",
    handle_exceptions: bool = False,
    config: AppConfig = Config,
) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
    """Create a search_one node for any graph.

    Args:
        model_config: ModelConfig with provider, model, and settings
        deps_type: Type of dependencies for the agent
        with_step_wrapper: Whether to wrap with agui_emitter start/finish step
        success_message_format: Format string for success activity message
        handle_exceptions: Whether to handle exceptions with fallback answer
        config: AppConfig object (defaults to global Config)

    Returns:
        Async function that can be used as a graph step
    """

    async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
        state: GraphState = ctx.state  # type: ignore[assignment]
        deps: GraphDeps = ctx.deps  # type: ignore[assignment]
        sub_q = ctx.inputs

        # Create unique step name from question text
        step_name = f"search: {sub_q}"

        if deps.agui_emitter and with_step_wrapper:
            deps.agui_emitter.start_step(step_name)

        try:
            # Create semaphore if not already provided
            if deps.semaphore is None:
                deps.semaphore = asyncio.Semaphore(state.max_concurrency)

            # Use semaphore to control concurrency
            async with deps.semaphore:
                return await _do_search(
                    state,
                    deps,
                    sub_q,
                    model_config,
                    deps_type,
                    success_message_format,
                    handle_exceptions,
                    config,
                )
        finally:
            if deps.agui_emitter and with_step_wrapper:
                deps.agui_emitter.finish_step()

    return search_one


async def _do_search[AgentDepsT: GraphAgentDeps](
    state: GraphState,
    deps: GraphDeps,
    sub_q: str,
    model_config: ModelConfig,
    deps_type: type[AgentDepsT],
    success_message_format: str,
    handle_exceptions: bool,
    config: AppConfig,
) -> SearchAnswer:
    """Internal search implementation."""
    if deps.agui_emitter:
        deps.agui_emitter.update_activity(
            "searching",
            {
                "stepName": "search_one",
                "message": f"Searching: {sub_q}",
                "query": sub_q,
            },
        )

    agent = Agent(
        model=get_model(model_config, config),
        output_type=ToolOutput(RawSearchAnswer, max_retries=3),
        instructions=SEARCH_AGENT_PROMPT,
        retries=3,
        deps_type=deps_type,
    )

    @agent.tool
    async def search_and_answer(
        ctx2: RunContext[AgentDepsT], query: str, limit: int = 5
    ) -> str:
        """Search the knowledge base for relevant documents.

        Returns results with chunk IDs and relevance scores.
        Reference results by their chunk_id in cited_chunks.
        """
        results = await ctx2.deps.client.search(query, limit=limit)
        results = await ctx2.deps.client.expand_context(results)
        # Store results for citation resolution
        ctx2.deps.search_results = results

        parts = []
        for r in results:
            parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")

        if not parts:
            return f"No relevant information found in the knowledge base for: {query}"

        return "\n\n".join(parts)

    # Tool is registered via decorator above
    _ = search_and_answer

    agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]

    try:
        result = await agent.run(sub_q, deps=agent_deps)
        raw_answer = result.output
        if raw_answer:
            # Convert RawSearchAnswer to SearchAnswer with resolved citations
            answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
            state.context.add_qa_response(answer)
            # State updated with new answer - emit state update and narrate
            if deps.agui_emitter:
                deps.agui_emitter.update_state(state)
                # Format the success message
                if "{confidence" in success_message_format:
                    message = success_message_format.format(
                        sub_q=sub_q, confidence=answer.confidence
                    )
                else:
                    message = success_message_format.format(sub_q=sub_q)
                deps.agui_emitter.update_activity(
                    "searching",
                    {
                        "stepName": "search_one",
                        "message": message,
                        "query": sub_q,
                        "confidence": answer.confidence,
                    },
                )
            return answer
        # Return empty SearchAnswer if no result
        return SearchAnswer(query=sub_q, answer="", confidence=0.0)
    except Exception as e:
        if handle_exceptions:
            # Narrate the error
            if deps.agui_emitter:
                deps.agui_emitter.update_activity(
                    "searching",
                    {
                        "stepName": "search_one",
                        "message": f"Search failed: {e}",
                        "query": sub_q,
                        "error": str(e),
                    },
                )
            failure_answer = SearchAnswer(
                query=sub_q,
                answer=f"Search failed after retries: {str(e)}",
                confidence=0.0,
            )
            return failure_answer
        else:
            raise