from __future__ import annotations

import asyncio
import logging
import random
from typing import TYPE_CHECKING, Optional

from pydantic_ai.exceptions import ModelHTTPError

from .config import LLMAgentConfig
from .models import (
    BusinessContext,
    LLMQueryRequest,
    RealIncident,
    RiskAssessment,
    RiskAssessmentLevel,
    TechnicalFinding,
)
from .providers import create_provider

if TYPE_CHECKING:
    from puli_mcp_server.vertex_token_manager import VertexTokenManager

logger = logging.getLogger(__name__)

_RATE_LIMIT_STATUS = 429
_RETRY_ATTEMPTS = 2
_RETRY_DELAY_SECONDS = 1.0


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()
        last_exc: Optional[Exception] = None
        for attempt in range(1 + _RETRY_ATTEMPTS):
            try:
                result = await self.agent.run(prompt)
                return result.output
            except ModelHTTPError as exc:
                if exc.status_code == _RATE_LIMIT_STATUS and attempt < _RETRY_ATTEMPTS:
                    logger.warning(
                        "Rate limit hit (attempt %d/%d), retrying in %.0fs…",
                        attempt + 1,
                        1 + _RETRY_ATTEMPTS,
                        _RETRY_DELAY_SECONDS,
                    )
                    await asyncio.sleep(_RETRY_DELAY_SECONDS + random.uniform(0, 0.5))
                    last_exc = exc
                else:
                    raise
        # All retries exhausted — return a meaningful fallback instead of crashing.
        logger.error("LLM rate limit exceeded after %d attempts: %s", 1 + _RETRY_ATTEMPTS, last_exc)
        return [
            RiskAssessment(
                risk_assessment_level=RiskAssessmentLevel.LOW_RISK,
                risk_score=100,
                business_flow_name="Analysis unavailable",
                technical_finding=TechnicalFinding(
                    file_path="N/A",
                    line_number="N/A",
                    change_description="Puli could not analyse this change.",
                    technical_reason=(
                        "The LLM service returned a rate-limit error (429) and all "
                        f"{1 + _RETRY_ATTEMPTS} attempts were exhausted."
                    ),
                ),
                business_context=BusinessContext(
                    process_description="LLM-based code risk analysis",
                    flow_steps=["Request received", "LLM called", "[RATE LIMIT — retries exhausted]", "Fallback returned"],
                ),
                consequence="Analysis skipped. Re-run the check in a few seconds.",
                closing_line="Rate limit hit — please retry shortly.",
            )
        ]
