import os
import io
from PIL import Image
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS 
import re 
import torch 
from transformers import pipeline

# =================================================================
# 1. SETUP FLASK SERVER & CORS
# =================================================================
app = Flask(__name__)
CORS(app) 

# =================================================================
# 2. SETUP MODEL AI (Kembali ke Pipeline Sederhana)
# =================================================================
MODEL_NAME = "prithivMLmods/BrainTumor-Classification-Mini"
classifier = None 

def load_model():
    """Memuat model AI menggunakan pipeline standar."""
    global classifier
    try:
        print("⏳ Sedang memuat model AI ({})...".format(MODEL_NAME))
        device = "cuda" if torch.cuda.is_available() else "cpu"
        classifier = pipeline(
            "image-classification",
            model=MODEL_NAME,
            device=device 
        )
        print("✅ Model {} berhasil dimuat ke device: {}".format(MODEL_NAME, device))
    except Exception as e:
        print("❌ Gagal memuat model: {}".format(e))
        classifier = None 

load_model()

# Pemetaan label
LABEL_MAPPING = {
    'glioma': {'status': 'Tumor Terdeteksi', 'jenis': 'Glioma'},
    'meningioma': {'status': 'Tumor Terdeteksi', 'jenis': 'Meningioma'},
    'notumor': {'status': 'Tidak Ada Tumor', 'jenis': 'Tidak Ada'},
    'pituitary': {'status': 'Tumor Terdeteksi', 'jenis': 'Pituitary Tumor'}
}

def clean_label(raw_label):
    """Membersihkan label mentah dari model."""
    match = re.search(r'(glioma|meningioma|notumor|pituitary)', raw_label, re.IGNORECASE)
    if match:
        return match.group(0).lower()
    return raw_label.lower()

# =================================================================
# 3. ENDPOINT WEB SERVER (Host HTML)
# =================================================================

@app.route('/', methods=['GET'])
def serve_html():
    """Endpoint untuk menampilkan file index.html di tab 'App'."""
    try:
        return send_file('index.html')
    except Exception as e:
        return f"<h1>Error: File index.html tidak ditemukan di server.</h1><p>Pastikan file index.html sudah di-COPY ke Docker container.</p><p>{e}</p>", 500

# =================================================================
# 4. ENDPOINT API (Predict)
# =================================================================

@app.route('/predict', methods=['POST'])
def predict():
    """Endpoint utama untuk prediksi."""
    if classifier is None:
        return jsonify({"error": "Model AI belum dimuat atau gagal dimuat."}), 500

    if 'file' not in request.files:
        return jsonify({"error": "no file uploaded"}), 400
    
    file = request.files['file']
    
    try:
        image_bytes = file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
    except Exception as e:
        return jsonify({"error": f"Gagal memproses gambar: {e}"}), 400

    # Lakukan Prediksi
    try:
        # Gunakan pipeline
        results = classifier(images=image, top_k=1) 
        
        result = results[0]
        raw_label = result['label']
        confidence = result['score'] * 100

        clean_key = clean_label(raw_label)
        result_data = LABEL_MAPPING.get(clean_key, {
            'status': 'Hasil Tidak Dikenal', 
            'jenis': 'N/A'
        })
        
        return jsonify({
            "prediction_status": result_data['status'],
            "tumor_type": result_data['jenis'],
            "confidence": round(confidence, 2)
        })

    except Exception as e:
        return jsonify({"error": "Terjadi kesalahan saat menjalankan prediksi model."}), 500