main.py

"""
KhmerCrop AI — FastAPI backend (Proposal §6: FastAPI + Python + CNN)

Flow implemented here (Proposal §5):
  1. Client uploads a photo of a diseased leaf.
  2. Server runs it through the matching crop's model.
  3. Server responds with the predicted disease/pest name + confidence.

Inference backend is controlled by USE_GPU below:
  - USE_GPU = True  -> loads khmercrop_{crop}_model.keras via plain TensorFlow (uses GPU/CUDA
                       automatically if available; falls back to CPU if no GPU is present)
  - USE_GPU = False -> loads khmercrop_{crop}_model.tflite via ai_edge_litert (lightweight,
                       CPU-only — this is what Hugging Face Spaces free tier should use)

Explicitly OUT OF SCOPE for this file (owned by teammates, per your instructions):
  - Web Frontend (HTML/CSS/JS, SolidJS, Tailwind)
  - Database (SQLite/PostgreSQL) — see the TODOs marked `# DB HOOK` below for exactly
    where a teammate should plug in persistence for upload history / results.

Run locally:
    uv run uvicorn main:app --reload
"""

import io
import os
import pathlib
from typing import Dict, List

import numpy as np
from fastapi import FastAPI, File, HTTPException, UploadFile
from PIL import Image
from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Toggle: True = Keras model on GPU (local dev with CUDA), False = TFLite on CPU
# (use this for Hugging Face Spaces / any CPU-only deployment)
# ---------------------------------------------------------------------------
USE_GPU = True

try:
    import tensorflow as tf
except ImportError:
    USE_GPU = False
    import ai_edge_litert.interpreter as tflite

try:
    import spaces
except ImportError:
    class MockSpaces:
        def GPU(self, func):
            return func
    spaces = MockSpaces()


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

# Parent directory containing potato/export, rice/export, corn/export, citrus/export
# — i.e. the same PROJECT_DIR the training notebook writes into.
MODELS_DIR = pathlib.Path(os.environ.get("MODELS_DIR", "./models"))

# Matches the crop names used in the training notebook (Train_pipeline_ipynb.ipynb)
SUPPORTED_CROPS = ["potato", "rice", "corn", "citrus"]


# ---------------------------------------------------------------------------
# Model registry — loads all 4 models once at startup, not per-request
# ---------------------------------------------------------------------------

class CropModel:
    def __init__(self, crop: str, export_dir: pathlib.Path):
        self.crop = crop
        labels_path = export_dir / f"class_names_{crop}.txt"
        if not labels_path.exists():
            raise FileNotFoundError(f"[{crop}] Missing labels file: {labels_path}.")
        with open(labels_path, encoding="utf-8") as f:
            self.class_names = [line.strip() for line in f if line.strip()]

        if USE_GPU:
            keras_path = export_dir / f"khmercrop_{crop}_model.keras"
            if not keras_path.exists():
                raise FileNotFoundError(
                    f"[{crop}] Missing model file: {keras_path}. "
                    f"USE_GPU=True requires a .keras export for this crop."
                )
            self.model = tf.keras.models.load_model(str(keras_path))
            _, h, w, _ = self.model.input_shape
            self.img_size = (int(w), int(h))
        else:
            tflite_path = export_dir / f"khmercrop_{crop}_model.tflite"
            if not tflite_path.exists():
                raise FileNotFoundError(
                    f"[{crop}] Missing model file: {tflite_path}. "
                    f"USE_GPU=False requires a .tflite export for this crop."
                )
            self.interpreter = tflite.Interpreter(model_path=str(tflite_path))
            self.interpreter.allocate_tensors()
            self.input_details = self.interpreter.get_input_details()
            self.output_details = self.interpreter.get_output_details()
            # img_size comes straight from the model's own input tensor shape, so it can
            # never drift out of sync with what the notebook actually trained on.
            _, h, w, _ = self.input_details[0]["shape"]
            self.img_size = (int(w), int(h))

    @spaces.GPU
    def predict(self, image: Image.Image) -> Dict:
        image = image.convert("RGB").resize(self.img_size)
        # Preprocessing (MobileNetV2 preprocess_input) is baked into the exported model
        # graph itself (see build_model() in the notebook), so we feed raw 0-255 floats.
        arr = np.asarray(image, dtype=np.float32)[np.newaxis, ...]

        if USE_GPU:
            output = self.model.predict(arr, verbose=0)[0]
        else:
            self.interpreter.set_tensor(self.input_details[0]["index"], arr)
            self.interpreter.invoke()
            output = self.interpreter.get_tensor(self.output_details[0]["index"])[0]

        top_idx = int(np.argmax(output))
        return {
            "crop": self.crop,
            "label": self.class_names[top_idx],
            "confidence": float(output[top_idx]),
            "all_scores": {
                name: float(score) for name, score in zip(self.class_names, output)
            },
        }


