from flask import Flask, request, jsonify
from PIL import Image
from io import BytesIO
import numpy as np

import os
from keras.models import load_model
from keras.preprocessing import image
import warnings
import random
import base64


# Folder de imagens de Raios-X 
RANDOM_XRAYS_DIR = 'random_xrays'


warnings.filterwarnings("ignore")

app = Flask(__name__, static_folder='.', static_url_path='/')

# ===========================
# Carregamento do modelo
# ===========================
MODEL_PATH = 'modelo_raiox.keras'
MODEL = None

try:
    if not os.path.exists(MODEL_PATH):
        print(f"❌ Arquivo de modelo '{MODEL_PATH}' não encontrado. A predição não funcionará.")
    else:
        MODEL = load_model(MODEL_PATH, compile=False)
        print("✅ Modelo carregado com sucesso.")
except Exception as e:
    print(f"❌ Erro ao carregar o modelo: {e}")

# ===========================
# Função de predição (Adaptada para Flask)
# ===========================
def predict_image(img_data):
    """
    Realiza a predição em um objeto PIL Image.
    Retorna um dicionário com a classe, probabilidade e cores.
    """
    if MODEL is None:
        return {
            "classe": "Erro",
            "prob": 0,
            "cor": "#ff0000",
            "mensagem": "Modelo não carregado."
        }

    IMAGE_HEIGHT, IMAGE_WIDTH = 96, 96
    categories = ['Normal', 'Tuberculose']

    # 1. Pré-processamento
    img = img_data.convert("RGB").resize((IMAGE_HEIGHT, IMAGE_WIDTH))
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0) / 255.0

    # 2. Predição
    pred = MODEL.predict(img, verbose=0)
    classe_idx = np.argmax(pred)
    prob = np.max(pred) * 100

    # 3. Formatação do resultado
    classe = categories[classe_idx]
    cor_classe = "#5dade2" if classe_idx == 0 else "#e74c3c"

    return {
        "classe": classe,
        "prob": round(prob),
        "cor": cor_classe,
        "mensagem": "Predição realizada com sucesso."
    }

# ===========================
# Rotas do Flask

@app.route('/random_predict', methods=['GET'])
def random_predict():
    """Endpoint para buscar uma imagem aleatória local, realizar a predição e retornar a imagem em base64 e o resultado."""
    
    if not os.path.exists(RANDOM_XRAYS_DIR) or not os.listdir(RANDOM_XRAYS_DIR):
        return jsonify({"error": "Imagens aleatórias não encontradas. Verifique o download."}), 500

    # 1. Escolher uma imagem localmente
    imagens_locais = os.listdir(RANDOM_XRAYS_DIR)
    nome_imagem = random.choice(imagens_locais)
    local_path = os.path.join(RANDOM_XRAYS_DIR, nome_imagem)
    
    try:
        # 2. Carregar a imagem para predição
        with open(local_path, 'rb') as f:
            img_data = f.read()
        
        # 3. Realizar a predição
        img = Image.open(BytesIO(img_data))
        result = predict_image(img)
        
        # 4. Codificar a imagem em Base64 para envio ao frontend
        img_base64 = base64.b64encode(img_data).decode('utf-8')
        
        # 5. Combinar resultado da predição com a imagem
        result['image_base64'] = img_base64
        result['image_mimetype'] = 'image/png' # Assumindo que as imagens são PNG
        
        return jsonify(result)

    except Exception as e:
        print(f"Erro interno no random_predict (local): {e}")
        return jsonify({"error": f"Erro interno: {e}"}), 500


@app.route('/')
def index():
    """Serve o arquivo HTML principal."""
    # O Flask está configurado para servir arquivos estáticos da pasta atual.
    # A rota '/' irá servir o index.html.
    return app.send_static_file('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    """Endpoint para receber o arquivo e realizar a predição."""
    if 'file' not in request.files:
        return jsonify({"error": "Nenhum arquivo enviado"}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({"error": "Nome de arquivo vazio"}), 400
    
    if file and file.filename.endswith('.png'):
        try:
            # Lê o arquivo diretamente para um objeto PIL Image
            img = Image.open(BytesIO(file.read()))
            
            # Realiza a predição
            result = predict_image(img)
            
            return jsonify(result)
        
        except Exception as e:
            print(f"Erro durante a predição: {e}")
            return jsonify({"error": f"Erro interno durante a predição: {e}"}), 500
    
    return jsonify({"error": "Formato de arquivo inválido. Apenas .png é aceito."}), 400

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=7860)