from __future__ import annotations

import os

from pydantic import BaseModel

from safe_genai.exceptions import SafeGenAIProviderError
from safe_genai.providers.base import BaseProvider
from safe_genai.result import Usage


def _extract_usage(response: object) -> Usage:
    meta = getattr(response, "usage_metadata", None)
    if meta is None:
        return Usage()
    return Usage(
        input_tokens=getattr(meta, "prompt_token_count", 0) or 0,
        output_tokens=getattr(meta, "candidates_token_count", 0) or 0,
    )


class GoogleGenAIProvider(BaseProvider):
    def __init__(
        self,
        api_key: str | None = None,
        *,
        vertexai: bool = False,
        vertex_project: str | None = None,
        vertex_location: str = "us-central1",
    ) -> None:
        try:
            import google.genai as genai
            import google.genai.types as genai_types
        except ImportError as e:
            raise ImportError("google-genai package is required: pip install google-genai") from e

        if vertexai:
            resolved_project = vertex_project or os.environ.get("GOOGLE_CLOUD_PROJECT")
            if not resolved_project:
                raise ValueError(
                    "vertex_project is required for Vertex AI. "
                    "Pass it explicitly or set the GOOGLE_CLOUD_PROJECT environment variable."
                )
            self._client = genai.Client(
                vertexai=True,
                project=resolved_project,
                location=vertex_location,
            )
        else:
            resolved_key = api_key or os.environ.get("GEMINI_API_KEY")
            if not resolved_key:
                raise ValueError(
                    "GEMINI_API_KEY environment variable is not set and no api_key was provided. "
                    "If you want to use Vertex AI with gcloud ADC, set vertexai=True."
                )
            self._client = genai.Client(api_key=resolved_key)

        self._genai = genai
        self._genai_types = genai_types

    def generate_text(
        self,
        model: str,
        system_prompt: str,
        user_prompt: str,
        *,
        enable_grounding: bool = False,
    ) -> tuple[str, Usage]:
        try:
            config = self._build_text_config(system_prompt, enable_grounding)
            response = self._client.models.generate_content(
                model=model,
                contents=user_prompt,
                config=config,
            )
            return response.text or "", _extract_usage(response)
        except Exception as e:
            raise SafeGenAIProviderError("GoogleGenAI", e) from e

    def generate_structured(
        self,
        model: str,
        system_prompt: str,
        user_prompt: str,
        output_schema: type[BaseModel],
    ) -> tuple[BaseModel, Usage]:
        try:
            config = self._genai_types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=output_schema,
            )
            response = self._client.models.generate_content(
                model=model,
                contents=user_prompt,
                config=config,
            )
            return output_schema.model_validate_json(response.text), _extract_usage(response)
        except Exception as e:
            raise SafeGenAIProviderError("GoogleGenAI", e) from e

    async def agenerate_text(
        self,
        model: str,
        system_prompt: str,
        user_prompt: str,
        *,
        enable_grounding: bool = False,
    ) -> tuple[str, Usage]:
        try:
            config = self._build_text_config(system_prompt, enable_grounding)
            response = await self._client.aio.models.generate_content(
                model=model,
                contents=user_prompt,
                config=config,
            )
            return response.text or "", _extract_usage(response)
        except Exception as e:
            raise SafeGenAIProviderError("GoogleGenAI", e) from e

    async def agenerate_structured(
        self,
        model: str,
        system_prompt: str,
        user_prompt: str,
        output_schema: type[BaseModel],
    ) -> tuple[BaseModel, Usage]:
        try:
            config = self._genai_types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=output_schema,
            )
            response = await self._client.aio.models.generate_content(
                model=model,
                contents=user_prompt,
                config=config,
            )
            return output_schema.model_validate_json(response.text), _extract_usage(response)
        except Exception as e:
            raise SafeGenAIProviderError("GoogleGenAI", e) from e

    def _build_text_config(self, system_prompt: str, enable_grounding: bool) -> object:
        tools = (
            [self._genai_types.Tool(google_search=self._genai_types.GoogleSearch())]
            if enable_grounding
            else None
        )
        return self._genai_types.GenerateContentConfig(
            system_instruction=system_prompt,
            tools=tools,
        )
