"""
Gemini client — the Gemini side of llm.py's LLM_BACKEND switch.

Activated when LLM_BACKEND=gemini (auto-selected if GCP_API_KEY is set).
Routes Claude-shaped calls through the google-genai SDK with an API key. Exposes:

- GeminiBedrockClient: .invoke_model() drop-in for boto3 bedrock-runtime.
  Translates Anthropic-style request body to Gemini and shapes the response
  back to the Anthropic JSON the rest of the repo expects.
- GeminiChatModel: langchain ChatBedrock-style .invoke().
- GeminiConverseClient: converse-API wrapper for RLM / tearsheet loops.
- gemini_model_call / gemini_invoke_model: functional helpers used by llm.py.
"""

import base64
import json
import logging
import os
import uuid

logger = logging.getLogger(__name__)

try:
    from utils.token_tracker import record_usage as _record_token_usage
except ImportError:
    try:
        from task.utils.token_tracker import record_usage as _record_token_usage
    except ImportError:
        def _record_token_usage(_model, _input_tokens, _output_tokens):  # type: ignore
            pass


GEMINI_MODEL_ID = os.environ.get("GEMINI_MODEL_ID", "gemini-3.5-flash")
GCP_API_KEY = os.environ.get("GCP_API_KEY") or os.environ.get("GEMINI_API_KEY", "")


_client_init_logged = False
_gemini_client_singleton = None


def _get_gemini_client():
    """Return a google-genai Client authenticated with GCP_API_KEY.

    Lazily constructed and cached: the genai client (and the GCP_API_KEY
    requirement) is only resolved on the FIRST LLM call, not at module import.
    This keeps a missing GCP_API_KEY from crashing the whole worker at boot —
    it surfaces only on the specific request that needs the model.
    """
    global _gemini_client_singleton
    if _gemini_client_singleton is not None:
        return _gemini_client_singleton
    from google import genai
    if not GCP_API_KEY:
        raise RuntimeError(
            "GCP_API_KEY (or GEMINI_API_KEY) is not set. "
            "Either set it or switch LLM_BACKEND."
        )
    global _client_init_logged
    if not _client_init_logged:
        _client_init_logged = True
        logger.info("[gemini_client] init: model=%s (api-key auth)", GEMINI_MODEL_ID)
        print(f"[gemini_client] init: model={GEMINI_MODEL_ID} (api-key auth)", flush=True)
    _gemini_client_singleton = genai.Client(api_key=GCP_API_KEY)
    return _gemini_client_singleton


def _to_gemini_contents(messages: list[dict]):
    """Convert Anthropic-style messages to google-genai Content list.

    Anthropic role "assistant" maps to Gemini role "model". Image blocks
    in Anthropic format ({"type": "image", "source": {"type": "base64",
    "media_type": ..., "data": ...}}) become inline Parts.
    """
    from google.genai import types

    contents = []
    for msg in messages:
        role = "model" if msg.get("role") == "assistant" else "user"
        content = msg.get("content")
        parts = []
        if isinstance(content, str):
            parts.append(types.Part.from_text(text=content))
        elif isinstance(content, list):
            for block in content:
                btype = block.get("type")
                if btype == "text":
                    parts.append(types.Part.from_text(text=block.get("text", "")))
                elif btype == "image":
                    src = block.get("source", {}) or {}
                    if src.get("type") == "base64":
                        parts.append(types.Part.from_bytes(
                            data=base64.b64decode(src.get("data", "")),
                            mime_type=src.get("media_type", "image/png"),
                        ))
                else:
                    # Tearsheet/Bedrock converse-style block — {"text": ...} or {"image": ...}.
                    if "text" in block:
                        parts.append(types.Part.from_text(text=block["text"]))
                    elif "image" in block:
                        img = block["image"]
                        fmt = img.get("format", "png")
                        data = img.get("source", {}).get("bytes", b"")
                        parts.append(types.Part.from_bytes(
                            data=data, mime_type=f"image/{fmt}",
                        ))
        else:
            parts.append(types.Part.from_text(text=str(content)))
        contents.append(types.Content(role=role, parts=parts))
    return contents


