"""Style Agent

Formats verbose review comments into concise, structured format.
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING, Annotated, TypeVar

from pydantic import BaseModel, Field

from ..schema import (
    BasicReviewOutput,
    GitHubReviewComment,
    GitLabReviewComment,
    ReviewOutput,
)
from .base import AgentConfig, AgentResponse, BaseAgentImpl
from .output_parser import OutputParserAgent
from .utils import format_template

if TYPE_CHECKING:
    from ..workflows.context import PipelineContext

logger = logging.getLogger(__name__)

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

# Type alias for response model types
ReviewOutputType = (
    type[ReviewOutput[GitHubReviewComment]]
    | type[ReviewOutput[GitLabReviewComment]]
    | type[ReviewOutput[BasicReviewOutput]]
)


class StyleAgent:
    """Agent that formats verbose reviews into concise, structured, scannable format"""

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

        Args:
            config: LLM configuration
        """
        self.base = BaseAgentImpl(
            name="Style",
            description="Formats verbose review comments into concise, structured, scannable format",
            config=config,
        )

    async def stylize(
        self,
        reviews: str,
        output_parser: OutputParserAgent,
        response_model: ReviewOutputType,
    ) -> AgentResponse[
        ReviewOutput[GitHubReviewComment]
        | ReviewOutput[GitLabReviewComment]
        | ReviewOutput[BasicReviewOutput]
    ]:
        """Stylize review comments to be more concise and scannable

        Args:
            reviews: JSON array of review objects to format/stylize
            output_parser: Output parser agent for fixing invalid JSON responses
            response_model: Pydantic model class for response

        Returns:
            AgentResponse with parsed output and token usage

        Raises:
            ReviewateError: If stylization fails
        """
        # Build context with reviews to format
        ctx = {
            "reviews": reviews,
        }

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

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

        # Add platform-specific output format instructions
        output_format = self.base.get_output_format(response_model)

        # Call base analyze with strong typing
        return await self.base.analyze(
            system_prompt=system_prompt,
            user_prompt=output_format,
            response_model=response_model,
            output_parser=output_parser,
        )

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

        Formats ctx.reviews to be more concise and scannable.
        """
        if ctx.response_model is None:
            raise ValueError("response_model must be set before running style")
        if not ctx.reviews:
            logger.info("Style skipped: no reviews")
            return

        logger.info(f"Starting Style on {len(ctx.reviews)} reviews")

        # Convert reviews to JSON for the prompt
        reviews_json = json.dumps([r.model_dump() for r in ctx.reviews], indent=2)

        result = await self.stylize(
            reviews=reviews_json,
            output_parser=output_parser,
            response_model=ctx.response_model,
        )

        styled = list(result.output.reviews)

        # Only take the reformatted body from the style agent.
        # Preserve path/line/side from the originals so the LLM can't drop them.
        if len(styled) == len(ctx.reviews):
            for original, new in zip(ctx.reviews, styled, strict=True):
                original.body = new.body  # type: ignore[attr-defined]
        else:
            # Count mismatch (style agent split/merged reviews) — use styled as-is
            logger.warning(
                f"Style review count changed ({len(ctx.reviews)} → {len(styled)}), "
                "using styled output as-is"
            )
            ctx.reviews = styled

        ctx.track("style", result.metadata)
        logger.info(f"Style completed: {len(ctx.reviews)} reviews")

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

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

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