"""Chamadas ao Gemini (API paga: os dados nao sao usados para treino)."""
import json
import os
import time

from google import genai
from google.genai import errors, types

from .banco import agora
from .config import carregar_env

_cliente = None


def cliente():
    global _cliente
    if _cliente is None:
        carregar_env()
        _cliente = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    return _cliente


def _registrar(con, etapa, modelo, resp):
    u = resp.usage_metadata
    if con is not None and u:
        con.execute("INSERT INTO uso_llm VALUES (?, ?, ?, ?, ?)",
                    (agora(), etapa, modelo, u.prompt_token_count or 0,
                     (u.candidates_token_count or 0) + (u.thoughts_token_count or 0)))
        con.commit()


def gerar(modelo, conteudo, *, sistema=None, esquema=None, temperatura=0.2, max_saida=8192,
          con=None, etapa="geral"):
    """Gera texto (ou JSON, se `esquema` for um JSON Schema). Tenta de novo em erros temporarios."""
    cfg = types.GenerateContentConfig(
        system_instruction=sistema, temperature=temperatura, max_output_tokens=max_saida,
        **({"response_mime_type": "application/json", "response_json_schema": esquema} if esquema else {}))
    for tentativa in range(5):
        try:
            resp = cliente().models.generate_content(model=modelo, contents=conteudo, config=cfg)
            _registrar(con, etapa, modelo, resp)
            texto = resp.text or ""
            return json.loads(texto) if esquema else texto
        except errors.APIError as e:
            if e.code in (429, 500, 502, 503, 504) and tentativa < 4:
                time.sleep(5 * 2 ** tentativa)
                continue
            raise
        except json.JSONDecodeError:
            if tentativa < 2:
                continue
            raise


def parte_bytes(dados: bytes, mime: str):
    return types.Part.from_bytes(data=dados, mime_type=mime)