from typing import Optional, List

from pydantic_ai import Agent
from pydantic_ai.models.anthropic import AnthropicModel
from pydantic_ai.providers.anthropic import AnthropicProvider

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


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()

        if self.config.provider == "google-vertex-anthropic":
            self.agent = self.create_anthropic_vertex_model()
            return

        # Build the model name string (e.g., "openai:gpt-4" or "anthropic:claude-3")
        model_name = f"{self.config.provider}:{self.config.model}"

        model_settings = None
        if self.config.provider == "google-vertex":
            if "gemini-3-flash" in self.config.model:
                model_settings = {
                    "google_thinking_config": {
                        "thinking_level": "MINIMAL" 
                    }
                }
            elif "gemini-3-pro" in self.config.model:
                model_settings = {
                    "google_thinking_config": {
                        "include_thoughts": False, # Hides the text, but doesn't stop the delay
                        "thinking_level": "LOW",     # The lowest allowed for Pro
                    }
                }

        # 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],
            model_settings=model_settings,
        )

    def create_anthropic_vertex_model(self):
        vertex_client = AsyncAnthropicVertex(
            project_id=self.config.vertex_project,
            region=self.config.vertex_location,
        )
        provider = AnthropicProvider(anthropic_client=vertex_client)
        model = AnthropicModel(self.config.model, provider=provider)
        return Agent(
            model,
            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