def _build_generate_config(*, system=None, max_tokens=4096, temperature=0,
                           use_thinking=False, thinking_budget=None,
                           json_mode=False, response_schema=None,
                           use_search=False):
    """Assemble a GenerateContentConfig with optional thinking + system prompt.

    json_mode=True forces structured JSON output via the Gemini decoder
    (`response_mime_type="application/json"`) — the model cannot emit
    prose / markdown / table rows even if the prompt is sloppy.

    response_schema (optional) binds the output to a specific shape. Pass
    a Python type (`list[str]`, a TypedDict, a pydantic BaseModel) or a
    `types.Schema` dict. Implies json_mode. The decoder will refuse to
    emit anything that doesn't conform.

    use_search=True attaches Gemini's built-in Google Search grounding tool
    so the model searches the live web before answering. OFF by default —
    callers opt in per request (it adds latency + search cost and makes
    output web-dependent). Verified on gemini-2.5-pro to coexist
    with json_mode / response_schema and with thinking, so it composes with
    the other flags here.
    """
    from google.genai import types

    cfg_kwargs = {
        "temperature": temperature,
        "max_output_tokens": max_tokens,
    }
    if system:
        cfg_kwargs["system_instruction"] = system
    if use_thinking or (thinking_budget and thinking_budget > 0):
        budget = thinking_budget if (thinking_budget and thinking_budget > 0) else 5000
        cfg_kwargs["thinking_config"] = types.ThinkingConfig(thinking_budget=budget)
    if json_mode or response_schema is not None:
        cfg_kwargs["response_mime_type"] = "application/json"
    if response_schema is not None:
        cfg_kwargs["response_schema"] = response_schema
    if use_search:
        # Gemini 2.0+ grounding uses GoogleSearch(); serializes to the same
        # {"google_search": {}} tool the vendor-research REST path sends.
        cfg_kwargs["tools"] = [types.Tool(google_search=types.GoogleSearch())]
    return types.GenerateContentConfig(**cfg_kwargs)


def _extract_text(response) -> str:
    """Pull the concatenated text from a google-genai GenerateContentResponse."""
    if getattr(response, "text", None):
        return response.text
    out = []
    for cand in getattr(response, "candidates", []) or []:
        content = getattr(cand, "content", None)
        if not content:
            continue
        for part in getattr(content, "parts", []) or []:
            if getattr(part, "text", None):
                out.append(part.text)
    return "".join(out)


def _extract_usage(response) -> tuple[int, int]:
    """Pull token counts from response.usage_metadata per Gemini API docs.

    Documented fields (https://ai.google.dev/gemini-api/docs/tokens):
      - prompt_token_count        — input tokens
      - candidates_token_count    — output tokens (excluding thinking)
      - thoughts_token_count      — thinking tokens (billed as OUTPUT per
                                    https://ai.google.dev/gemini-api/docs/pricing
                                    "Output price (including thinking tokens)")
      - cached_content_token_count, total_token_count — informational

    Returns (input_tokens, output_tokens) using:
      input  = prompt_token_count
      output = candidates_token_count + thoughts_token_count
    """
    um = getattr(response, "usage_metadata", None)
    if um is None:
        return 0, 0
    in_toks = int(getattr(um, "prompt_token_count", 0) or 0)
    out_toks = (
        int(getattr(um, "candidates_token_count", 0) or 0)
        + int(getattr(um, "thoughts_token_count", 0) or 0)
    )
    return in_toks, out_toks


def _record_usage(response, model_id: str) -> tuple[int, int]:
    """Manual token tracking — read documented fields, record once."""
    in_toks, out_toks = _extract_usage(response)
    if in_toks or out_toks:
        _record_token_usage(model_id, in_toks, out_toks)
    return in_toks, out_toks


def _stream_generate(client, *, model: str, contents, config):
    """Run generate_content_stream and return (full_text, last_chunk).

    Streaming avoids the 10-minute timeout on long requests. The last chunk
    carries the cumulative ``usage_metadata`` for token tracking.
    """
    text_parts: list[str] = []
    last_chunk = None
    for chunk in client.models.generate_content_stream(
        model=model, contents=contents, config=config,
    ):
        last_chunk = chunk
        chunk_text = _extract_text(chunk)
        if chunk_text:
            text_parts.append(chunk_text)
    return "".join(text_parts), last_chunk


class _StreamBody:
    """Mimic boto3 StreamingBody so callers can do response['body'].read()."""

    def __init__(self, data: bytes):
        self._data = data

    def read(self):
        return self._data


