import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from PIL import Image, ImageOps
import io
import tensorflow as tf
import numpy as np
from tensorflow.keras.applications.mobilenet_v3 import preprocess_input
import os
import traceback
from huggingface_hub import hf_hub_download

os.environ['YOLO_CONFIG_DIR'] = '/tmp'
os.environ["YOLO_OFFLINE"] = "True"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

app = FastAPI(title="ChiliGuard API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

CLASS_NAMES = ['Anthracnose', 'Cercospora', 'Fresh Leaf', 'Leaf Curl']

yolo_model = None
mobilenet_model = None

def load_yolo():
    global yolo_model
    print("🔄 Loading YOLOv11...")
    try:
        # Download first then load - more reliable in HF Space
        model_path = hf_hub_download(
            repo_id="MdMahamudulHasan/chili-leaf-detection", 
            filename="YOLOV11nbest.pt"
        )
        yolo_model = YOLO(model_path)
        print("✅ YOLOv11 Loaded Successfully!")
        return True
    except Exception as e:
        print("❌ YOLO Load Failed:")
        print(traceback.format_exc())
        return False

def load_mobilenet():
    global mobilenet_model
    print("🔄 Loading MobileNetV3...")
    try:
        model_path = hf_hub_download(
            repo_id="MdMahamudulHasan/chili-leaf-classification", 
            filename="mobilenetv3_chili_leaf_global.keras"
        )
        mobilenet_model = tf.keras.models.load_model(model_path, compile=False)
        print("✅ MobileNetV3 Loaded Successfully!")
        return True
    except Exception as e:
        print("❌ MobileNet Load Failed:")
        print(traceback.format_exc())
        return False

# Load models
print("=== Model Loading Started ===")
yolo_loaded = load_yolo()
mobilenet_loaded = load_mobilenet()
print("=== Model Loading Finished ===\n")

@app.get("/")
def read_root():
    return {
        "message": "ChiliGuard API is Running",
        "yolo_status": "Loaded" if yolo_model else "Not Loaded",
        "mobilenet_status": "Loaded" if mobilenet_model else "Not Loaded"
    }

@app.post("/predict")
async def predict(image: UploadFile = File(...)):
    if not yolo_model or not mobilenet_model:
        return {"error": "Models are still initializing. Try again in a few seconds."}

    # 1. Read and fix image orientation
    image_bytes = await image.read()
    img = Image.open(io.BytesIO(image_bytes))
    img = ImageOps.exif_transpose(img).convert("RGB")

    results_data = {}

    # 2. YOLO Bounding Box Extraction
    try:
        yolo_results = yolo_model(img, imgsz=640, conf=0.15, verbose=False) 
        yolo_res = yolo_results[0]
        boxes = []
        if hasattr(yolo_res, 'boxes') and yolo_res.boxes is not None:
            for box in yolo_res.boxes:
                coords = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
                conf = float(box.conf.item()) * 100
                cls_id = int(box.cls.item())
                boxes.append({
                    "id": cls_id,
                    "bbox": coords,
                    "yolo_label": yolo_res.names[cls_id],
                    "yolo_confidence": conf
                })
        results_data["yolo"] = {"boxes": boxes}
    except Exception as e:
        print(f"YOLO inference error: {e}")
        results_data["yolo"] = {"error": str(e), "boxes": []}

    # 3. MobileNet Disease Detection
    try:
        # Preprocessing
        img_resized = img.resize((224, 224), Image.NEAREST) 
        img_array = np.asarray(img_resized, dtype=np.float32)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = preprocess_input(img_array)
        
        # Inference
        dataset_pred = mobilenet_model(img_array, training=False)
        probability = dataset_pred.numpy()[0]
        predicted_class_index = np.argmax(probability)
        predicted_class = CLASS_NAMES[predicted_class_index]
        confidence = float(probability[predicted_class_index]) * 100
        
        results_data["mobilenet"] = {
            "disease": predicted_class,
            "confidence": confidence
        }

        # 4. Integrate YOLO boxes with MobileNet labels
        # We use YOLO for WHERE the disease is (with YOLO confidence), but MobileNet for WHAT the disease is.
        for box in results_data["yolo"]["boxes"]:
            box["label"] = predicted_class  # MobileNet's disease name
            # Keep YOLO's confidence (box already has "yolo_confidence" key)

    except Exception as e:
        print(f"MobileNet inference error: {e}")
        results_data["mobilenet"] = {"error": str(e)}

    return results_data

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