import gradio as gr
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import os
import unicodedata
import re
import requests
from sentence_transformers import SentenceTransformer

# Modelo oficial da Google (Visão Pesada)
model_id = "google/paligemma-3b-mix-224"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
processor = AutoProcessor.from_pretrained(model_id)

# Modelo CLIP para Memória Visual (Leve e rápido)
clip_model = SentenceTransformer('clip-ViT-B-32')

# =============================================================================
# PALAVRAS GENÉRICAS QUE NÃO DIFERENCIAM PRODUTOS
# =============================================================================
STOPWORDS = {
    "de", "da", "do", "das", "dos", "com", "sem", "para", "por", "em",
    "e", "a", "o", "os", "as", "um", "uma", "uns", "umas", "no", "na",
    "nos", "nas", "ao", "aos", "se", "que", "kit", "par", "pares",
    "original", "oficial", "novo", "nova", "nacional", "importado",
    "feminino", "masculino", "infantil", "adulto", "unissex",
    "premium", "super", "ultra", "plus", "max", "mini", "pro", "lite",
    "conforto", "macio", "macias", "suave", "respiravel", "respiravel",
    "moda", "intima", "basica", "casual", "social",
    "polegada", "polegadas", "pol", "bivolt", "127v", "220v",
}

def normalize(text):
    if not text:
        return ""
    return unicodedata.normalize('NFD', text).encode('ascii', 'ignore').decode('utf-8').lower()

def title_to_fingerprint(title):
    if not title:
        return ""
    t = title.strip()
    specs_extracted = []
    for m in re.finditer(r'(\d+\.?\d*)\s*(?:"|pol\.?|polegadas?)\b', t, re.IGNORECASE):
        specs_extracted.append(f'{m.group(1)}pol')
    for m in re.finditer(r'(\d+)\s*hz\b', t, re.IGNORECASE):
        specs_extracted.append(f'{m.group(1)}hz')
    for m in re.finditer(r'(\d+)\s*(gb|tb)\b', t, re.IGNORECASE):
        specs_extracted.append(f'{m.group(1)}{m.group(2).lower()}')
    for m in re.finditer(r'\b(i[3579]|ryzen\s*[579]|celeron|core\s*ultra\s*\d)\b', t, re.IGNORECASE):
        specs_extracted.append(normalize(m.group(1)).replace(' ', ''))
    for m in re.finditer(r'\b(pp|xg|gg|eg|3g|4g)\b', t, re.IGNORECASE):
        specs_extracted.append(m.group(1).lower())
    for m in re.finditer(r'\b([pmg])\b', t, re.IGNORECASE):
        specs_extracted.append(m.group(1).lower())
    for m in re.finditer(r'\b(3[0-9]|4[0-9]|5[0-9])\b', t):
        specs_extracted.append(m.group(1))
    for m in re.finditer(r'\b([a-zA-Z]{1,4}[-]?\d{3,8}[a-zA-Z0-9-]*|[a-zA-Z]{2,6}\d{2,6}[a-zA-Z0-9]*)\b', t):
        code = normalize(m.group(1))
        if not re.fullmatch(r'\d+(gb|tb|hz|pol)', code):
            specs_extracted.append(code)

    t_norm = normalize(t)
    t_norm = re.sub(r'(\d+\.?\d*)\s*(?:"|pol\.?|polegadas?)\b', '', t_norm)
    t_norm = re.sub(r'(\d+)\s*hz\b', '', t_norm)
    t_norm = re.sub(r'(\d+)\s*(gb|tb)\b', '', t_norm)
    t_norm = re.sub(r'\b(i[3579]|ryzen\s*[579]|celeron|core\s*ultra\s*\d)\b', '', t_norm)
    
    words = [w for w in re.split(r'\W+', t_norm) if w and len(w) > 1 and w not in STOPWORDS]
    base_text = " ".join(words[:4])
    
    final_parts = [base_text] + specs_extracted
    return " ".join(final_parts).strip()


def get_visual_memory(image):
    supabase_url = os.environ.get("AI_SUPABASE_URL")
    supabase_key = os.environ.get("AI_SUPABASE_KEY")
    if not supabase_url or not supabase_key:
        return None
        
    embedding = clip_model.encode(image).tolist()
    
    headers = {
        "apikey": supabase_key,
        "Authorization": f"Bearer {supabase_key}",
        "Content-Type": "application/json"
    }
    
    try:
        res = requests.post(
            f"{supabase_url}/rest/v1/rpc/match_visual_memory",
            headers=headers,
            json={"query_embedding": embedding, "match_threshold": 0.90, "match_count": 1},
            timeout=5
        )
        if res.status_code == 200 and len(res.json()) > 0:
            return res.json()[0]['correct_type']
    except:
        pass
    return None