class GeminiBedrockClient:
    """Drop-in replacement for boto3 bedrock-runtime client.

    Translates an Anthropic-shaped request body to Gemini, then shapes the
    response back to Anthropic JSON so the rest of the repo (which parses
    response_body['content'][i]['text'] / 'usage') keeps working.

    Token tracking — pure shim: returns the real ``usage`` in the response
    and does NOT record itself. The caller
    parses ``response_body['usage']`` and records via the documented
    ``record_usage`` API (``task/utils/token_tracker.py:218``). Token counts
    are computed from the documented ``usage_metadata`` fields per the
    Gemini docs (``_extract_usage``).
    """

    def __init__(self):
        # Client built lazily on first invoke_model() so a missing GCP_API_KEY
        # does not crash the worker at import/construction time.
        self._model = GEMINI_MODEL_ID

    @property
    def _client(self):
        return _get_gemini_client()

    def invoke_model(self, modelId: str, body: str, **_kwargs):
        logger.info("[gemini_client.GeminiBedrockClient.invoke_model] model=%s (caller passed modelId=%s)", self._model, modelId)
        request = json.loads(body) if isinstance(body, str) else body
        messages = request.get("messages", [])
        max_tokens = request.get("max_tokens", 16384)
        temperature = request.get("temperature", 0)
        system = request.get("system")
        thinking = request.get("thinking") or {}

        use_thinking = thinking.get("type") == "enabled"
        budget = thinking.get("budget_tokens") if use_thinking else None
        # Optional Google Search grounding — set "use_search": true in the
        # request body (non-Anthropic field; absent → off, behavior unchanged).
        use_search = bool(request.get("use_search", False))
        # Unlike Anthropic, Gemini accepts thinking WITH temperature=0 — keep
        # the caller's temperature (default 0) so thinking calls stay
        # deterministic. Never force temperature=1 here.

        contents = _to_gemini_contents(messages)
        config = _build_generate_config(
            system=system, max_tokens=max_tokens, temperature=temperature,
            use_thinking=use_thinking, thinking_budget=budget,
            use_search=use_search,
        )

        text, last_chunk = _stream_generate(
            self._client, model=self._model, contents=contents, config=config,
        )

        in_toks, out_toks = _extract_usage(last_chunk) if last_chunk else (0, 0)

        response_dict = {
            "id": f"gemini-{uuid.uuid4().hex[:24]}",
            "type": "message",
            "role": "assistant",
            "content": [{"type": "text", "text": text}],
            "model": self._model,
            "stop_reason": "end_turn",
            "usage": {
                "input_tokens": in_toks,
                "output_tokens": out_toks,
            },
        }
        return {"body": _StreamBody(json.dumps(response_dict).encode("utf-8"))}


class _InvokeResult:
    def __init__(self, content: str):
        self.content = content


class GeminiChatModel:
    """Drop-in replacement for LangChain ChatBedrock.

    The ``model_id`` argument is accepted for signature compatibility with
    callers that hardcode a Bedrock id (e.g. ``us.anthropic.claude-sonnet-4-6``)
    but is **ignored** — under ``LLM_BACKEND=gemini`` we always call
    ``GEMINI_MODEL_ID``.
    """

    def __init__(self, model_id: str = None, model_kwargs: dict = None):
        del model_id  # accepted for signature compat; we always use GEMINI_MODEL_ID
        # Client built lazily on first invoke() so a missing GCP_API_KEY does
        # not crash the worker at import/construction time.
        self._model = GEMINI_MODEL_ID
        self._kwargs = model_kwargs or {}

    @property
    def _client(self):
        return _get_gemini_client()

    def invoke(self, prompt):
        logger.info("[gemini_client.GeminiChatModel.invoke] model=%s", self._model)
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        else:
            messages = prompt

        config = _build_generate_config(
            max_tokens=self._kwargs.get("max_tokens", 16384),
            temperature=self._kwargs.get("temperature", 0),
            use_search=bool(self._kwargs.get("use_search", False)),
        )
        text, last_chunk = _stream_generate(
            self._client, model=self._model,
            contents=_to_gemini_contents(messages), config=config,
        )
        if last_chunk is not None:
            _record_usage(last_chunk, self._model)
        return _InvokeResult(text)


