from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Optional

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

if TYPE_CHECKING:
    from puli_mcp_server.vertex_token_manager import VertexTokenManager

logger = logging.getLogger(__name__)


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

    Delegates agent construction and credential management to a
    :class:`~providers.base.BaseLLMProvider` implementation, selected
    automatically by :func:`~providers.create_provider`.
    """

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

        Args:
            config: LLMAgentConfig instance. If None, loads from environment variables.
            vertex_token_manager: Optional token manager that provides short-lived
                SA credentials for Vertex AI. When ``None``, ADC (gcloud login) is used.
        """
        self.config = config or LLMAgentConfig.from_env()
        self._provider = create_provider(self.config, vertex_token_manager)
        self.agent = self._provider.build_agent()

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

        When using SA tokens, the agent is rebuilt before each query to
        ensure the latest (possibly refreshed) token is used. Agent
        construction is cheap (no network calls), so the overhead is
        negligible.

        Note: self.agent reassignment is safe because the MCP server runs
        as a single-user process (one per IDE session) with sequential
        query execution. There is no concurrent access to self.agent.

        Args:
            request: LLMQueryRequest containing the data to analyze.

        Returns:
            The LLM response as a list of RiskAssessment objects.
        """
        if self._provider.needs_rebuild():
            self.agent = self._provider.build_agent()
            logger.debug("Rebuilt agent with fresh credentials")

        prompt = request.to_prompt_str()
        result = await self.agent.run(prompt)
        return result.output
