"""
Google Cloud Vertex AI provider implementation.

Uses Vertex AI for model access and GCP Billing API for cost tracking.
"""

import logging
import os
from datetime import UTC, datetime, timedelta
from typing import Any

from openscientist.exceptions import ProviderError
from openscientist.providers.base import AirgapEgress, AirgapPosture, ClaudeCompatible, CostInfo
from openscientist.settings import get_settings

from ._anthropic_common import (
    send_anthropic_message,
    send_anthropic_message_with_tools,
)
from ._env_cleanup import clear_env_vars, clear_provider_mode_flags

logger = logging.getLogger(__name__)


class VertexProvider(ClaudeCompatible):
    """Google Cloud Vertex AI provider."""

    @property
    def id(self) -> str:
        return "vertex"

    @property
    def display_name(self) -> str:
        return "Vertex AI"

    def validate_required_config(self) -> list[str]:
        """Check required Vertex AI configuration."""
        errors = []
        settings = get_settings()

        if not settings.provider.anthropic_vertex_project_id:
            errors.append("ANTHROPIC_VERTEX_PROJECT_ID not set (GCP project ID)")

        creds_env = settings.provider.google_application_credentials
        if not creds_env:
            errors.append("GOOGLE_APPLICATION_CREDENTIALS not set (path to service account JSON)")
        else:
            # Check if file exists
            creds_path = os.path.expanduser(creds_env)
            if not os.path.exists(creds_path):
                errors.append(f"GOOGLE_APPLICATION_CREDENTIALS file not found: {creds_path}")

        if not settings.provider.gcp_billing_account_id:
            errors.append("GCP_BILLING_ACCOUNT_ID not set (needed for cost tracking)")

        if not settings.provider.cloud_ml_region:
            errors.append("CLOUD_ML_REGION not set (e.g., us-east5)")

        return errors

    def claude_sdk_env(self) -> dict[str, str]:
        """Vertex routing/auth env vars the claude-agent-sdk CLI must see.

        Mirrors `ProviderSettings._apply_vertex_env_vars` so the explicit
        per-provider env stays consistent with the container-env builder.
        """
        p = get_settings().provider
        env: dict[str, str] = {"CLAUDE_CODE_USE_VERTEX": "1"}
        for key, value in (
            ("ANTHROPIC_VERTEX_PROJECT_ID", p.anthropic_vertex_project_id),
            ("GCP_BILLING_ACCOUNT_ID", p.gcp_billing_account_id),
            ("CLOUD_ML_REGION", p.cloud_ml_region),
            ("VERTEX_REGION_CLAUDE_4_5_SONNET", p.vertex_region_claude_4_5_sonnet),
            ("VERTEX_REGION_CLAUDE_4_5_HAIKU", p.vertex_region_claude_4_5_haiku),
            ("GOOGLE_APPLICATION_CREDENTIALS", p.google_application_credentials),
        ):
            if value:
                env[key] = value
        return env

    def claude_model_name(self) -> str:
        """Model name for ClaudeAgentOptions.model."""
        return get_settings().provider.model or "claude-sonnet-4-5@20250929"

    def airgap_egress(self) -> AirgapPosture:
        region = get_settings().provider.cloud_ml_region
        return AirgapPosture(
            AirgapEgress.DIRECT,
            direct_endpoints=(
                (f"{region}-aiplatform.googleapis.com", 443),
                ("oauth2.googleapis.com", 443),
            ),
        )

    def _validate_optional_config(self) -> list[str]:
        """Check optional Vertex AI configuration."""
        warnings = []
        settings = get_settings()

        if not settings.provider.model:
            warnings.append("OPENSCIENTIST_MODEL not set (will use claude-sonnet-4-5@20250929)")

        if not settings.provider.vertex_region_claude_4_5_sonnet:
            warnings.append("VERTEX_REGION_CLAUDE_4_5_SONNET not set (may cause region issues)")

        if not settings.provider.vertex_region_claude_4_5_haiku:
            warnings.append("VERTEX_REGION_CLAUDE_4_5_HAIKU not set (may cause region issues)")

        return warnings

    def setup_environment(self) -> None:
        """Vertex AI environment should be configured via .env and docker-compose.yml."""
        # Keep existing Vertex mode flag (if configured) and clear conflicting ones.
        clear_provider_mode_flags(logger, active_flag="CLAUDE_CODE_USE_VERTEX")
        clear_env_vars(logger, ("ANTHROPIC_API_KEY",))
        logger.info("Vertex AI provider initialized (configuration from environment)")

    def get_cost_info(self, lookback_hours: int = 24) -> CostInfo:
        """
        Get Vertex AI cost information from Cloud Billing API.

        Args:
            lookback_hours: Time window for recent spend calculation

        Returns:
            CostInfo with Vertex AI spend data

        Note:
            Requires BigQuery billing export to be enabled in GCP Console.
            Data has 1-6 hour lag due to GCP billing pipeline.
        """
        try:
            from google.cloud import bigquery
            from google.oauth2 import service_account
        except ImportError:
            logger.error(
                "google-cloud-bigquery not installed. "
                "Install with: pip install google-cloud-bigquery"
            )
            raise

        # Load credentials
        settings = get_settings()
        creds_env = settings.provider.google_application_credentials
        if not creds_env:
            raise ProviderError("GOOGLE_APPLICATION_CREDENTIALS not set")
        creds_path = os.path.expanduser(creds_env)
        credentials = service_account.Credentials.from_service_account_file(creds_path)  # type: ignore[no-untyped-call]

        project_id = settings.provider.anthropic_vertex_project_id
        billing_account = settings.provider.gcp_billing_account_id
        if not billing_account:
            raise ProviderError("GCP_BILLING_ACCOUNT_ID not set")

        # Initialize BigQuery client
        bq_client = bigquery.Client(credentials=credentials, project=project_id)

        # Calculate time windows
        now = datetime.now(UTC)
        recent_start = now - timedelta(hours=lookback_hours)

        # BigQuery billing export table name
        # Format: billing_export.gcp_billing_export_v1_{billing_account_with_underscores}
        billing_table = (
            f"{project_id}.billing_export.gcp_billing_export_v1_{billing_account.replace('-', '_')}"
        )

        # Query for total spend (all time for Vertex AI in this project only)
        total_query = f"""
        SELECT SUM(cost) as total_cost
        FROM `{billing_table}`
        WHERE service.description = 'Vertex AI'
          AND project.id = '{project_id}'
        """

        # Query for recent spend (filter by project AND time window)
        recent_query = f"""
        SELECT SUM(cost) as recent_cost
        FROM `{billing_table}`
        WHERE service.description = 'Vertex AI'
          AND project.id = '{project_id}'
          AND usage_start_time >= TIMESTAMP('{recent_start.isoformat()}')
        """

        try:
            # Execute queries
            total_result = list(bq_client.query(total_query).result())
            total_spend = float(total_result[0].total_cost or 0)

            recent_result = list(bq_client.query(recent_query).result())
            recent_spend = float(recent_result[0].recent_cost or 0)

            # Estimate data lag (GCP billing typically has 1-6 hour delay)
            # Assume ~3 hour average lag
            lag_time = now - timedelta(hours=3)
            data_lag_note = f"Data current as of ~{lag_time.strftime('%I:%M %p %Z')}"

        except Exception as e:
            logger.warning("Could not fetch Vertex AI billing data: %s", e)
            logger.warning(
                "Ensure BigQuery billing export is enabled in GCP Console. "
                "See docs/VERTEX_SETUP.md for setup instructions."
            )
            # Return None if billing data unavailable (permissions, export not enabled, etc.)
            total_spend = None
            recent_spend = None
            data_lag_note = "Billing data unavailable (ensure BigQuery export is enabled)"

        return CostInfo(
            provider_name="Vertex AI",
            total_spend_usd=total_spend,
            recent_spend_usd=recent_spend,
            recent_period_hours=lookback_hours,
            last_updated=now,
            data_lag_note=data_lag_note,
            metadata={"project_id": project_id},
        )

    async def send_message(
        self,
        messages: list[dict[str, str]],
        system: str | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
    ) -> str:
        """
        Send message using Vertex AI via Anthropic SDK.

        This bypasses the Claude Code CLI and its local pre-flight content
        filter, which can produce false positives on legitimate scientific content.
        """
        import anthropic

        settings = get_settings()

        # These are validated as required in _validate_required_config
        project_id = settings.provider.anthropic_vertex_project_id
        region = settings.provider.cloud_ml_region
        if not project_id or not region:
            raise ValueError("Vertex AI project_id and region are required")

        client = anthropic.AnthropicVertex(
            project_id=project_id,
            region=region,
        )
        return send_anthropic_message(
            client=client,
            messages=messages,
            system=system,
            model=model,
            configured_model=settings.provider.model,
            provider_default_model="claude-sonnet-4-5@20250929",
            max_tokens=max_tokens,
        )

    async def send_message_with_tools(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        system: str | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
    ) -> dict[str, Any]:
        """
        Send message with tool definitions using Vertex AI via Anthropic SDK.

        Returns full response including stop_reason and content blocks.
        """
        import anthropic
        from anthropic.types import ToolUseBlock

        settings = get_settings()

        # These are validated as required in _validate_required_config
        project_id = settings.provider.anthropic_vertex_project_id
        region = settings.provider.cloud_ml_region
        if not project_id or not region:
            raise ValueError("Vertex AI project_id and region are required")

        client = anthropic.AnthropicVertex(
            project_id=project_id,
            region=region,
        )
        return send_anthropic_message_with_tools(
            client=client,
            messages=messages,
            tools=tools,
            system=system,
            model=model,
            configured_model=settings.provider.model,
            provider_default_model="claude-sonnet-4-5@20250929",
            max_tokens=max_tokens,
            tool_use_block_type=ToolUseBlock,
        )
