"""Google Gemini adapter (google-genai): JSON output constrained by response_json_schema."""
from __future__ import annotations
import time

from pydantic import ValidationError

from .base import ImagePart, ProviderError, Result, SchemaError


class GoogleProvider:
    name = "google"

    def __init__(self, api_key: str | None = None):
        from google import genai
        from google.genai import errors, types
        self._types, self._errors = types, errors
        self.client = genai.Client(api_key=api_key)

    def structured(self, *, model, system, parts, schema, max_tokens=8000, cache_system=True) -> Result:
        t = self._types
        contents = [t.Part.from_bytes(data=p.data, mime_type=p.media_type) if isinstance(p, ImagePart)
                    else t.Part.from_text(text=p.text) for p in parts]
        cfg = t.GenerateContentConfig(system_instruction=system, max_output_tokens=max_tokens,
                                      temperature=0, response_mime_type="application/json",
                                      response_json_schema=schema.model_json_schema())
        t0 = time.monotonic()
        try:
            resp = self.client.models.generate_content(model=model, contents=contents, config=cfg)
        except self._errors.APIError as e:
            code = getattr(e, "code", 0) or 0
            raise ProviderError(str(e), retryable=code == 429 or code >= 500) from e
        raw = resp.text or ""
        try:
            parsed = schema.model_validate_json(raw)
        except ValidationError as e:
            raise SchemaError(str(e), raw) from e
        u = resp.usage_metadata
        return Result(parsed=parsed, raw_text=raw, provider=self.name, model=model,
                      tokens_in=u.prompt_token_count or 0, tokens_out=u.candidates_token_count or 0,
                      cached_tokens=u.cached_content_token_count or 0,
                      latency_ms=int((time.monotonic() - t0) * 1000))