from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from ultralytics import YOLO
import numpy as np
from PIL import Image
import io
import json
import os
import logging
import threading

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = FastAPI(title="Plastic Detection API", version="1.0.0")

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

# Global model state
model = None
model_loading = False
model_ready = False

def preload_model():
    """Load model in background thread"""
    global model, model_loading, model_ready
    if model is None and not model_loading:
        model_loading = True
        logger.info("🔄 Pre-loading YOLO model in background...")
        try:
            model_path = "runs/plastic_detection/weights/best.pt"
            if not os.path.exists(model_path):
                logger.warning(f"Model not found: {model_path}, using yolov8n.pt")
                model_path = "yolov8n.pt"
            model = YOLO(model_path)
            model_ready = True
            logger.info(f"✅ Model loaded: {model_path}")
        except Exception as e:
            logger.error(f"❌ Failed to load model: {e}")
            model_loading = False

def get_model():
    """Get model, load if needed"""
    global model, model_loading, model_ready
    if model is None:
        if not model_loading:
            preload_model()
        # Wait for model to load (with timeout)
        timeout = 120  # 2 minutes
        waited = 0
        while model is None and waited < timeout:
            import time
            time.sleep(0.5)
            waited += 0.5
        if model is None:
            raise RuntimeError("Model failed to load")
    return model

@app.on_event("startup")
async def startup_event():
    """Pre-load model on startup"""
    logger.info("🚀 Starting Plastic Detection API...")
    # Load model in background thread
    thread = threading.Thread(target=preload_model, daemon=True)
    thread.start()

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "message": "Plastic Detection API is running",
        "status": "ok",
        "model_ready": model_ready
    }

@app.post("/predict")
@app.post("/Predict")
async def predict(file: UploadFile = File(None)):
    """
    Predict plastic type in an image
    
    Args:
        file: Image file (multipart form upload)
    
    Returns:
        JSON with predictions
    """
    logger.info(f"📥 Prediction request received, model_ready={model_ready}")
    try:
        if not file:
            raise HTTPException(status_code=400, detail="No file provided")
            
        model_instance = get_model()
        
        # Read image file
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        logger.info(f"✅ Image loaded: {image.size}")
        
        # Run inference
        logger.info("🔮 Running YOLO inference...")
        results = model_instance.predict(source=image, conf=0.25, verbose=False)
        
        # Extract predictions
        predictions = []
        if results and len(results) > 0:
            result = results[0]
            class_names = result.names
            
            if result.boxes is not None and len(result.boxes) > 0:
                for box in result.boxes:
                    pred = {
                        "class": class_names[int(box.cls)],
                        "class_id": int(box.cls),
                        "confidence": float(box.conf),
                        "bbox": {
                            "x1": float(box.xyxy[0][0]),
                            "y1": float(box.xyxy[0][1]),
                            "x2": float(box.xyxy[0][2]),
                            "y2": float(box.xyxy[0][3]),
                        }
                    }
                    predictions.append(pred)
        
        logger.info(f"✅ Found {len(predictions)} detections")
        return {
            "success": True,
            "detections": predictions,
            "count": len(predictions),
            "model": "YOLOv8 Best",
            "filename": file.filename
        }
    
    except Exception as e:
        logger.error(f"❌ Error: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")

@app.get("/classes")
async def get_classes():
    """Get all detectable classes"""
    logger.info("Classes request")
    try:
        model_instance = get_model()
        classes = model_instance.names
        logger.info(f"Returning {len(classes)} classes")
        return {
            "classes": classes,
            "total": len(classes)
        }
    except Exception as e:
        logger.error(f"Error: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=f"Error: {str(e)}")

@app.get("/stats")
async def get_stats():
    """Get backend statistics"""
    return {
        "status": "running",
        "model_ready": model_ready,
        "model_loading": model_loading,
        "available_models": ["best.pt", "yolov8n.pt"]
    }

if __name__ == "__main__":
    print("🚀 Starting Plastic Detection API...")
    print("🌐 Available at: http://0.0.0.0:8000")
    print("📖 Docs: http://localhost:8000/docs")
    print("⏳ Model loads on first request...")
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")