import os
import sys
import torch
import numpy as np
import supervision as sv
from inference import get_model
from PIL import Image
from typing import List, Dict
from collections import Counter
import shutil
import tempfile
from flask import Flask, request, jsonify, render_template
from dotenv import load_dotenv

# Carrega as variáveis do arquivo .env para o ambiente
load_dotenv()

# Adiciona o diretório 'Long-CLIP' ao path para encontrar a pasta 'model'
long_clip_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Long-CLIP")
sys.path.append(long_clip_path)

try:
    from model import longclip
except ImportError:
    print(f"Erro: A pasta 'model' do LongCLIP não foi encontrada em: {long_clip_path}")

# Suprime avisos de dependências ausentes da biblioteca inference (opcional)
os.environ["CORE_MODEL_SAM_ENABLED"] = "False"
os.environ["CORE_MODEL_SAM3_ENABLED"] = "False"
os.environ["CORE_MODEL_GAZE_ENABLED"] = "False"
os.environ["CORE_MODEL_YOLO_WORLD_ENABLED"] = "False"

app = Flask(__name__)

# Variável global para armazenar as análises em lote (Artigo)
resultados_globais_artigo = []

# --- CONFIGURAÇÕES INICIAIS ---
# A chave da API do Roboflow agora é carregada automaticamente do arquivo .env
device = "cuda" if torch.cuda.is_available() else "cpu"
output_dir = "static/outputs"
app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()

# --- FUNÇÕES DO LONGCLIP ---
def load_longclip_model(checkpoint_path, device="cpu"):
    checkpoint_path = os.path.normpath(checkpoint_path)
    print(f"Carregando LongCLIP de {checkpoint_path}...")
    
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(f"Erro: O arquivo de checkpoint não foi encontrado em: {checkpoint_path}")
        
    model, preprocess = longclip.load(checkpoint_path, device=device)
    return model, preprocess

def get_text_features(model, text_descriptions: Dict[str, str], device="cpu"):
    class_names = list(text_descriptions.keys())
    descriptions = list(text_descriptions.values())
    
    print(f"Tokenizando {len(descriptions)} descrições detalhadas...")
    tokens = longclip.tokenize(descriptions).to(device)
    
    with torch.no_grad():
        features = model.encode_text(tokens)
        features /= features.norm(dim=-1, keepdim=True)
    
    return features, class_names

# --- DICIONÁRIO DE DESCRIÇÕES OTIMIZADO PARA LONGCLIP ---
text_descriptions_longclip = {
    "Mancha de Olho Pardo (Cercospora)": (
        "A highly magnified view showing a distinct circular 'bullseye' or eye geometry. "
        "The absolute center must present a lighter, silvery, pale, or white core. "
        "This central pale area is obligatorily surrounded by prominent, dark concentric rings. "
        "Outside these dark rings lies a diffuse yellow halo."
    ),
    "Ferrugem do Cafeeiro": (
        "A highly magnified view of bright, vivid cadmium-orange textures. "
        "The surface is densely covered in thick powdery granular dust or accumulated pollen-like spores. "
        "The bright orange fuzzy material forms soft, blurry, and diffuse shapes."
    ),
    "Bicho Mineiro do Café": (
        "A highly magnified view of parched, light-tan or beige blisters. "
        "The texture consists of flat, withered material with a highly wrinkled, translucent, and papery appearance. "
        "It looks exactly like dry, flaky skin or torn parched parchment."
    ),
    "Mancha de Phoma": (
        "A highly magnified view of a massive, solid, dark-brown or pitch-black necrotic mass. "
        "The affliction appears as a severe, thick dark rot or burn. "
        "It forms a continuous, opaque block of dark dead tissue, often causing the affected margins to become torn."
    )
}


@app.route('/')
def home():
    supabase_url = os.environ.get('SUPABASE_URL')
    supabase_key = os.environ.get('SUPABASE_ANON_KEY')
    return render_template('index.html', supabase_url=supabase_url, supabase_key=supabase_key)

@app.route('/redefinir-senha')
def redefinir_senha():
    supabase_url = os.environ.get('SUPABASE_URL')
    supabase_key = os.environ.get('SUPABASE_ANON_KEY')
    return render_template('redefinir_senha.html', supabase_url=supabase_url, supabase_key=supabase_key)

