"""Triage Agent

Analyzes merge requests to determine optimal reasoning levels for each review agent.
Runs first in the pipeline to optimize token costs based on PR complexity.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Annotated

from pydantic import BaseModel, ConfigDict, Field

from ..schema import MergeRequest
from .base import AgentConfig, AgentResponse, BaseAgentImpl, ReasoningEffort
from .utils import format_template

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from ..workflows.context import PipelineContext


class ReasoningConfig(BaseModel):
    """Reasoning configuration for agents based on PR analysis.

    Simplified for ensemble approach: one level for all reviewers, one for fact checker.
    """

    model_config = ConfigDict(extra="forbid")

    reviewer: Annotated[
        ReasoningEffort,
        Field(description="Reasoning level for all ensemble reviewers"),
    ] = None
    fact_checker: Annotated[
        ReasoningEffort,
        Field(description="Reasoning level for fact checker agent"),
    ] = None


class TriageAgent:
    """Agent that analyzes PRs to determine optimal reasoning levels.

    This agent runs first in the review pipeline and outputs reasoning
    configuration for each downstream agent based on:
    - PR complexity (lines changed, files affected)
    - Change types (security-sensitive, performance-critical, etc.)
    - Risk assessment

    The triage agent itself uses no reasoning to minimize costs.
    """

    def __init__(
        self,
        config: Annotated[AgentConfig, Field(description="LLM configuration")],
    ):
        """Initialize TriageAgent.

        Args:
            config: LLM configuration
        """
        # Triage agent uses no reasoning itself - it's meant to be cheap
        self.base = BaseAgentImpl(
            name="Triage",
            description="Analyzes PRs to determine optimal reasoning levels for review agents",
            config=config,
            reasoning_effort=None,
        )

    async def analyze(
        self,
        merge_info: MergeRequest,
        diffs: str,
    ) -> AgentResponse[ReasoningConfig]:
        """Analyze PR to determine reasoning levels for each agent.

        Args:
            merge_info: The merge request information
            diffs: The code diffs

        Returns:
            AgentResponse with ReasoningConfig containing reasoning levels per agent
        """
        # Build context for the prompt
        ctx = {
            "title": merge_info.title,
            "description": merge_info.description,
            "target_branch": merge_info.target_branch,
            "source_branch": merge_info.source_branch,
            "project_id": merge_info.repo,
        }

        # Format diffs with guardrail for caching
        cached_diffs = self.base.format_cached_diff(diffs)

        # Load and format prompt
        template = self.base.load_prompt("triage_agent.txt")
        system_prompt = format_template(self.base.jinja_env, template, ctx)

        # Get output format
        output_format = self.base.get_output_format(ReasoningConfig)

        # Call LLM - no retry needed, this is a simple classification task
        output, token_usage = await self.base.call_llm(
            system_prompt=system_prompt,
            user_prompt=output_format,
            response_model=ReasoningConfig,
            cached_content=cached_diffs,
        )

        return AgentResponse(output=output, metadata=token_usage)

    async def run(self, ctx: PipelineContext) -> None:
        """Execute as a pipeline step.

        Reads merge_request from ctx, analyzes it, and sets reasoning_effort.
        """
        if ctx.merge_request is None:
            raise ValueError("merge_request must be set before running triage")

        logger.info("Starting Triage")
        result = await self.analyze(
            merge_info=ctx.merge_request,
            diffs=ctx.merge_request.diffs,
        )

        # Store reasoning config
        ctx.reasoning_config = result.output
        ctx.reasoning_effort = result.output.reviewer  # Used by ensemble reviewers
        ctx.track("triage", result.metadata)
        logger.info(
            f"Triage: reviewer={result.output.reviewer}, fact_checker={result.output.fact_checker}"
        )

    @classmethod
    def default(cls) -> TriageAgent:
        """Create TriageAgent with default config from environment."""
        from ..config import Config

        config = Config.from_env()
        return cls(config.create_agent_config("triage_agent"))
