"""
ai.py — Visao + raciocinio/estruturacao de texto, ambas via NVIDIA NIM (ou HF como alternativa).

Arquitetura (cada modelo individualmente < 32B — teto e por modelo, nao somado):
  - Visao   : meta/llama-3.2-11b-vision-instruct    (11B, NVIDIA NIM)  -> NVIDIA_API_KEY
              (ou meta-llama/Llama-3.2-11B-Vision-Instruct via HF, se so houver HF_TOKEN)
  - Cerebro : nvidia/nemotron-3-nano-30b-a3b         (MoE, 3B ATIVOS,  NVIDIA NIM) -> NVIDIA_API_KEY
  - Embed   : sentence-transformers/all-MiniLM-L6-v2 (22M, local/CPU)

O Nemotron 3 Nano e o nucleo do app: alem de extrair o JSON estrito a partir da
descricao livre da visao (especie, raca da lista permitida, cor, condicao, frase de
matching), ele tambem REDIGE a "historia" de cada animal exibida no perfil
(generate_story). E um modelo de reasoning MoE (30B totais / 3B ativos por chamada,
256k de contexto) servido na API hospedada da NVIDIA — so 3B ativos o tornam rapido
e barato, e o total de 30B fica abaixo do teto de 32B por modelo do hackathon.

Fluxo: o Llama-3.2-11B-Vision enxerga a foto e descreve o animal (texto livre —
modelos de visao raramente cospem JSON limpo). Em seguida o Nemotron 3 Nano LE essa
descricao, raciocina e estrutura. Cada modelo no que e bom: Llama ve, Nemotron pensa.

Basta a NVIDIA_API_KEY: ela serve tanto para a visao quanto para o Nemotron.
Sem NVIDIA_API_KEY mas com HF_TOKEN, a visao usa o HF Serverless e a estruturacao
cai para o parser local (regex de JSON).

Variaveis de ambiente:
  NVIDIA_API_KEY      — chave NVIDIA NIM (free tier em build.nvidia.com)
  HF_TOKEN            — token HuggingFace (alternativa de visao via HF Serverless)
  NVIDIA_VISION_MODEL — override do modelo de visao NIM (default meta/llama-3.2-11b-vision-instruct)
  NVIDIA_MODEL        — override do cerebro de texto NIM (default nvidia/nemotron-3-nano-30b-a3b;
                        aponte para um Nemotron 3 Nano 4B self-hosted se rodar local)
  HF_VISION_MODEL     — override do modelo de visao HF (default meta-llama/Llama-3.2-11B-Vision-Instruct)
"""
import base64
import io
import json
import logging
import os
import re

import numpy as np

log = logging.getLogger(__name__)

_HF_MODEL         = "meta-llama/Llama-3.2-11B-Vision-Instruct"  # id HF Serverless
_NIM_VISION_MODEL = "meta/llama-3.2-11b-vision-instruct"        # id NVIDIA NIM (11B, free)
_NIM_TEXT_MODEL   = "nvidia/nemotron-3-nano-30b-a3b"            # id NVIDIA NIM (Nemotron 3 Nano MoE, 3B ativos)