def process_leaf_image(image_path, model_roboflow, long_model, long_preprocess, text_features, class_names, device):
    # LÓGICA DE LIMPEZA DA PASTA
    if os.path.exists(output_dir):
        for filename in os.listdir(output_dir):
            file_path = os.path.join(output_dir, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f'Falha ao deletar {file_path}. Motivo: {e}')
    else:
        os.makedirs(output_dir)

    # 1. DETECÇÃO (ROBOFLOW)
    image = Image.open(image_path)
    predictions = model_roboflow.infer(image, confidence=0.6)[0]
    detections = sv.Detections.from_inference(predictions)

    img_width, img_height = image.size
    padding_percentage = 0.20
    min_dim_for_padding = 100

    for i, xyxy in enumerate(detections.xyxy):
        x_min_orig, y_min_orig, x_max_orig, y_max_orig = xyxy
        bbox_width_orig = x_max_orig - x_min_orig
        bbox_height_orig = y_max_orig - y_min_orig
        final_x_min, final_y_min, final_x_max, final_y_max = x_min_orig, y_min_orig, x_max_orig, y_max_orig

        if bbox_width_orig < min_dim_for_padding or bbox_height_orig < min_dim_for_padding:
            pad_x = bbox_width_orig * padding_percentage
            pad_y = bbox_height_orig * padding_percentage
            final_x_min = max(0, x_min_orig - pad_x)
            final_y_min = max(0, y_min_orig - pad_y)
            final_x_max = min(img_width, x_max_orig + pad_x)
            final_y_max = min(img_height, y_max_orig + pad_y)

        final_x_min, final_y_min, final_x_max, final_y_max = int(final_x_min), int(final_y_min), int(final_x_max), int(final_y_max)
        cropped_image = image.crop((final_x_min, final_y_min, final_x_max, final_y_max))
        cropped_image.save(os.path.join(output_dir, f"recorte_{i}.jpg"))

    # 2. CLASSIFICAÇÃO (LONGCLIP)
    results_list = []
    image_urls = []

    for file_name in os.listdir(output_dir):
        if file_name.startswith("recorte_") and file_name.endswith(".jpg"):
            img_path = os.path.join(output_dir, file_name)
            img_input = Image.open(img_path)
            image_input = long_preprocess(img_input).unsqueeze(0).to(device)
            
            with torch.no_grad():
                image_features = long_model.encode_image(image_input)
                image_features /= image_features.norm(dim=-1, keepdim=True)
                logits = (100.0 * image_features @ text_features.T)
                probs = logits.softmax(dim=-1).cpu().numpy()[0]
            
            top_idx = np.argmax(probs)
            predicted_class = class_names[top_idx]
            top_prob = float(probs[top_idx])
            
            results_list.append(predicted_class)
            image_urls.append({
                "url": f"/static/outputs/{file_name}",
                "classe": predicted_class,
                "probabilidade": top_prob
            })

    contagem = Counter(results_list)
    mais_frequente = contagem.most_common(1)[0][0] if results_list else "Nenhuma detecção encontrada"

    return {
        "contagem": dict(contagem),
        "mais_frequente": mais_frequente,
        "total": sum(contagem.values()),
        "imagens": image_urls
    }

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'Nenhum arquivo enviado'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'Arquivo vazio selecionado'}), 400
        
    try:
        image_path = os.path.join(app.config['UPLOAD_FOLDER'], "temp_leaf.jpg")
        file.save(image_path)

        model_roboflow = get_model("bracol-validado-region-detect-oupfv/1")
        checkpoint_abs_path = os.path.join(long_clip_path, "checkpoints", "longclip-L.pt")
        long_model, long_preprocess = load_longclip_model(checkpoint_abs_path, device)
        text_features, class_names = get_text_features(long_model, text_descriptions_longclip, device)

        result = process_leaf_image(
            image_path, model_roboflow, long_model, long_preprocess, 
            text_features, class_names, device
        )

        return jsonify(result)

    except Exception as e:
        return jsonify({'error': str(e)}), 500

def get_true_class_from_filename(filename):
    name_lower = filename.lower()
    if 'ferrugem' in name_lower:
        return "Ferrugem do Cafeeiro"
    elif 'cercospora' in name_lower or 'olho pardo' in name_lower:
        return "Mancha de Olho Pardo (Cercospora)"
    elif 'mineiro' in name_lower or 'bicho' in name_lower:
        return "Bicho Mineiro do Café"
    elif 'phoma' in name_lower:
        return "Mancha de Phoma"
    else:
        return None

