import os
from typing import TypeVar

from google import genai
from google.genai import errors as genai_errors
from google.genai import types
from pydantic import BaseModel

from ..exceptions import OcrError, OcrFailureCode

DEFAULT_GEMINI_MODEL = "gemini-2.5-flash"
ModelT = TypeVar("ModelT", bound=BaseModel)

# Gemini's 4xx codes that mean "the request/credentials were bad", not "the file
# was bad" -- everything else in ClientError (e.g. 400 for an unreadable/corrupt
# image) is treated as an unsupported file below.
_GEMINI_CONFIGURATION_ERROR_CODES = {401, 403, 404}


def _get_gemini_client() -> genai.Client:
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise OcrError(OcrFailureCode.CONFIGURATION_ERROR, "GEMINI_API_KEY is not set in the environment")
    return genai.Client(api_key=api_key)


def generate_structured_content(
    *,
    file_path: str,
    mime_type: str,
    prompt: str,
    response_schema: type[ModelT],
    model: str = DEFAULT_GEMINI_MODEL,
    thinking_budget: int | None = None,
) -> ModelT:
    """Sends an image + prompt to Gemini and parses the JSON response into response_schema.

    thinking_budget caps the model's thinking tokens (in the .usage_metadata sense); left
    as None, the model's own default thinking behavior applies.
    """
    client = _get_gemini_client()

    with open(file_path, "rb") as f:
        image_bytes = f.read()

    thinking_config = (
        types.ThinkingConfig(thinking_budget=thinking_budget) if thinking_budget is not None else None
    )

    try:
        response = client.models.generate_content(
            model=model,
            contents=[types.Part.from_bytes(data=image_bytes, mime_type=mime_type), prompt],
            config=types.GenerateContentConfig(
                response_mime_type="application/json",
                response_schema=response_schema,
                media_resolution=types.MediaResolution.MEDIA_RESOLUTION_HIGH,
                thinking_config=thinking_config,
            ),
        )
    except genai_errors.ClientError as exc:
        code = (
            OcrFailureCode.CONFIGURATION_ERROR
            if exc.code in _GEMINI_CONFIGURATION_ERROR_CODES
            else OcrFailureCode.UNSUPPORTED_FILE
        )
        raise OcrError(code, f"Gemini rejected the request: {exc.message or exc}") from exc
    except genai_errors.ServerError as exc:
        raise OcrError(OcrFailureCode.PROVIDER_ERROR, f"Gemini call failed: {exc.message or exc}") from exc

    usage = response.usage_metadata
    if usage is not None:
        print(
            f"[gemini] model={model} prompt_tokens={usage.prompt_token_count} "
            f"candidates_tokens={usage.candidates_token_count} total_tokens={usage.total_token_count}"
        )

    if response.parsed is None:
        raise OcrError(
            OcrFailureCode.NO_DATA_EXTRACTED,
            f"Gemini response did not match schema {response_schema.__name__}. Raw text: {response.text!r}",
        )
    return response.parsed
