# ==============================
# Imports
# ==============================
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import numpy as np
import cv2
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import (
    InputLayer, TimeDistributed, GlobalAveragePooling2D,
    LSTM, Dense, Dropout
)
from tensorflow.keras.applications import MobileNetV2
import os

# ==============================
# Parameters
# ==============================
IMG_H, IMG_W, CHANNELS = 224, 224, 3
NUM_FRAMES = 30
MODEL_PATH = "video_mobilenetv2_final_backup.keras"

model = None  # Lazy loading

# ==============================
# Model Loader
# ==============================
def load_model():
    global model
    if model is not None:
        return model

    print("🔄 Loading model...")

    base_cnn = MobileNetV2(
        include_top=False, 
        weights='imagenet',
        input_shape=(IMG_H, IMG_W, CHANNELS)
    )

    for layer in base_cnn.layers[:-20]:
        layer.trainable = False

    model = Sequential([
        InputLayer(shape=(NUM_FRAMES, IMG_H, IMG_W, CHANNELS)),
        TimeDistributed(base_cnn),
        TimeDistributed(GlobalAveragePooling2D()),
        LSTM(128),
        Dropout(0.5),
        Dense(128, activation='relu'),
        Dropout(0.5),
        Dense(1, activation='sigmoid')
    ])

    model.load_weights(MODEL_PATH)
    print("✅ Model loaded successfully!")
    return model

# ==============================
# Frame Extraction
# ==============================
def extract_video_frames(video_path, n_frames=NUM_FRAMES):
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total <= 0:
        blank = np.zeros((IMG_H, IMG_W, 3), dtype=np.float32)
        return np.stack([blank]*n_frames)

    idxs = np.linspace(0, total - 1, n_frames).astype(int)
    frames = []

    for i in idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ok, frame = cap.read()

        if not ok:
            frames.append(np.zeros((IMG_H, IMG_W, 3), dtype=np.float32))
            continue

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = cv2.resize(frame, (IMG_W, IMG_H))
        frame = frame.astype(np.float32) / 255.0
        frames.append(frame)

    cap.release()
    return np.array(frames, dtype=np.float32)

# ==============================
# Predict
# ==============================
def predict_video(path):
    model = load_model()
    frames = extract_video_frames(path)
    frames = np.expand_dims(frames, 0)
    score = model.predict(frames, verbose=0)[0][0]
    label = "Violence" if score >= 0.5 else "NonViolence"
    conf = score*100 if score >= 0.5 else (1 - score)*100
    return label, conf

# ==============================
# FastAPI App
# ==============================
app = FastAPI(
    title="Video Violence Detector",
    description="API for detecting violence in videos using deep learning",
    version="1.0.0"
)

# ==============================
# Startup Event - Preload Model
# ==============================
@app.on_event("startup")
async def startup_event():
    """Preload model on startup to avoid cold start delays"""
    print("🚀 Starting up application...")
    try:
        load_model()
        print("✅ Model preloaded successfully!")
    except Exception as e:
        print(f"⚠️ Warning: Could not preload model: {e}")

# ==============================
# Health Check Endpoint
# ==============================
@app.get("/health")
def health_check():
    """Health check endpoint for monitoring"""
    return {
        "status": "healthy",
        "model_loaded": model is not None
    }

# ==============================
# Root Endpoint
# ==============================
@app.get("/")
def home():
    """Root endpoint with API information"""
    return {
        "message": "Video Violence Detector API",
        "endpoints": {
            "POST /predict": "Upload a video file for violence detection",
            "GET /health": "Check API health status"
        },
        "usage": "Send a POST request to /predict with a video file"
    }

# ==============================
# Prediction Endpoint
# ==============================
@app.post("/predict")
async def upload(file: UploadFile = File(...)):
    """
    Upload a video file and get violence prediction
    
    Returns:
        JSON with prediction (Violence/NonViolence) and confidence score
    """
    # Validate file type
    if not file.content_type.startswith('video/'):
        return JSONResponse(
            status_code=400, 
            content={"error": "File must be a video"}
        )
    
    temp = f"temp_{file.filename}"
    
    try:
        # Save uploaded file
        with open(temp, "wb") as f:
            f.write(await file.read())
        
        # Make prediction
        label, conf = predict_video(temp)
        
        # Clean up temp file
        if os.path.exists(temp):
            os.remove(temp)
        
        return {
            "prediction": label,
            "confidence": f"{conf:.2f}%",
            "filename": file.filename
        }
        
    except Exception as e:
        # Clean up on error
        if os.path.exists(temp):
            os.remove(temp)
        
        return JSONResponse(
            status_code=500, 
            content={
                "error": "Prediction failed",
                "details": str(e)
            }
        )

# ==============================
# Run locally (not used in HF)
# ==============================
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)