"""FastAPI application for phenology classification."""

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import torch
from PIL import Image
import io
import json
import os
from typing import Dict

from src.inference import load_inference_model, preprocess_image, predict

# Initialize FastAPI app
app = FastAPI(
    title="ResNet Phenology Classifier API",
    description="API for classifying plant images by phenological phase",
    version="1.0.0"
)

# Global variables for model
model = None
classes = None
device = None

# Configuration
MODEL_PATH = os.getenv("MODEL_PATH", "models/best_model.pth")
CLASS_MAPPING_PATH = os.getenv("CLASS_MAPPING_PATH", "models/class_mapping.json")


@app.on_event("startup")
async def load_model():
    """Load model on startup."""
    global model, classes, device
    try:
        model, classes, device = load_inference_model(MODEL_PATH, CLASS_MAPPING_PATH)
        print(f"Model loaded successfully from {MODEL_PATH}")
        print(f"Classes: {classes}")
    except Exception as e:
        print(f"Error loading model: {e}")
        raise


@app.get("/")
async def root():
    """Root endpoint."""
    return {
        "message": "ResNet Phenology Classifier API",
        "version": "1.0.0",
        "endpoints": {
            "classify": "/classify",
            "health": "/health",
            "classes": "/classes"
        }
    }


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    return {"status": "healthy", "model_loaded": True}


@app.get("/classes")
async def get_classes():
    """Get list of available classes."""
    if classes is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    return {"classes": classes, "num_classes": len(classes)}


@app.post("/classify")
async def classify_image(file: UploadFile = File(...)) -> Dict:
    """
    Classify a plant image.
    
    Args:
        file: Uploaded image file
        
    Returns:
        Dictionary with classification results
    """
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    # Validate file type
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    try:
        # Read image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert('RGB')
        
        # Preprocess
        from torchvision import transforms
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image_tensor = transform(image).unsqueeze(0)
        
        # Predict
        result = predict(model, image_tensor, classes, device)
        
        return JSONResponse(content=result)
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")


@app.post("/classify/batch")
async def classify_batch(files: list[UploadFile] = File(...)) -> Dict:
    """
    Classify multiple images.
    
    Args:
        files: List of uploaded image files
        
    Returns:
        Dictionary with classification results for each image
    """
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    if len(files) > 10:
        raise HTTPException(status_code=400, detail="Maximum 10 images per batch")
    
    results = []
    
    for idx, file in enumerate(files):
        if not file.content_type.startswith("image/"):
            results.append({
                "filename": file.filename,
                "error": "File must be an image"
            })
            continue
        
        try:
            # Read image
            contents = await file.read()
            image = Image.open(io.BytesIO(contents)).convert('RGB')
            
            # Preprocess
            from torchvision import transforms
            transform = transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])
            image_tensor = transform(image).unsqueeze(0)
            
            # Predict
            result = predict(model, image_tensor, classes, device)
            result["filename"] = file.filename
            results.append(result)
        
        except Exception as e:
            results.append({
                "filename": file.filename,
                "error": str(e)
            })
    
    return JSONResponse(content={"results": results, "total": len(files)})


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