def save_visual_memory(image, correct_type):
    supabase_url = os.environ.get("AI_SUPABASE_URL")
    supabase_key = os.environ.get("AI_SUPABASE_KEY")
    if not supabase_url or not supabase_key:
        return "Erro: Chaves do Supabase não configuradas no Hugging Face."
        
    embedding = clip_model.encode(image).tolist()
    
    headers = {
        "apikey": supabase_key,
        "Authorization": f"Bearer {supabase_key}",
        "Content-Type": "application/json",
        "Prefer": "return=minimal"
    }
    
    data = {
        "correct_type": correct_type,
        "embedding": embedding
    }
    
    try:
        res = requests.post(f"{supabase_url}/rest/v1/visual_memory", headers=headers, json=data, timeout=5)
        if res.status_code in [200, 201]:
            return "LEARN_OK"
        return f"Erro Supabase: {res.text}"
    except Exception as e:
        return f"Erro Requisição: {str(e)}"


def analyze_offer(image, title, secret_key="", dynamic_memory=""):
    if image is None:
        return "Erro: Nenhuma imagem enviada"

    valid_key = os.environ.get("VAL_SECRET_KEY")
    if valid_key and secret_key != valid_key:
        return "Erro: Chave Secreta Inválida"

    try:
        image = image.convert("RGB")

        # Se for comando para APRENDER memória visual:
        if dynamic_memory and dynamic_memory.startswith("LEARN_VISUAL:"):
            correct_type = dynamic_memory.replace("LEARN_VISUAL:", "").strip()
            return save_visual_memory(image, correct_type)

        # PASSO 1: Gera fingerprint do título
        fingerprint = title_to_fingerprint(title)

        # PASSO 2: Busca na Memória Visual (Super-rápido, via CLIP)
        visual_type = get_visual_memory(image)
        if visual_type:
            visual_type_norm = normalize(visual_type)
            fp_parts = fingerprint.split()
            if fp_parts:
                fp_parts[0] = visual_type_norm
            return " ".join(fp_parts) if fp_parts else visual_type_norm + " " + fingerprint

        # PASSO 3: Se não tem memória visual, usa a IA Pesada (PaliGemma)
        # Token <image> obrigatório no início para o PaliGemmaProcessor (transformers >= 4.40)
        prompt = (
            "<image> answer pt Qual é o tipo deste produto? Responda em 2 palavras no máximo. "
            "Seja específico: calcinha, cueca, tênis, monitor, notebook, camisa, etc."
        )
        
        inputs = processor(text=prompt, images=image, return_tensors="pt")
        generate_ids = model.generate(
            **inputs,
            max_new_tokens=10,
            repetition_penalty=1.1,
            no_repeat_ngram_size=2,
            do_sample=False,
            temperature=None
        )
        resposta_pura = processor.batch_decode(generate_ids, skip_special_tokens=True)[0]
        # O batch_decode remove o <image> automaticamente; removemos o resto do prompt
        prompt_limpo = prompt.replace("<image> ", "")
        ai_type = resposta_pura.replace(prompt_limpo, "").replace(prompt, "").strip()

        texto_min = normalize(ai_type)
        is_invalid = (not ai_type or "nao" in texto_min or "sorry" in texto_min or "desculpe" in texto_min or len(ai_type) < 2)

        if not is_invalid:
            ai_type_norm = normalize(ai_type)
            fp_parts = fingerprint.split()
            if fp_parts:
                fp_parts[0] = ai_type_norm
            fingerprint = " ".join(fp_parts) if fp_parts else ai_type_norm + " " + fingerprint

        return fingerprint if fingerprint else (title[:60] if title else "Fallback")

    except Exception as e:
        return f"CRASH_REAL: {str(e)}"

demo = gr.Interface(
    fn=analyze_offer,
    inputs=[
        gr.Image(type="pil", label="Foto do Produto"),
        gr.Textbox(label="Título original (opcional)"),
        gr.Textbox(label="Secret Key (opcional)", type="password"),
        gr.Textbox(label="Memória Dinâmica / Comando (opcional)")
    ],
    outputs=gr.Textbox(label="Fingerprint / Status"),
    title="🤖 IA Analítica de Produtos c/ Memória Visual",
    description="Gera fingerprint baseado no TÍTULO (marca/modelo) e IA visual para o tipo. Suporta RAG de imagens."
)

demo.launch()