@app.route('/predict_batch', methods=['POST'])
def predict_batch():
    if 'files' not in request.files:
        return jsonify({'error': 'Nenhum arquivo enviado'}), 400
    
    files = request.files.getlist('files')
    if not files or files[0].filename == '':
        return jsonify({'error': 'Nenhum arquivo selecionado'}), 400
        
    try:
        # Carregar modelos uma única vez para o lote inteiro
        model_roboflow = get_model("bracol-validado-region-detect-oupfv/1")
        checkpoint_abs_path = os.path.join(long_clip_path, "checkpoints", "longclip-L.pt")
        long_model, long_preprocess = load_longclip_model(checkpoint_abs_path, device)
        text_features, class_names = get_text_features(long_model, text_descriptions_longclip, device)

        batch_results = []
        y_true = []
        y_pred = []

        for idx, file in enumerate(files):
            image_path = os.path.join(app.config['UPLOAD_FOLDER'], f"temp_leaf_batch_{idx}.jpg")
            file.save(image_path)

            result = process_leaf_image(
                image_path, model_roboflow, long_model, long_preprocess, 
                text_features, class_names, device
            )
            
            predicted_class = result["mais_frequente"]
            batch_results.append({
                "filename": file.filename,
                "mais_frequente": predicted_class
            })
            
            # Extrair classe real: prioriza a seleção manual se existir, senão usa o nome do arquivo
            manual_true_class = request.form.get('true_class')
            if manual_true_class:
                true_class = manual_true_class
            else:
                true_class = get_true_class_from_filename(file.filename)
                
            if true_class and predicted_class != "Nenhuma detecção encontrada":
                y_true.append(true_class)
                y_pred.append(predicted_class)

        # Consolidar a contagem da doença principal de cada folha do lote
        doencas_principais = [res["mais_frequente"] for res in batch_results if res["mais_frequente"] != "Nenhuma detecção encontrada"]
        contagem_geral = Counter(doencas_principais)
        
        # Calcular porcentagens e garantir que todas as classes (doenças) apareçam no relatório
        total_validas = len(doencas_principais)
        estatisticas = {}
        for doenca in class_names:
            qtd = contagem_geral.get(doenca, 0)
            pct = (qtd / total_validas * 100) if total_validas > 0 else 0
            estatisticas[doenca] = {
                "quantidade": qtd,
                "porcentagem": round(pct, 1)
            }

        # Calcular métricas de avaliação
        metricas_avaliacao = {}
        if len(y_true) > 0:
            acertos = sum(1 for yt, yp in zip(y_true, y_pred) if yt == yp)
            acuracia_geral = (acertos / len(y_true)) * 100
            
            metricas_por_classe = {}
            for cls in class_names:
                tp = sum(1 for yt, yp in zip(y_true, y_pred) if yt == cls and yp == cls)
                fp = sum(1 for yt, yp in zip(y_true, y_pred) if yt != cls and yp == cls)
                fn = sum(1 for yt, yp in zip(y_true, y_pred) if yt == cls and yp != cls)
                
                precision = (tp / (tp + fp)) * 100 if (tp + fp) > 0 else 0
                recall = (tp / (tp + fn)) * 100 if (tp + fn) > 0 else 0
                f1_score = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
                
                metricas_por_classe[cls] = {
                    "precision": round(precision, 2),
                    "recall": round(recall, 2),
                    "f1_score": round(f1_score, 2),
                    "suporte": sum(1 for yt in y_true if yt == cls)
                }
            
            matriz_confusao = {}
            for cls_real in class_names:
                matriz_confusao[cls_real] = {}
                for cls_pred in class_names:
                    matriz_confusao[cls_real][cls_pred] = sum(1 for yt, yp in zip(y_true, y_pred) if yt == cls_real and yp == cls_pred)

            metricas_avaliacao = {
                "acuracia_geral": round(acuracia_geral, 2),
                "por_classe": metricas_por_classe,
                "matriz_confusao": matriz_confusao,
                "total_avaliado": len(y_true),
                "dados_brutos": {
                    "y_true": y_true,
                    "y_pred": y_pred
                }
            }

        resultado_final = {
            "estatisticas": estatisticas,
            "detalhes": batch_results,
            "total_folhas": len(files),
            "folhas_com_deteccao": total_validas,
            "metricas_avaliacao": metricas_avaliacao
        }
        
        import datetime
        timestamp_str = datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S")
        
        doencas_no_lote = list(set(y_true)) if len(y_true) > 0 else ["Desconhecida"]
        
        for doenca in doencas_no_lote:
            erros_confusao = {}
            suporte = len(files)
            if doenca != "Desconhecida" and "matriz_confusao" in metricas_avaliacao:
                erros_confusao = metricas_avaliacao["matriz_confusao"].get(doenca, {})
                suporte = metricas_avaliacao.get("por_classe", {}).get(doenca, {}).get("suporte", len(files))

            resultados_globais_artigo.append({
                "timestamp": timestamp_str,
                "doenca_alvo": doenca,
                "total_folhas": suporte,
                "estatisticas": estatisticas,
                "acuracia_geral": metricas_avaliacao.get("acuracia_geral", 0),
                "metricas_classe": metricas_avaliacao.get("por_classe", {}).get(doenca, {}) if doenca != "Desconhecida" else {},
                "erros_confusao": erros_confusao
            })

        return jsonify(resultado_final)

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/resultados_artigo')
def resultados_artigo():
    return render_template('resultados_artigo.html', dados=resultados_globais_artigo)

if __name__ == '__main__':
    app.run(debug=True, use_reloader=False)