"""Vertex AI provider for Anthropic Claude models (google-vertex-anthropic).

Uses ``AsyncAnthropicVertex`` from the Anthropic SDK, wrapped by pydantic-ai's
``AnthropicProvider`` / ``AnthropicModel``.

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

Note: ``AsyncAnthropicVertex`` captures the access token at construction time,
so the agent must be rebuilt whenever the token refreshes.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, List, Optional

from anthropic import AsyncAnthropicVertex
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 RiskAssessment
from .base import BaseLLMProvider

if TYPE_CHECKING:
    from puli_mcp_server.vertex_token_manager import VertexTokenManager

logger = logging.getLogger(__name__)


class VertexAnthropicProvider(BaseLLMProvider):
    """Provider for Anthropic Claude 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]]:
        kwargs = {
            "project_id": self.config.vertex_project,
            "region": self.config.vertex_location,
        }
        if self._vertex_token_manager:
            # SA-token path: pass the latest access token to the Anthropic SDK.
            # AsyncAnthropicVertex captures the token at construction time,
            # so we must rebuild the agent when the token refreshes.
            kwargs["access_token"] = self._vertex_token_manager.access_token

        vertex_client = AsyncAnthropicVertex(**kwargs)
        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],
        )

    def needs_rebuild(self) -> bool:
        """Rebuild when using SA tokens (token captured at construction time)."""
        return self._vertex_token_manager is not None

