"""Output Parser Agent

Fixes invalid JSON responses to match expected schema.
"""

from __future__ import annotations

from typing import Annotated, TypeVar

from pydantic import BaseModel, Field

from .base import AgentConfig, BaseAgentImpl, TokenUsage
from .utils import format_template

T = TypeVar("T", bound=BaseModel)


class OutputParserAgent:
    """Agent that fixes invalid JSON responses to match expected schema"""

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

        Args:
            config: LLM configuration
        """
        self.base = BaseAgentImpl(
            name="OutputParser",
            description="Fixes invalid JSON responses to match expected schema",
            config=config,
        )

    async def parse_output(
        self,
        invalid_response: Annotated[str, Field(description="The invalid response text from LLM")],
        json_schema: Annotated[str, Field(description="The expected JSON schema description")],
        response_model: Annotated[type[T], Field(description="Pydantic model class")],
    ) -> tuple[T, TokenUsage]:
        """Parse invalid response to match JSON schema

        Args:
            invalid_response: The invalid response text from LLM
            json_schema: The expected JSON schema description
            response_model: Pydantic model class to parse into

        Returns:
            Tuple of (parsed output, token usage)

        Raises:
            ReviewateError: If parsing fails
        """
        # Build context with invalid response and schema
        ctx = {
            "invalid_response": invalid_response,
            "json_schema": json_schema,
        }

        # Load system prompt template
        template = self.base.load_prompt("output_parser.txt")

        # Format system prompt
        system_prompt = format_template(self.base.jinja_env, template, ctx)

        # Call LLM directly (no retry loop to avoid recursion)
        output, token_usage = await self.base.call_llm(
            system_prompt,
            user_prompt="Please fix the following JSON response to match the schema.",
            response_model=response_model,
        )

        return output, token_usage

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

        Returns:
            OutputParserAgent with configuration from environment variables
        """
        from ..config import Config

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