from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import numpy as np
import cv2
import base64

from sentinelscan.modeling.infer import CrackModel
from sentinelscan.utils.viz import overlay_mask
from sentinelscan.utils.severity import crack_metrics, severity_from_metrics

app = FastAPI(title="SentinelScan API", version="0.1.0")

model = CrackModel(ckpt_path="models/best.pt", size=512)

def _read_image(file_bytes: bytes):
    arr = np.frombuffer(file_bytes, np.uint8)
    bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if bgr is None:
        raise ValueError("Could not decode image")
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    return rgb

def _to_base64_png(rgb: np.ndarray):
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    ok, buf = cv2.imencode(".png", bgr)
    if not ok:
        raise ValueError("Could not encode image")
    return base64.b64encode(buf.tobytes()).decode("utf-8")

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        rgb = _read_image(await file.read())
        pred = model.predict(rgb, threshold=0.5)

        m = crack_metrics(pred["mask"])
        sev = severity_from_metrics(m)

        crack_detected = m["area_px"] > 50  # tiny specks ignored

        overlay = overlay_mask(rgb, pred["mask"])
        overlay_b64 = _to_base64_png(overlay)

        return JSONResponse({
            "crack_detected": bool(crack_detected),
            "confidence": float(pred["confidence"]),
            "severity": sev if crack_detected else "None",
            "metrics": m,
            "overlay_png_base64": overlay_b64,
        })
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=400)