import streamlit as st
import shutil # <-- Nayi library for file copy (Fixes cross-device link error)
import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
import os

# --- 1. Configuration (Constants) ---
# Yahan aapka confirmed Repo ID use ho raha hai
HF_REPO_ID = "saad1BM/brain-tumor-detection" 
MODEL_FILENAME = "brain_tumor_model.keras" 
MODEL_PATH = MODEL_FILENAME 

IMAGE_SIZE = (224, 224)
CLASS_NAMES = ['glioma_tumor', 'meningioma_tumor', 'no_tumor', 'pituitary_tumor']


# --- 2. Model Loading (FINAL FIX: Cross-Device Link Resolved) ---
@st.cache_resource
def load_trained_model():
    """Trained model ko Hugging Face se download aur load karta hai."""
    
    # Global variables ko function ke andar access karna
    global HF_REPO_ID, MODEL_FILENAME, MODEL_PATH
    
    # Minimum expected model size (approx 100 MB, for size check)
    MIN_SIZE = 100 * 1024 * 1024 

    try:
        # Check karein ki model pehle se download hai ya nahi ya file corrupt/incomplete hai
        if not os.path.exists(MODEL_PATH) or os.path.getsize(MODEL_PATH) < MIN_SIZE: 
            st.info(f"Downloading model from Hugging Face: {HF_REPO_ID}/{MODEL_FILENAME}")
            
            # hf_hub_download model ko /tmp mein download karega
            downloaded_path = hf_hub_download(
                repo_id=HF_REPO_ID, 
                filename=MODEL_FILENAME, 
                # Cache ko /tmp mein rakhein
                cache_dir="/tmp/hf_cache", 
            )
            
            st.info("Copying model file locally to avoid cross-device link error...")
            
            # os.rename ki jagah shutil.copyfile use karein (Yeh cross-device error ko theek karta hai)
            shutil.copyfile(downloaded_path, MODEL_PATH) 

            st.success("Model download and copy successful!")
            
        # Load the model
        model = load_model(MODEL_PATH)
        return model
        
    except Exception as e:
        # Error message mein zaroori details daal di gayi hain
        st.error(f"Model Load Failed. Ensure 'huggingface-hub' is in requirements.txt and file is public. Error: {e}")
        return None

model = load_trained_model()

# --- 3. Prediction Function ---
def predict_image(image_file, model):
    """Uploaded image par prediction karta hai."""
    if model is None:
        return "Model Load Failed", 0.0

    # PIL image ko array mein convert karein
    img = Image.open(image_file).convert("RGB")
    # Image ko model ke input size mein resize karein (224x224)
    img = img.resize(IMAGE_SIZE)
    img_array = np.array(img)
    
    # Batch dimension add karein: (224, 224, 3) se (1, 224, 224, 3)
    img_array = np.expand_dims(img_array, axis=0)
    
    # Normalization (jaisa training mein kiya tha: 1./255)
    img_array = img_array / 255.0

    # Prediction karein
    predictions = model.predict(img_array)
    
    # Highest probability index aur score nikalien
    predicted_index = np.argmax(predictions, axis=1)[0]
    confidence_score = np.max(predictions) * 100

    predicted_class = CLASS_NAMES[predicted_index]
    
    return predicted_class, confidence_score

# --- 4. Streamlit UI ---

st.set_page_config(page_title="Brain Tumor Detection", layout="wide")

# a) Page Title
st.title("🧠 Brain Tumor Detection System (AI Powered)")
st.write("Upload an MRI image below to classify it as one of the tumor types or no tumor.")
st.markdown("---")

col1, col2 = st.columns(2)

with col1:
    # b) Image Upload Section
    uploaded_file = st.file_uploader("Upload MRI Image:", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        # Image Preview
        st.image(uploaded_file, caption="Uploaded MRI Image", use_column_width=True)
        st.markdown("---")

        # c) Prediction Button
        if st.button("Detect Tumor"):
            st.spinner("Analyzing image and detecting tumor...")
            
            # Prediction
            predicted_class, confidence_score = predict_image(uploaded_file, model)
            
            # d) Output Section
            
            # Simplified classification for the app display (Tumor / No Tumor)
            if predicted_class == 'no_tumor':
                result_label = f"🟢 **Prediction: No Tumor**"
            elif predicted_class == 'Model Load Failed':
                result_label = f"❌ **Prediction: Model Initialization Error**"
                st.error("Model could not be loaded for prediction. Please check logs for Hugging Face download errors.")
            else:
                result_label = f"🔴 **Prediction: Tumor ({predicted_class.replace('_', ' ').title()})**"
            
            st.success("✅ Analysis Complete")
            st.subheader(result_label)
            st.metric(label="Confidence Score", value=f"{confidence_score:.2f}%")
            
            st.write("---")

with col2:
    st.header("Results and Interpretation")
    st.info("The system uses Transfer Learning (VGG16) to classify the image into four categories: Glioma, Meningioma, Pituitary, or No Tumor.")
    
    if uploaded_file is None:
        st.warning("Please upload an image and click 'Detect Tumor' to see the results.")