_DOG_BREEDS = (
    "SRD, Labrador Retriever, Golden Retriever, German Shepherd, Bulldog, French Bulldog, "
    "Poodle, Beagle, Rottweiler, Pitbull, American Staffordshire Terrier, Dachshund, "
    "Yorkshire Terrier, Boxer, Chihuahua, Shih Tzu, Siberian Husky, Border Collie, "
    "Australian Shepherd, Cavalier King Charles Spaniel, Doberman, Great Dane, Maltese, "
    "Pug, Pomeranian, Bernese Mountain Dog, Cocker Spaniel, Shiba Inu, Akita, Chow Chow, "
    "Dalmatian, Bichon Frise, Mastiff, Cane Corso, Saint Bernard, Weimaraner, Basset Hound, "
    "Newfoundland, Bloodhound, Whippet, Greyhound, Vizsla, Samoyed, Jack Russell Terrier, "
    "Australian Cattle Dog, English Setter, Pekingese, Lhasa Apso, Schnauzer, "
    "Belgian Malinois, Rhodesian Ridgeback, Brazilian Terrier, Fila Brasileiro, Other"
)
_CAT_BREEDS = (
    "Domestic Shorthair, Domestic Longhair, Siamese, Persian, Maine Coon, Bengal, "
    "British Shorthair, Ragdoll, Scottish Fold, Turkish Angora, Sphynx, Abyssinian, "
    "Russian Blue, Norwegian Forest Cat, Birman, Oriental Shorthair, Devon Rex, Cornish Rex, "
    "American Shorthair, Exotic Shorthair, Burmese, Tonkinese, Himalayan, Manx, Savannah, "
    "Munchkin, Egyptian Mau, Somali, Balinese, Ocicat, Singapura, Chartreux, Selkirk Rex, "
    "Turkish Van, Bombay, Other"
)
_COLORS = (
    "Black, White, Gray, Brown, Chocolate, Caramel, Tan, Golden, Cream, Fawn, Orange, "
    "Ginger, Red, Sable, Cinnamon, Blue, Lilac, Tabby, Tuxedo, Calico, Tortoiseshell, "
    "Tricolor, Bicolor, Brindle, Merle, Spotted, Mixed"
)
_SIZES = "Tiny, Small, Medium, Large, Extra Large, Giant"

# Prompt da VISAO — texto livre permitido (o Nemotron estrutura depois).
PROMPT = (
    "Look at this image. Is there a dog or a cat in it?\n"
    "If there is NO dog and NO cat, reply with exactly: NO ANIMAL\n"
    "If there is a dog or a cat, describe the most prominent one in 2-3 sentences: "
    "species (dog or cat), likely breed, size, main color and secondary colors, "
    "any distinctive marks (collar, scar, patch, missing ear), and apparent condition "
    "(healthy, thin or injured)."
)

# Prompt do NEMOTRON — extrai JSON estrito a partir da descricao livre do Llama.
EXTRACT_PROMPT = (
    "You convert a free-text animal description into STRICT JSON for a stray-animal database.\n"
    "Output ONLY a JSON object — no markdown, no explanation.\n\n"
    "If the description says there is no dog or cat (e.g. 'NO ANIMAL'), output exactly:\n"
    "{\"is_animal\": false}\n\n"
    "Otherwise output exactly these keys:\n"
    "{\"is_animal\": true,"
    " \"species\": \"dog or cat\","
    " \"breed_estimate\": \"single best match from the breed lists below\","
    " \"size\": \"single best match from the size list below\","
    " \"primary_color\": \"single best match from the color list below\","
    " \"secondary_colors\": [\"other colors from the color list, or empty list\"],"
    " \"distinctive_marks\": [\"notable features or empty list\"],"
    " \"condition\": \"healthy, thin or injured\","
    " \"description_text\": \"one concise English sentence identifying this specific animal\"}\n\n"
    "All values in English, matching EXACTLY one option from these lists:\n"
    "DOG breeds: " + _DOG_BREEDS + "\n"
    "CAT breeds: " + _CAT_BREEDS + "\n"
    "COLORS: " + _COLORS + "\n"
    "SIZES: " + _SIZES + "\n"
    "PRIMARY COLOR RULE: if the coat shows a recognizable PATTERN "
    "(Calico, Tortoiseshell, Tabby, Tuxedo, Bicolor, Tricolor, Brindle, Merle, Spotted), "
    "set primary_color to that pattern and put the actual solid colors (e.g. White, Black, Orange) "
    "in secondary_colors. Use a plain solid color as primary_color ONLY when there is no clear pattern.\n"
    "Use SRD (dog) or Domestic Shorthair (cat) only if the breed is truly unclear.\n"
    "If several animals appear, structure only the most prominent one.\n"
    "Do NOT invent details that are not in the description.\n\n"
    "DESCRIPTION:\n"
)


