import streamlit as st
import google.generativeai as genai
from PIL import Image
import io

class ForensicVLMOrchestrator:
    def __init__(self, mode="api"):
        self.mode = mode
        # Configura a API do Google Gemini
        api_key = st.secrets.get("GOOGLE_API_KEY")
        if api_key:
            genai.configure(api_key=api_key)
            self.model = genai.GenerativeModel('gemini-2.5-flash')

    def _determine_typology(self, prob_swin, prob_clip):
        if prob_swin > prob_clip:
            return "uma imagem totalmente criada por computador (não existe na vida real)"
        else:
            return "uma fotografia verdadeira onde o rosto foi alterado ou trocado"

    def _crop_artifact_zone(self, img_rgb, bbox):
        if not bbox:
            return img_rgb
        h, w = img_rgb.shape[:2]
        x1, y1, x2, y2 = bbox
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)
        crop = img_rgb[y1:y2, x1:x2]
        if crop.size == 0:
            return img_rgb
        return crop

    def build_forensic_prompt(self, prob_final, prob_swin, prob_clip, zones_list):
        tipo_fraude = self._determine_typology(prob_swin, prob_clip)

        prompt = (
            f"Gere um relatório pericial objetivo e acessível a um público leigo.\n"
            f"Analise o recorte da imagem.\n\n"
            f"--- DADOS DO SISTEMA ---\n"
            f"Certeza de fraude: {prob_final * 100:.1f}%\n"
            f"Tipo de fraude: {tipo_fraude}\n"
            f"Zonas com anomalias detetadas: {zones_list}\n"
            f"-----------------------\n\n"
            f"REGRAS OBRIGATÓRIAS:\n"
            f"1. Escreva exatamente 4 frases num único parágrafo, em Português de Portugal.\n"
            f"2. PROIBIDO usar a primeira pessoa (eu, vejo, noto, acho). Use exclusivamente voz impessoal e formal.\n"
            f"3. PROIBIDO usar calão ou jargão técnico (IA, algoritmo, síntese, CLIP, Swin, artefactos, píxeis).\n"
            f"4. FRASE 1: Identifique as zonas afetadas ({zones_list}) e descreva os defeitos visuais presentes. Utilize conceitos reais de análise de imagem.\n"
            f"5. FRASE 2: Indique que estas inconsistências visuais evidenciam tratar-se de {tipo_fraude}.\n"
            f"6. FRASE 3: Descreva características globais da imagem que não são condizentes com uma fotografia genuína.\n"
            f"7. FRASE 4: Conclua EXATAMENTE com: 'O sistema tem {prob_final * 100:.1f}% de certeza de que esta imagem é falsa.'\n"
        )
        return prompt

    def generate_justification(self, img_rgb, prob_final, prob_swin, prob_clip, zone_name, bbox):
        crop_rgb = self._crop_artifact_zone(img_rgb, bbox)
        prompt = self.build_forensic_prompt(prob_final, prob_swin, prob_clip, zone_name)

        if not hasattr(self, 'model'):
            return "Erro: Google API Key não configurada."

        # Preparar imagem (Thumb para otimizar token/latência)
        pil_img = Image.fromarray(crop_rgb)
        pil_img.thumbnail((512, 512), Image.Resampling.BICUBIC)

        try:
            # Geração com streaming nativo do SDK do Gemini
            response = self.model.generate_content(
                [prompt, pil_img],
                stream=True,
                generation_config={"temperature": 0.1}
            )
            
            def stream_generator():
                for chunk in response:
                    if chunk.text:
                        yield chunk.text
            
            return stream_generator()
            
        except Exception as e:
            erro_str = str(e)
            if "429" in erro_str or "quota" in erro_str.lower():
                return "O sistema está a processar demasiados pedidos em simultâneo (Rate Limit atingido). Por favor, aguarde 30 segundos e tente analisar a imagem novamente."
            else:
                return f"Falha de comunicação com o motor de inferência LVM. Detalhe técnico: {erro_str}"