import os
import shutil
import zipfile
import gradio as gr
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
import tempfile
import subprocess

# ==========================================
# 1. DEEPISLES SETUP (Environment Trickery)
# ==========================================
os.environ['nnUNet_raw_data_base'] = '/tmp/nnUNet_raw_data_base'
os.environ['nnUNet_preprocessed'] = '/tmp/nnUNet_preprocessed'
os.environ['RESULTS_FOLDER'] = '/tmp/nnUNet_trained_models'

TASK_NAME = "Task012_Ischemic_Stroke_TM_Fullset"
TRAINER = "nnUNetTrainerV2_DDP__nnUNetPlansv2.1"
MODEL_PATH_INTERNAL = os.path.join(os.environ['RESULTS_FOLDER'], "nnUNet", "3d_fullres", TASK_NAME, TRAINER, "fold_0")

print("--- 🧠 Initializing DeepISLES Engine ---")

UPLOADED_WEIGHTS = "fold_0.zip" 

if os.path.exists(UPLOADED_WEIGHTS):
    if not os.path.exists(MODEL_PATH_INTERNAL):
        print(f"📦 Installing {UPLOADED_WEIGHTS}...")
        os.makedirs(MODEL_PATH_INTERNAL, exist_ok=True)
        try:
            with zipfile.ZipFile(UPLOADED_WEIGHTS, 'r') as zip_ref:
                zip_ref.extractall(MODEL_PATH_INTERNAL)
            print("✅ DeepISLES Weights Installed.")
        except Exception as e:
            print(f"❌ Error extracting weights: {e}")
else:
    print(f"⚠️ CRITICAL: '{UPLOADED_WEIGHTS}' not found. Please upload it to the Files tab.")

# ==========================================
# 2. FILE HANDLING
# ==========================================
def find_sequences(extract_dir):
    dwi = None
    adc = None
    for root, _, files in os.walk(extract_dir):
        for f in files:
            name = f.lower()
            if "dwi" in name and ".nii" in name: dwi = os.path.join(root, f)
            elif "adc" in name and ".nii" in name: adc = os.path.join(root, f)
    return dwi, adc

def generate_overlay(dwi_path, mask_path):
    try:
        img = nib.load(dwi_path)
        mask = nib.load(mask_path)
        dwi = img.get_fdata()
        seg = mask.get_fdata()
        
        dwi = np.rot90(dwi, k=1, axes=(0, 1))
        seg = np.rot90(seg, k=1, axes=(0, 1))
        
        sums = np.sum(seg, axis=(0, 1))
        idx = np.argmax(sums) if np.max(sums) > 0 else dwi.shape[2]//2
        
        fig, ax = plt.subplots(figsize=(5, 5))
        ax.imshow(dwi[:, :, idx], cmap="gray")
        if np.max(seg) > 0:
            overlay = np.zeros((*seg.shape[:2], 4))
            overlay[seg[:, :, idx] > 0] = [1, 0, 0, 0.6] 
            ax.imshow(overlay)
        ax.axis("off")
        
        out_path = "/tmp/overlay.png"
        plt.savefig(out_path, bbox_inches='tight')
        plt.close()
        return out_path
    except:
        return None

# ==========================================
# 3. INFERENCE ENGINE (nnU-Net Wrapper)
# ==========================================
def segment_stroke(zip_file):
    if not zip_file: return {"error": "No file uploaded"}, None
    
    work_dir = tempfile.mkdtemp()
    input_dir = os.path.join(work_dir, "input")
    output_dir = os.path.join(work_dir, "output")
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # Handle Gradio 5 file path string
        file_path = zip_file.name if hasattr(zip_file, 'name') else zip_file
        
        with zipfile.ZipFile(file_path, 'r') as z:
            z.extractall(work_dir)
            
        dwi, adc = find_sequences(work_dir)
        if not dwi or not adc:
            return {"status": "FAILED", "error": "DWI or ADC missing in zip"}, None
            
        # Format for nnU-Net: 0000=DWI, 0001=ADC
        shutil.copy(dwi, os.path.join(input_dir, "patient_0000.nii.gz"))
        shutil.copy(adc, os.path.join(input_dir, "patient_0001.nii.gz"))
        
        print("🚀 Running DeepISLES Inference...")
        cmd = [
            "nnUNet_predict",
            "-i", input_dir,
            "-o", output_dir,
            "-t", "012",
            "-m", "3d_fullres",
            "-f", "0",
            "--save_npz" 
        ]
        
        process = subprocess.run(cmd, capture_output=True, text=True)
        
        if process.returncode != 0:
            print("NNUNET ERROR:", process.stderr)
            return {"status": "FAILED", "error": "Model inference failed. See logs."}, None
            
        result_file = os.path.join(output_dir, "patient.nii.gz")
        if os.path.exists(result_file):
            mask_img = nib.load(result_file)
            mask_data = mask_img.get_fdata()
            zooms = mask_img.header.get_zooms()
            vol_ml = (np.sum(mask_data > 0) * np.prod(zooms)) / 1000.0
            
            risk = "LOW"
            if vol_ml > 70: risk = "HIGH (Futile Recanalization?)"
            elif vol_ml > 50: risk = "MODERATE"
            
            report = {
                "status": "SUCCESS",
                "core_volume_ml": round(vol_ml, 2),
                "risk_assessment": risk,
                "model_used": "DeepISLES (nnU-Net Ensembled)"
            }
            
            overlay = generate_overlay(dwi, result_file)
            return report, overlay
        else:
            return {"status": "FAILED", "error": "No output mask generated"}, None

    except Exception as e:
        return {"status": "CRASHED", "error": str(e)}, None
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)

# ==========================================
# 4. UI
# ==========================================
with gr.Blocks() as demo:
    gr.Markdown("# 👁️ Neuro-Triage: Vision Agent")
    gr.Markdown("Upload Patient ZIP (DWI + ADC). Calculates Infarct Core Volume.")
    
    with gr.Row():
        with gr.Column():
            inp = gr.File(label="Patient Data (.zip)", type="filepath")
            btn = gr.Button("Calculate Core Volume", variant="primary")
        with gr.Column():
            out_json = gr.JSON(label="Clinical Report")
            out_img = gr.Image(label="Lesion Map")
            
    btn.click(segment_stroke, inputs=inp, outputs=[out_json, out_img])

if __name__ == "__main__":
    demo.launch()