from fastapi import FastAPI, File, UploadFile, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from plantdoc_predictor import Predictor
from PIL import Image
import io

app = FastAPI(title="PlantDoc API", version="1.0.0")

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

MODELS = [
    {"name": "mobilenetv2_v1",          "accuracy": 96.8,  "description": "Mobile-optimized, fastest inference"},
    {"name": "densenet169_v1",           "accuracy": 99.68, "description": "Best classical CNN"},
    {"name": "rec_add_attention_v1",     "accuracy": 99.70, "description": "Best overall — Recursive Additive Attention (research model)"},
    {"name": "convnext_small_v1",        "accuracy": 99.50, "description": "High accuracy + efficient"},
    {"name": "convnext_tiny_v1",         "accuracy": 99.23, "description": "Lightweight ConvNeXt"},
    {"name": "convnext_base_v1",         "accuracy": 99.10, "description": "ConvNeXt base"},
    {"name": "swin_tiny_patch4_window7", "accuracy": 99.10, "description": "Lightweight Swin Transformer"},
    {"name": "swin_base_patch4_window7", "accuracy": 99.10, "description": "Swin Transformer"},
    {"name": "vit_tiny_16_v1",           "accuracy": 99.10, "description": "Lightweight Vision Transformer"},
    {"name": "vit_small_16_v1",          "accuracy": 99.10, "description": "Small Vision Transformer"},
    {"name": "vit_base_16_v1",           "accuracy": 99.10, "description": "Vision Transformer base"},
    {"name": "vit_large_16_v1",          "accuracy": 99.10, "description": "Vision Transformer large"},
    {"name": "densenet121_v1",           "accuracy": 98.68, "description": "Efficient DenseNet"},
    {"name": "vgg19_v1",                 "accuracy": 98.98, "description": "Deep VGG"},
    {"name": "inceptionv3_v1",           "accuracy": 98.20, "description": "Multi-scale InceptionV3"},
    {"name": "resnet50_v1",              "accuracy": 97.80, "description": "Residual learning CNN"},
    {"name": "efficientnetb50_v1",       "accuracy": 97.80, "description": "EfficientNet balanced"},
    {"name": "densenet210_v1",           "accuracy": 97.00, "description": "Deep DenseNet"},
    {"name": "vgg16_v1",                 "accuracy": 96.80, "description": "Classic CNN"},
    {"name": "alexnet_v1",               "accuracy": 92.80, "description": "Baseline CNN"},
]

MODEL_NAMES = [m["name"] for m in MODELS]

_predictor_cache: dict[str, Predictor] = {}

def get_predictor(model_name: str) -> Predictor:
    if model_name not in _predictor_cache:
        _predictor_cache[model_name] = Predictor(model_name=model_name, verbose=False)
    return _predictor_cache[model_name]


@app.on_event("startup")
def warmup():
    get_predictor("mobilenetv2_v1")


@app.get("/health")
def health():
    return {"status": "ok", "loaded_models": list(_predictor_cache.keys())}


@app.get("/models")
def list_models():
    return {"models": MODELS}


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    model: str = Query(default="mobilenetv2_v1"),
):
    if model not in MODEL_NAMES:
        raise HTTPException(status_code=400, detail=f"Unknown model '{model}'. Call /models for the list.")

    content = await file.read()
    if len(content) == 0:
        raise HTTPException(status_code=400, detail="Empty file uploaded.")

    try:
        image = Image.open(io.BytesIO(content)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file.")

    try:
        predictor = get_predictor(model)
        result = predictor.predict(image)
        plant, disease = (result["label"].split("___") + [""])[:2]
        result["plant"] = plant.replace("_", " ")
        result["disease"] = disease.replace("_", " ")
        result["is_healthy"] = "healthy" in result["label"].lower()
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))