MODELS: Dict[str, CropModel] = {}


def load_models() -> List[str]:
    """Load every crop model that's actually present, skipping (and logging) any
    that aren't exported yet instead of crashing the whole server on startup."""
    loaded, skipped = [], []
    for crop in SUPPORTED_CROPS:
        export_dir = MODELS_DIR / crop / "export"
        try:
            MODELS[crop] = CropModel(crop, export_dir)
            loaded.append(crop)
        except FileNotFoundError as e:
            skipped.append(crop)
            print(f"[startup] Skipping '{crop}': {e}")
    backend = "Keras/GPU" if USE_GPU else "TFLite/CPU"
    print(f"[startup] Backend: {backend} | Loaded models: {loaded or 'NONE'} | Skipped: {skipped or 'none'}")
    return loaded


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------

app = FastAPI(title="KhmerCrop AI Backend", version="1.0.0")


class PredictionResponse(BaseModel):
    crop: str
    label: str
    confidence: float


@app.on_event("startup")
def on_startup():
    load_models()


@app.get("/health")
def health():
    return {
        "status": "ok",
        "backend": "keras-gpu" if USE_GPU else "tflite-cpu",
        "models_loaded": list(MODELS.keys()),
    }


@app.get("/crops")
def list_crops():
    """What crops/classes are currently servable — lets the frontend build its
    crop-selector without hardcoding class lists."""
    return {
        crop: {"num_classes": len(m.class_names), "class_names": m.class_names}
        for crop, m in MODELS.items()
    }


@app.post("/predict/{crop}", response_model=PredictionResponse)
async def predict(crop: str, file: UploadFile = File(...)):
    crop = crop.lower()
    if crop not in SUPPORTED_CROPS:
        raise HTTPException(
            status_code=404,
            detail=f"Unknown crop '{crop}'. Supported: {SUPPORTED_CROPS}",
        )
    if crop not in MODELS:
        raise HTTPException(
            status_code=503,
            detail=f"Model for '{crop}' isn't loaded (not exported yet on this server).",
        )
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    try:
        raw = await file.read()
        image = Image.open(io.BytesIO(raw))
    except Exception:
        raise HTTPException(status_code=400, detail="Could not read the uploaded image.")

    result = MODELS[crop].predict(image)

    # DB HOOK: a teammate should persist this prediction here for the "History" /
    # "Result Dashboard" web features, e.g.:
    #   db.save_prediction(crop=crop, label=result["label"],
    #                       confidence=result["confidence"], image=raw, user=...)
    # Left out intentionally — Database (SQLite/PostgreSQL) is out of scope for this file.

    return PredictionResponse(
        crop=result["crop"], label=result["label"], confidence=result["confidence"]
    )


if __name__ == "__main__":
    import uvicorn
    print("[main] Starting server on http://localhost:8001")
    uvicorn.run("main:app", host="0.0.0.0", port=3002, reload=True)