"""Vertex AI provider for Gemini models (google-vertex).

Supports two authentication modes:
  - **SA token**: Explicit credentials from ``VertexTokenManager`` (production).
  - **ADC**: Application Default Credentials via ``gcloud login`` (local dev).
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, List, Optional

from pydantic_ai import Agent
from pydantic_ai.models.google import GoogleModel
from pydantic_ai.providers.google import GoogleProvider

from ..config import LLMAgentConfig
from ..models import RiskAssessment
from .base import BaseLLMProvider

if TYPE_CHECKING:
    from puli_mcp_server.vertex_token_manager import VertexTokenManager

logger = logging.getLogger(__name__)

# Model-specific thinking configurations for Gemini models.
_THINKING_CONFIGS = {
    "gemini-3-flash": {
        "google_thinking_config": {"thinking_level": "MINIMAL"},
    },
    "gemini-3-pro": {
        "google_thinking_config": {
            "include_thoughts": False,
            "thinking_level": "LOW",
        },
    },
}


def _resolve_model_settings(model_name: str) -> Optional[dict]:
    """Return model-specific settings (thinking config) for a Gemini model."""
    for key, settings in _THINKING_CONFIGS.items():
        if key in model_name:
            return settings
    return None


class VertexGeminiProvider(BaseLLMProvider):
    """Provider for Gemini models on Vertex AI."""

    def __init__(
        self,
        config: LLMAgentConfig,
        vertex_token_manager: Optional[VertexTokenManager] = None,
    ) -> None:
        super().__init__(config)
        self._vertex_token_manager = vertex_token_manager

    def build_agent(self) -> Agent[None, List[RiskAssessment]]:
        model_settings = _resolve_model_settings(self.config.model)

        if self._vertex_token_manager:
            # SA-token path: pass fresh credentials + project/location.
            # get_credentials() returns a new Credentials object with the
            # latest token (auto-refreshed if near expiry).
            provider = GoogleProvider(
                credentials=self._vertex_token_manager.get_credentials(),
                project=self.config.vertex_project,
                location=self.config.vertex_location,
            )
            model = GoogleModel(self.config.model, provider=provider)
        else:
            # ADC path: pydantic-ai reads GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION
            model = f"google-vertex:{self.config.model}"

        return Agent(
            model,
            system_prompt=self.config.system_prompt,
            output_type=List[RiskAssessment],
            model_settings=model_settings,
        )

    def needs_rebuild(self) -> bool:
        """Rebuild when using SA tokens (credentials may have refreshed)."""
        return self._vertex_token_manager is not None

