from typing import Optional, List

from pydantic_ai import Agent

from .config import LLMAgentConfig
from .models import LLMQueryRequest, RiskAssessment


class LLMAgent:
    """Agent for querying an LLM using pydantic-ai."""

    def __init__(self, config: Optional[LLMAgentConfig] = None):
        """
        Initialize the LLM agent.

        Args:
            config: LLMAgentConfig instance. If None, loads from environment variables.
        """
        self.config = config or LLMAgentConfig.from_env()
        
        # Build the model name string (e.g., "openai:gpt-4" or "anthropic:claude-3")
        model_name = f"{self.config.provider}:{self.config.model}"

        # Initialize the Agent with the model name string
        # pydantic-ai will automatically read API keys from environment variables
        self.agent = Agent(
            model_name,
            system_prompt=self.config.system_prompt,
            output_type=List[RiskAssessment],
        )

    async def query(self, request: LLMQueryRequest) -> RiskAssessment:
        """
        Query the LLM with a ChangeSet and historical incidents.

        Args:
            request: LLMQueryRequest containing the data to analyze.

        Returns:
            The LLM response as a RiskAssessmentList object.
        """
        # Format the prompt with the request data
        prompt = request.to_prompt_str()
        result = await self.agent.run(prompt)
        return result.output