class AnimalAI:
    def __init__(self):
        self.vision_mode    = None   # "nim" | "hf"
        self.vision_model   = None
        self.vision_client  = None   # OpenAI (NIM) ou InferenceClient (HF)
        self.nim_text_model = _NIM_TEXT_MODEL
        self.nim_client     = None   # OpenAI -> NVIDIA NIM, p/ extracao Nemotron

        nvidia_key = os.environ.get("NVIDIA_API_KEY", "")
        hf_token   = os.environ.get("HF_TOKEN", "")

        # Cliente NVIDIA NIM (serve para visao Llama E texto Nemotron)
        if nvidia_key:
            try:
                from openai import OpenAI
                self.nim_client = OpenAI(
                    base_url="https://integrate.api.nvidia.com/v1",
                    api_key=nvidia_key,
                )
                self.nim_text_model = os.environ.get("NVIDIA_MODEL", _NIM_TEXT_MODEL)
                self.vision_mode   = "nim"
                self.vision_model  = os.environ.get("NVIDIA_VISION_MODEL", _NIM_VISION_MODEL)
                self.vision_client = self.nim_client
                log.info("Visao: %s + Texto: %s via NVIDIA NIM", self.vision_model, self.nim_text_model)
            except ImportError:
                log.warning("openai nao instalado")

        # Alternativa: visao via HF Serverless se nao houver NVIDIA
        if self.vision_client is None and hf_token:
            try:
                from huggingface_hub import InferenceClient
                self.vision_mode   = "hf"
                self.vision_model  = os.environ.get("HF_VISION_MODEL", _HF_MODEL)
                self.vision_client = InferenceClient(model=self.vision_model, token=hf_token)
                log.info("Visao: %s via HF InferenceClient", self.vision_model)
            except ImportError:
                log.warning("huggingface_hub nao instalado")

        if self.vision_client is None:
            log.warning("Sem NVIDIA_API_KEY nem HF_TOKEN — IA desabilitada. Configure um dos Secrets.")

        self.embedder = None
        try:
            from sentence_transformers import SentenceTransformer
            self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
            log.info("sentence-transformers: all-MiniLM-L6-v2")
        except Exception as e:
            log.warning("sentence-transformers nao carregou: %s", e)

    def analyze_image(self, image) -> dict:
        """Analisa imagem. _ai_success=False indica que a IA nao foi usada."""
        if self.vision_client is None:
            return self._fallback()

        try:
            img_b64 = self._to_b64(image)
            content = [
                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + img_b64}},
                {"type": "text", "text": PROMPT},
            ]

            if self.vision_mode == "nim":
                resp = self.vision_client.chat.completions.create(
                    model=self.vision_model,
                    messages=[{"role": "user", "content": content}],
                    max_tokens=400,
                    temperature=0.2,
                )
                raw = resp.choices[0].message.content or ""
            else:
                resp = self.vision_client.chat_completion(
                    messages=[{"role": "user", "content": content}],
                    max_tokens=400,
                    temperature=0.2,
                )
                raw = resp.choices[0].message.content or ""

            log.info("Visao resposta: %s", raw[:200])

            # Atalho: visao declarou que nao ha animal
            if "NO ANIMAL" in raw.upper() and self._extract_json(raw) is None:
                log.info("IA: nenhum animal detectado na imagem.")
                return {"is_animal": False, "_ai_success": True}

            # Estrutura a descricao em JSON (Nemotron preferido; robusto a prosa)
            if self.nim_client is not None:
                result = self._structure_with_nemotron(raw)
            else:
                result = self._parse(raw)

            if result.get("is_animal") is False:
                log.info("IA: nenhum animal detectado (estruturacao).")
                return {"is_animal": False, "_ai_success": True}

            result["is_animal"] = True
            result.setdefault("_ai_success", True)
            return result

        except Exception as e:
            log.error("Erro na API de visao: %s", e)
            return self._fallback()

    def _nemotron_chat(self, messages, max_tokens, temperature=0.2, top_p=0.9):
        """Chama o Nemotron 3 Nano DESLIGANDO o reasoning (enable_thinking=false).

        O modelo liga o thinking por padrao e emite <think>...</think>, gastando
        os tokens antes de chegar na resposta. O toggle correto vai no corpo da
        requisicao. Tentamos os formatos conhecidos e, se o endpoint recusar o
        campo, repetimos sem ele (a extracao tolera <think> de qualquer forma).
        """
        attempts = [
            {"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}},
            {"extra_body": {"enable_thinking": False}},
            {},
        ]
        last_err = None
        for kw in attempts:
            try:
                return self.nim_client.chat.completions.create(
                    model=self.nim_text_model,
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    **kw,
                )
            except Exception as e:
                last_err = e
                continue
        raise last_err

    def _structure_with_nemotron(self, raw_text: str) -> dict:
        """Extrai JSON estrito da descricao livre da visao usando Nemotron 3 Nano."""
        try:
            resp = self._nemotron_chat(
                messages=[
                    {"role": "user", "content": EXTRACT_PROMPT + raw_text +
                     "\n\nReply with ONLY the JSON object, nothing else."},
                ],
                max_tokens=2048,   # folga extra caso o reasoning ainda venha ligado
                temperature=0.2,
                top_p=0.9,
            )
            msg = resp.choices[0].message
            # Tenta o content (resposta final) primeiro; so depois o reasoning_content,
            # p/ nao pegar um JSON "ecoado" no meio do raciocinio.
            parsed = self._extract_json(msg.content or "")
            if parsed is None:
                parsed = self._extract_json(getattr(msg, "reasoning_content", "") or "")
            if parsed is not None:
                log.info("Nemotron 3 Nano estruturou o JSON.")
                parsed["_ai_success"] = True
                return parsed
            log.warning("Nemotron nao retornou JSON — usando heuristica sobre a descricao da visao")
        except Exception as e:
            log.warning("Nemotron structure falhou: %s", e)

        # Reservas: JSON direto da visao; senao, heuristica sobre a prosa da visao
        # (mantem _ai_success=True para o embedding usar o texto real, nao aleatorio).
        direct = self._extract_json(raw_text)
        if direct is not None:
            direct["_ai_success"] = True
            return direct
        return self._heuristic_from_prose(raw_text)

    def _heuristic_from_prose(self, raw_text: str) -> dict:
        """Estrutura minima a partir da descricao livre da visao, sem o Nemotron.

        Usado so quando o Nemotron nao devolve JSON. Faz casamento simples de
        especie/cor/condicao por palavra-chave e usa a propria prosa da visao como
        description_text — assim o embedding e o matching continuam significativos.
        """
        text = (raw_text or "").strip()
        low = text.lower()
        species = "cat" if ("cat" in low and "dog" not in low.split("cat")[0]) else \
                  ("dog" if "dog" in low else "cat" if "cat" in low else "dog")
        color = next((c for c in _COLORS.split(", ") if c.lower() in low), "Mixed")
        condition = "injured" if any(w in low for w in ("injured", "wound", "hurt", "limp")) else \
                    "thin" if any(w in low for w in ("thin", "skinny", "malnourished")) else "healthy"
        return {
            "is_animal": True,
            "_ai_success": True,   # temos descricao real da visao → embedding util
            "species": species,
            "breed_estimate": "SRD" if species == "dog" else "Domestic Shorthair",
            "size": "Medium",
            "primary_color": color,
            "secondary_colors": [],
            "distinctive_marks": [],
            "condition": condition,
            "description_text": text[:300] or ("stray " + species),
        }

    def generate_story(self, animal: dict, sighting_count: int,
                       help_count: int = 0, location_hint: str = "") -> str:
        """Redige uma breve historia do animal (1-2 frases) com o Nemotron 3 Nano.

        E o cerebro do app dando voz ao perfil: a partir dos dados estruturados e do
        historico, escreve uma narrativa curta em portugues para a pagina do animal.
        Retorna "" se o Nemotron nao estiver disponivel ou falhar (o frontend so
        mostra a historia quando ela existe — degrada sem quebrar).
        """
        if self.nim_client is None:
            return ""

        desc = animal if isinstance(animal, dict) else {}
        facts = {
            "species":       desc.get("species", ""),
            "breed":         desc.get("breed_estimate", ""),
            "color":         desc.get("primary_color", ""),
            "condition":     desc.get("condition", ""),
            "marks":         desc.get("distinctive_marks", []),
            "sightings":     sighting_count,
            "help_events":   help_count,
            "location_hint": location_hint,
        }
        prompt = (
            "You write a SHORT story (1-2 sentences, max 40 words) in BRAZILIAN PORTUGUESE "
            "for the public profile of a stray animal being tracked by a community map. "
            "Warm, factual, no invented details. Output ONLY the sentence(s), no preamble.\n"
            "FACTS (JSON):\n" + json.dumps(facts, ensure_ascii=False)
        )
        try:
            resp = self._nemotron_chat(
                messages=[{"role": "user", "content": prompt}],
                max_tokens=512,   # folga p/ reasoning antes da frase final
                temperature=0.6,
                top_p=0.9,
            )
            msg = resp.choices[0].message
            text = (msg.content or "").strip()
            if not text:  # tudo foi p/ reasoning_content? usa de la
                text = (getattr(msg, "reasoning_content", "") or "").strip()
            # Modelo de reasoning pode emitir bloco <think>...</think> — remove.
            text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
            return text
        except Exception as e:
            log.warning("Nemotron generate_story falhou: %s", e)
            return ""

    def get_embedding(self, description: dict) -> list:
        """Embedding da descricao. Aleatorio se IA falhou (evita falsos matches)."""
        if not description.get("_ai_success", True):
            log.info("Fallback IA — embedding aleatorio")
            v = np.random.randn(384).astype(np.float32)
            v /= np.linalg.norm(v)
            return v.tolist()

        if self.embedder is None:
            v = np.random.randn(384).astype(np.float32)
            v /= np.linalg.norm(v)
            return v.tolist()

        text = description.get("description_text") or self._desc_text(description)
        return self.embedder.encode(text, normalize_embeddings=True).tolist()

    @staticmethod
    def _to_b64(image) -> str:
        buf = io.BytesIO()
        img = image.copy()
        img.thumbnail((800, 800))
        img.save(buf, format="JPEG", quality=80)
        return base64.b64encode(buf.getvalue()).decode()

    @staticmethod
    def _extract_json(raw: str):
        """Retorna dict de um bloco JSON valido no texto, ou None.

        Robusto a modelos de reasoning: remove blocos <think>...</think> e cercas
        de codigo, varre todos os objetos {...} balanceados e tenta json.loads em
        cada um (do ultimo p/ o primeiro — a resposta final costuma vir por ultimo).
        """
        if not raw:
            return None
        cleaned = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL)
        cleaned = cleaned.replace("```json", "").replace("```", "")

        candidates = []
        depth = 0
        start = None
        for i, ch in enumerate(cleaned):
            if ch == "{":
                if depth == 0:
                    start = i
                depth += 1
            elif ch == "}" and depth > 0:
                depth -= 1
                if depth == 0 and start is not None:
                    candidates.append(cleaned[start:i + 1])
                    start = None

        for cand in reversed(candidates):
            try:
                obj = json.loads(cand)
                if isinstance(obj, dict):
                    return obj
            except json.JSONDecodeError:
                continue
        return None

    @staticmethod
    def _parse(raw: str) -> dict:
        parsed = AnimalAI._extract_json(raw)
        if parsed is not None:
            return parsed
        log.warning("JSON nao parseado — fallback")
        return AnimalAI._fallback()

    @staticmethod
    def _desc_text(d: dict) -> str:
        parts = [d.get("size",""), d.get("primary_color",""), d.get("species",""), d.get("breed_estimate","")]
        marks = d.get("distinctive_marks", [])
        if marks:
            parts.append("with " + ", ".join(marks))
        return " ".join(filter(None, parts))

    @staticmethod
    def _fallback() -> dict:
        return {
            "is_animal": True,
            "_ai_success": False,
            "species": "dog",
            "breed_estimate": "SRD",
            "size": "Medium",
            "primary_color": "Mixed",
            "secondary_colors": [],
            "distinctive_marks": [],
            "condition": "healthy",
            "description_text": "stray dog of unknown breed",
        }