"""
app.py — CropGuard FastAPI inference server
============================================
Loads the trained MobileNetV2 model and exposes a REST API used by both
the React frontend and the single-file HTML app.

Endpoints
---------
GET  /health           -> {"status": "ok", "model_loaded": bool}
POST /predict          -> multipart image -> diagnosis JSON
GET  /diseases         -> the full treatment knowledge base

Diagnosis JSON shape (consumed by the frontends):
{
  "class_id": "tomato_late",
  "confidence": 0.94,
  "severity": "moderate",            # null when healthy
  "diseased_ratio": 0.42,
  "disease": { ...full record from recommendations.json... }
}

Run:
    uvicorn app:app --host 0.0.0.0 --port 8000 --reload
"""
import io, json, os
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image

MODEL_DIR = os.getenv("MODEL_DIR", "model")
IMG_SIZE = 224
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

app = FastAPI(title="CropGuard GH API", version="1.0")
app.add_middleware(
    CORSMiddleware, allow_origins=["*"],
    allow_methods=["*"], allow_headers=["*"],
)

# --- lazy globals ---
_model = None
_classes = None
with open(os.path.join(os.path.dirname(__file__), "recommendations.json")) as f:
    RECS = json.load(f)


def get_model():
    """Load the Keras model once and keep it resident in memory (§3.10.2)."""
    global _model, _classes
    if _model is None:
        import tensorflow as tf
        _model = tf.keras.models.load_model(os.path.join(MODEL_DIR, "crop_model.keras"))
        with open(os.path.join(MODEL_DIR, "classes.json")) as f:
            _classes = json.load(f)
    return _model, _classes


def preprocess(img: Image.Image):
    img = img.convert("RGB").resize((IMG_SIZE, IMG_SIZE))
    arr = np.asarray(img, dtype=np.float32) / 255.0
    arr = (arr - MEAN) / STD
    return np.expand_dims(arr, 0)


def estimate_severity(img: Image.Image):
    """Diseased-area ratio via colour thresholding (§3.8).
    Returns (severity_label, diseased_ratio)."""
    small = np.asarray(img.convert("RGB").resize((128, 128)), dtype=np.float32)
    r, g, b = small[..., 0], small[..., 1], small[..., 2]
    lum = small.mean(axis=2)
    mx = small.max(axis=2); mn = small.min(axis=2); sat = mx - mn
    leaf = ~((lum > 235) & (sat < 25))                     # drop white background
    yellow = (r > 120) & (g > 100) & (b < 100) & (r - b > 40)
    dark = lum < 70
    brown = (r > 70) & (r > g) & (g > b) & (r - b > 18) & (lum < 170)
    leaf_px = max(int(leaf.sum()), 1)
    ratio = float(((yellow & leaf).sum() * 0.85 +
                   (brown & leaf).sum() * 1.05 +
                   (dark & leaf).sum() * 1.10) / leaf_px)
    ratio = min(ratio, 1.0)
    label = "early" if ratio < 0.20 else "moderate" if ratio < 0.55 else "severe"
    return label, round(ratio, 3)


@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _model is not None}


@app.get("/diseases")
def diseases():
    return RECS


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        raw = await file.read()
        img = Image.open(io.BytesIO(raw))
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid image file")

    model, classes = get_model()
    probs = model.predict(preprocess(img), verbose=0)[0]
    idx = int(np.argmax(probs))
    class_id = classes[idx]
    confidence = float(probs[idx])

    record = RECS.get(class_id, {})
    if record.get("healthy"):
        severity, ratio = None, 0.0
    else:
        severity, ratio = estimate_severity(img)

    # The uploaded image is never stored — it is discarded here (§3.10.2 privacy).
    return {
        "class_id": class_id,
        "confidence": round(confidence, 4),
        "severity": severity,
        "diseased_ratio": ratio,
        "disease": record,
    }