class GeminiConverseClient:
    """Drop-in replacement for pdf_to_xlsx / tearsheets BedrockClient.

    The ``model_id`` and ``region`` arguments are accepted for signature
    compatibility (RLMBedrockClient callers hardcode Bedrock model ids and
    AWS regions) but are **ignored** — under ``LLM_BACKEND=gemini`` we always
    call ``GEMINI_MODEL_ID``.
    """

    def __init__(
        self,
        model_id: str = None,
        region: str = None,
        max_tokens: int = 16384,
        thinking_budget: int = 4096,
        use_search: bool = False,
    ):
        del model_id, region  # accepted for signature compat; always GEMINI_MODEL_ID
        self.model_id = GEMINI_MODEL_ID
        self.max_tokens = max_tokens
        self.thinking_budget = thinking_budget if (thinking_budget and thinking_budget > 0) else 0
        self.use_search = use_search
        # Client built lazily on first completion() so a missing GCP_API_KEY
        # does not crash the worker at import/construction time.
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.call_count = 0

    @property
    def _client(self):
        return _get_gemini_client()

    def _generate(self, messages: list[dict], system: str | None) -> str:
        logger.info(
            "[gemini_client.GeminiConverseClient] model=%s thinking=%s call#%d",
            self.model_id, bool(self.thinking_budget), self.call_count + 1,
        )
        contents = _to_gemini_contents(messages)
        config = _build_generate_config(
            system=system,
            max_tokens=self.max_tokens,
            temperature=0,  # Gemini allows thinking at temp 0 — deterministic
            use_thinking=bool(self.thinking_budget),
            thinking_budget=self.thinking_budget or None,
            use_search=self.use_search,
        )
        text, last_chunk = _stream_generate(
            self._client, model=self.model_id, contents=contents, config=config,
        )
        if last_chunk is not None:
            in_toks, out_toks = _record_usage(last_chunk, self.model_id)
            self.total_input_tokens += in_toks
            self.total_output_tokens += out_toks
        self.call_count += 1
        return text

    def completion(self, messages: list[dict], system: str | None = None) -> str:
        return self._generate(messages, system)

    def completion_with_image(
        self,
        image_bytes: bytes,
        mime_type: str,
        prompt: str,
        system: str | None = None,
    ) -> str:
        media_type = mime_type if "/" in mime_type else f"image/{mime_type}"
        b64 = base64.standard_b64encode(image_bytes).decode("utf-8")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "source": {
                    "type": "base64", "media_type": media_type, "data": b64,
                }},
                {"type": "text", "text": prompt},
            ],
        }]
        return self._generate(messages, system)

    def usage_summary(self) -> str:
        return (
            f"Calls: {self.call_count} | "
            f"Input: {self.total_input_tokens:,} | "
            f"Output: {self.total_output_tokens:,}"
        )


def gemini_model_call(prompt, max_tokens=16384, use_thinking=False, system=None,
                      json_mode=False, response_schema=None, use_search=False):
    """Single-shot Gemini call. Returns the text response.

    use_search=True grounds the call with Google Search (off by default)."""
    logger.info("[gemini_client.gemini_model_call] model=%s thinking=%s json_mode=%s schema=%s search=%s", GEMINI_MODEL_ID, use_thinking, json_mode, response_schema is not None, use_search)
    client = _get_gemini_client()
    config = _build_generate_config(
        system=system, max_tokens=max_tokens,
        temperature=0,  # Gemini allows thinking at temp 0 — deterministic
        use_thinking=use_thinking,
        json_mode=json_mode,
        response_schema=response_schema,
        use_search=use_search,
    )
    contents = _to_gemini_contents([{"role": "user", "content": prompt}])
    try:
        text, last_chunk = _stream_generate(
            client, model=GEMINI_MODEL_ID, contents=contents, config=config,
        )
        if last_chunk is not None:
            _record_usage(last_chunk, GEMINI_MODEL_ID)
        return text
    except Exception as e:
        logger.error(f"Gemini model call failed: {e}")
        raise


def gemini_invoke_model(messages, system=None, max_tokens=16384, temperature=0,
                        use_thinking=False, use_search=False):
    """Multi-turn Gemini call. Returns the text response.

    use_search=True grounds the call with Google Search (off by default)."""
    logger.info("[gemini_client.gemini_invoke_model] model=%s msgs=%d thinking=%s search=%s", GEMINI_MODEL_ID, len(messages), use_thinking, use_search)
    client = _get_gemini_client()
    config = _build_generate_config(
        system=system, max_tokens=max_tokens,
        temperature=temperature,  # Gemini allows thinking at temp 0 — no forcing
        use_thinking=use_thinking,
        use_search=use_search,
    )
    try:
        text, last_chunk = _stream_generate(
            client, model=GEMINI_MODEL_ID,
            contents=_to_gemini_contents(messages), config=config,
        )
        if last_chunk is not None:
            _record_usage(last_chunk, GEMINI_MODEL_ID)
        return text
    except Exception as e:
        logger.error(f"Gemini invoke_model failed: {e}")
        raise