ml/serving/server.py

# -*- coding: utf-8 -*-
from __future__ import annotations
import io
import os
import time
import uuid
import json
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from fastapi.middleware.cors import CORSMiddleware
import numpy as np
import requests
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

from .model_loader import ModelBundle
# postprocess no longer required – handled inside ModelBundle
# from .postprocess import postprocess_yolov8_seg
from .schemas import (
    PredictRequest, PredictResponse, Box, FeedbackRequest, FeedbackResponse,
    AnalyzeRequest, AnalyzeResponse, AnalyzeSummary, DetectedObject,
)
from .labels import to_bucket
from .weights import WEIGHT_SOURCE, estimate_weight_g, load_priors
from .materials import materials_catalog, display_for
from .products import lookup_product, valid_barcode
from . import dashboard as _dashboard

from .model_loader import dprint, DEBUG

# ---------- Config ----------
BUNDLE_DIR = Path(os.environ.get("ALAMI_AI_BUNDLE", "deploy/latest")).resolve()
if not BUNDLE_DIR.exists():
    raise RuntimeError(f"Bundle not found: {BUNDLE_DIR}")

FEEDBACK_LOG_DIR = Path(os.environ.get("ALAMI_FEEDBACK_DIR", "feedback_logs")).resolve()
FEEDBACK_LOG_DIR.mkdir(parents=True, exist_ok=True)
PRED_LOG = FEEDBACK_LOG_DIR / "predictions.jsonl"
FB_LOG = FEEDBACK_LOG_DIR / "feedback.jsonl"

# Supabase (optional; if ENV not set -> JSONL only)
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")

# Table names
SB_TABLE_PREDICTIONS = os.environ.get("SB_TABLE_PREDICTIONS", "trash_predictions")

# API keys (comma-separated). Unset => open access (dev / current mobile app).
# IMPORTANT: the live mobile app does not send a key yet — only enforce keys on
# a deployment AFTER the app ships the header, otherwise honest users stop
# getting paid (ai_status 'unavailable' => 0 TC, see AILLMHANDOFF.md §3).
API_KEYS = {k.strip() for k in os.environ.get("ALAMI_API_KEYS", "").split(",") if k.strip()}

# Outbound TLS verification for image fetches (default ON). Set
# ALAMI_INSECURE_FETCH=1 only for local testing against self-signed hosts.
INSECURE_FETCH = os.environ.get("ALAMI_INSECURE_FETCH") == "1"

app = FastAPI(title="Alami Vision API", version="2.1.0")
BUNDLE = ModelBundle(BUNDLE_DIR)
WEIGHT_PRIORS = load_priors(BUNDLE_DIR)


def require_api_key(x_api_key: Optional[str] = Header(None)):
    if not API_KEYS:
        return  # open mode
    if x_api_key not in API_KEYS:
        raise HTTPException(401, "invalid or missing API key (x-api-key header)")

dprint(f"bundle_dir={BUNDLE_DIR}")
dprint(f"imgsz={BUNDLE.imgsz}")
dprint(f"names(len)={len(BUNDLE.names)} => {BUNDLE.names[:10]}")
dprint(f"post_cfg={BUNDLE.post_cfg}")

# ---------- Supabase Client (lazy) ----------
_supabase = None

# CORS: env-driven (comma-separated origins). Default '*' — the API is consumed
# by native apps/servers; browser dashboards should pin their origin here.
_cors = [o.strip() for o in os.environ.get("ALAMI_CORS_ORIGINS", "*").split(",") if o.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors,
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

def get_supabase():
    """
    Initialize Supabase client only if all required ENV vars are present.
    Returns None if not usable.
    """
    global _supabase
    if _supabase is not None:
        return _supabase
    if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
        return None
    try:
        from supabase import create_client, Client  # pip install supabase
        _supabase = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)
        return _supabase
    except Exception as e:
    # No hard crash – keep logging to JSONL
        print(f"[WARN] Supabase client init failed: {e}")
        return None

# ---------- Utils ----------
def fetch_image_bytes(url: str) -> bytes:
    try:
        import requests

        if INSECURE_FETCH:
            import urllib3
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

        headers = {
            'User-Agent': 'alami-vision-api/2.0'
        }

        print(f"📥 Fetching image from: {url}")

        # TIMEOUTS setzen
        response = requests.get(
            url,
            timeout=(3.05, 8.0),  # (connect_timeout, read_timeout)
            headers=headers,
            verify=not INSECURE_FETCH
        )
        response.raise_for_status()
        
        max_size_mb = 10
        if len(response.content) > max_size_mb * 1024 * 1024:
            raise HTTPException(400, f"Image too large (> {max_size_mb}MB)")
        
        print(f"Successfully fetched {len(response.content)} bytes")
        return response.content
        
    except requests.exceptions.Timeout:
        print(f"Timeout fetching image: {url}")
        raise HTTPException(408, "Image fetch timeout")
    except requests.exceptions.RequestException as e:
        print(f"Network error: {e}")
        raise HTTPException(400, f"Failed to fetch image: {e}")
    except Exception as e:
        print(f"Unexpected error: {e}")
        raise HTTPException(400, f"Failed to fetch image: {e}")

def now_iso() -> str:
    import datetime as _dt
    return _dt.datetime.utcnow().isoformat() + "Z"

def append_jsonl(path: Path, obj: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(obj, ensure_ascii=False) + "\n")

def qc_weight_ok(w: Optional[float]) -> bool:
    if w is None:
        return True
    return 0.001 <= w <= 20.0

def save_latest_symlink(version_dir: Path):
    # convenience: deploy/latest -> deploy/<version>
    latest = version_dir.parent / "latest"
    if latest.exists() or latest.is_symlink():
        try:
            latest.unlink()
        except Exception:
            pass
    try:
        latest.symlink_to(version_dir.name)
    except Exception:
        pass

def top1_prediction(preds: List[Box]) -> Tuple[Optional[str], Optional[float]]:
    """
    Select highest confidence as Top-1 (label, confidence).
    """
    if not preds:
        return None, None
    best = max(preds, key=lambda p: p.conf)
    return best.label, float(best.conf)

# ---------- SB helpers ----------
def _is_missing_column_error(exc: Exception, column: str) -> bool:
    """
    True only when the PostgREST error clearly says `column` does not exist
    (schema not migrated yet, error code PGRST204 / "could not find ... column").
    Transient errors must NOT match — otherwise we'd silently store rows
    without their training payload.
    """
    msg = str(exc).lower()
    if column.lower() not in msg:
        return False
    return "pgrst204" in msg or "column" in msg or "schema cache" in msg


def sb_insert_prediction(
    prediction_id: str,
    user_id: Optional[str],
    image_url: str,
    predicted_type: Optional[str],
    predicted_weight_kg: Optional[float],
    confidence: Optional[float],
    model_version: str,
    predictions: Optional[List[Dict[str, Any]]] = None,
    source: Optional[str] = None,
) -> None:
    """
    Insert into trash_predictions. RLS must allow service role on server side.
    The full per-object `predictions` array (JSONB) is the training-data payload
    for the flywheel; if the column doesn't exist yet (schema not migrated),
    we retry with the minimal legacy payload so logging never breaks.
    """
    sb = get_supabase()
    if sb is None:
        return  # still OK (JSONL only)
    payload = {
        "prediction_id": prediction_id,
        "user_id": user_id,
        "image_url": image_url,
        "predicted_type": predicted_type,
        "predicted_weight_kg": predicted_weight_kg,  # aktuell None (Weight-Head später)
        "confidence": confidence,
        "model_version": model_version,
        # created_at via default now() in DB
    }
    # Optional columns (need the FLYWHEEL.md migration). If one is missing we
    # drop ONLY the column the error names and retry — transient errors never
    # silently downgrade the payload.
    optional: Dict[str, Any] = {}
    if predictions is not None:
        optional["predictions"] = predictions
    if source is not None:
        optional["source"] = source
    attempt = dict(payload, **optional)
    for _ in range(len(optional) + 1):
        try:
            sb.table(SB_TABLE_PREDICTIONS).insert(attempt).execute()
            dropped = set(optional) - set(attempt)
            if dropped:
                print(f"[WARN] Supabase insert: column(s) {sorted(dropped)} missing — stored "
                      "partial payload. Apply the trash_predictions migration in docs/FLYWHEEL.md.")
            return
        except Exception as e:
            missing = next((c for c in optional if c in attempt
                            and _is_missing_column_error(e, c)), None)
            if missing is None:
                print(f"[WARN] Supabase insert failed: {e}")
                return
            attempt.pop(missing)

def sb_update_feedback(
    prediction_id: str,
    corrected_type: Optional[str],
    corrected_weight_kg: Optional[float],
    corrected_items: Optional[List[Dict[str, Any]]] = None,
    source: Optional[str] = None,
    notes: Optional[str] = None,
    added_items: Optional[List[Dict[str, Any]]] = None,
    reasons: Optional[List[str]] = None,
) -> None:
    """
    Update corrected_* fields. Set corrected_at = now().
    Falls back to the minimal legacy payload if the new columns
    (corrected_items/feedback_source/notes/added_items/feedback_reasons) aren't
    migrated yet — the fallback drops only the missing columns, never the whole
    write.
    Notes must reach Supabase too: the local JSONL copy is ephemeral on the
    Space, and the weekly flywheel sync reads only from Supabase.
    """
    sb = get_supabase()
    if sb is None:
        return
    payload = {
        "corrected_type": corrected_type,
        "corrected_weight_kg": corrected_weight_kg,
        "corrected_at": now_iso()
    }
    full_payload = dict(payload)
    if corrected_items is not None:
        full_payload["corrected_items"] = corrected_items
    if source is not None:
        full_payload["feedback_source"] = source
    if notes is not None:
        full_payload["notes"] = notes
    if added_items is not None:
        full_payload["added_items"] = added_items
    if reasons is not None:
        full_payload["feedback_reasons"] = reasons
    # Try the full payload; on a missing-column error, drop ONLY the offending
    # column(s) and retry — never fall all the way back to minimal, so a not-yet-
    # migrated v2 column (added_items/feedback_reasons) can't take the already-
    # migrated ones (corrected_items/notes/…) down with it. The base `payload`
    # keys are never dropped (they aren't in the extra set).
    attempt = dict(full_payload)
    dropped_all: List[str] = []
    for _ in range(len(full_payload)):
        try:
            sb.table(SB_TABLE_PREDICTIONS).update(attempt).eq("prediction_id", prediction_id).execute()
            if dropped_all:
                print(f"[WARN] Supabase update: column(s) {sorted(dropped_all)} missing — stored the rest. "
                      "Apply the migration in docs/FLYWHEEL.md.")
            return
        except Exception as e:
            missing = [c for c in (attempt.keys() - payload.keys()) if _is_missing_column_error(e, c)]
            if not missing:
                print(f"[WARN] Supabase update failed: {e}")
                return
            for c in missing:
                attempt.pop(c, None)
            dropped_all.extend(missing)
    print("[WARN] Supabase update: exhausted column fallbacks; nothing written.")

# ---------- Routes ----------
REPO_ROOT = Path(__file__).resolve().parents[2]


@app.get("/dashboard", response_class=HTMLResponse)
def dashboard():
    """Zero-cost ML dashboard: model metrics, history, corpus monitor.
    Aggregate counts only — no PII — so it is safe to serve publicly."""
    data = _dashboard.gather_metrics(BUNDLE.bundle_dir, REPO_ROOT, get_supabase, SB_TABLE_PREDICTIONS)
    return HTMLResponse(_dashboard.render_html(data, model_version()))


@app.get("/v1/product/{barcode}", dependencies=[Depends(require_api_key)])
def product(barcode: str):
    """Barcode -> product sustainability info (Open Food/Products/Beauty Facts)
    + packaging mapped to Alami material buckets + German disposal guide.

    Completely separate from the model path — an upstream outage here can
    never affect /predict or the live camera preview. Always returns 200 with
    status: found | not_found | unavailable (except 400 for bad barcodes);
    the app branches on `status`, never on errors.
    The app MUST display the `attribution` (ODbL requirement)."""
    barcode = barcode.strip()
    if not valid_barcode(barcode):
        raise HTTPException(400, "invalid barcode (expected 6-14 digits)")
    result = lookup_product(barcode)
    append_jsonl(FEEDBACK_LOG_DIR / "product_scans.jsonl", {
        "ts": now_iso(), "barcode": barcode, "status": result["status"],
    })  # anonymous scan-interest stats — deliberately no user_id
    return result


@app.get("/v1/materials")
def materials():
    """Display metadata per material bucket (German label + overlay colour).
    The app fetches this once and caches it, so camera overlays are consistent.
    Static — no model call, no auth needed."""
    return {"materials": materials_catalog()}


@app.get("/healthz")
def healthz():
    return {
        "ok": True,
        "version": app.version,
        "imgsz": BUNDLE.imgsz,
        "names": BUNDLE.names,
        "conf_thr": BUNDLE.conf_thr,
        "supabase": bool(get_supabase() is not None),
        "auth": "api-key" if API_KEYS else "open",
        "timestamp": now_iso()
    }

def model_version() -> str:
    mv = BUNDLE.bundle_dir.name
    mc = (BUNDLE.bundle_dir / "model_card.json")
    if mc.exists():
        try:
            mvJson = json.loads(mc.read_text(encoding="utf-8"))
            mv = Path(mvJson.get("artifacts_dir", mv)).name
        except Exception:
            pass
    return mv


def run_inference(img_bytes: bytes) -> Tuple[List[Box], Dict[str, Any], int]:
    """
    Stage bytes to a temp file, run the bundle, return (boxes, raw_result, ms).
    Boxes carry label (material bucket) + raw_label (original model class).
    """
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        try:
            tmp.write(img_bytes)
            tmp.flush()
            tmp_path = Path(tmp.name)
            print(f"Image staged at: {tmp_path} ({len(img_bytes)} bytes)")
        except Exception as e:
            print(f"Failed to write temp file: {e}")
            raise HTTPException(400, f"Failed to stage image for inference: {e}")

    t0 = time.time()
    try:
        pred = BUNDLE.predict(tmp_path, return_masks=False)
    except Exception as e:
        raise HTTPException(500, f"Inference failed: {e}")
    finally:
        try:
            tmp_path.unlink(missing_ok=True)
        except Exception:
            pass
    dur = int((time.time() - t0) * 1000)

    preds: List[Box] = []
    for b in (pred.get("boxes") or []):
        cls_id = int(b.get("cls", 0))
        conf = float(b.get("conf", 0.0))
        xyxy = [float(x) for x in b.get("xyxy", [0, 0, 0, 0])]
        raw = BUNDLE.names[cls_id] if 0 <= cls_id < len(BUNDLE.names) else str(cls_id)
        preds.append(Box(xyxy=xyxy, cls=cls_id, conf=conf, label=to_bucket(raw), raw_label=raw))
    return preds, pred, dur


def log_prediction(prediction_id: str, image_ref: str, user_id: Optional[str],
                   mv: str, preds: List[Box], endpoint: str,
                   source: Optional[str] = None) -> None:
    append_jsonl(PRED_LOG, {
        "prediction_id": prediction_id,
        "ts": now_iso(),
        "image_url": image_ref,
        "user_id": user_id,
        "model_version": mv,
        "endpoint": endpoint,
        "source": source,
        "predictions": [p.model_dump() for p in preds]
    })
    top_label, top_conf = top1_prediction(preds)
    sb_insert_prediction(
        prediction_id=prediction_id,
        user_id=user_id,
        image_url=image_ref,
        predicted_type=top_label,
        predicted_weight_kg=None,
        confidence=top_conf,
        model_version=mv,
        predictions=[p.model_dump() for p in preds],
        source=source,
    )


@app.post("/predict", response_model=PredictResponse, dependencies=[Depends(require_api_key)])
def predict(req: PredictRequest):
    if not req.image_url:
        raise HTTPException(400, "image_url required")

    img_bytes = fetch_image_bytes(req.image_url)
    preds, _, dur = run_inference(img_bytes)
    mv = model_version()
    dprint(f"/predict result: count={len(preds)}")
    if preds:
        dprint(f"/predict sample[0]={preds[0]}")

    prediction_id = str(uuid.uuid4())
    log_prediction(prediction_id, req.image_url, req.user_id, mv, preds, endpoint="predict")

    return PredictResponse(
        model_version=mv,
        inference_ms=dur,
        predictions=preds,
        prediction_id=prediction_id
    )


def build_analyze_response(preds: List[Box], raw: Dict[str, Any], dur: int,
                           domain: str, prediction_id: str, mv: str) -> AnalyzeResponse:
    orig_h, orig_w = (raw.get("orig_shape") or [0, 0])
    img_area = float(max(orig_w, 1) * max(orig_h, 1))

    objects: List[DetectedObject] = []
    materials: Dict[str, int] = {}
    total_w = 0.0
    W = float(max(orig_w, 1))
    H = float(max(orig_h, 1))
    for p in preds:
        x1, y1, x2, y2 = p.xyxy
        area_frac = max(0.0, (x2 - x1)) * max(0.0, (y2 - y1)) / img_area
        w_est = estimate_weight_g(p.label, WEIGHT_PRIORS)
        total_w += w_est
        materials[p.label] = materials.get(p.label, 0) + 1
        disp = display_for(p.label)
        # normalized [x,y,w,h] in 0..1 so the app can draw the box on any
        # preview size without knowing the original image dimensions.
        bbox_norm = [
            round(min(max(x1 / W, 0.0), 1.0), 5),
            round(min(max(y1 / H, 0.0), 1.0), 5),
            round(min(max((x2 - x1) / W, 0.0), 1.0), 5),
            round(min(max((y2 - y1) / H, 0.0), 1.0), 5),
        ]
        objects.append(DetectedObject(
            label=p.label,
            label_en=disp["label_en"],
            label_de=disp["label_de"],
            raw_label=p.raw_label,
            confidence=p.conf,
            bbox_xyxy=p.xyxy,
            bbox_norm=bbox_norm,
            color=disp["color"],
            area_fraction=round(area_frac, 6),
            weight_estimate_g=w_est,
            weight_source=WEIGHT_SOURCE,
        ))

    return AnalyzeResponse(
        prediction_id=prediction_id,
        model_version=mv,
        domain=domain,
        inference_ms=dur,
        image={"width": int(orig_w), "height": int(orig_h)},
        objects=objects,
        summary=AnalyzeSummary(
            item_count=len(objects),
            trash_detected=len(objects) > 0,
            materials=materials,
            total_weight_estimate_g=round(total_w, 1) if objects else 0.0,
            weight_source=WEIGHT_SOURCE if objects else None,
        ),
    )


@app.post("/v1/analyze", response_model=AnalyzeResponse, dependencies=[Depends(require_api_key)])
def analyze(req: AnalyzeRequest):
    """
    Standalone Vision API: full scene analysis from an image URL.
    Superset of /predict: material buckets, per-object weight estimates,
    materials summary. Signals only — never computes TC/rewards.
    """
    if req.domain != "trash":
        raise HTTPException(400, f"Unknown domain '{req.domain}'; supported: trash")
    img_bytes = fetch_image_bytes(req.image_url)
    preds, raw, dur = run_inference(img_bytes)
    mv = model_version()
    prediction_id = str(uuid.uuid4())
    # Preview frames (log=false) are NOT persisted — they must never flood the
    # training flywheel; only the final snapped photo is logged.
    if req.log:
        log_prediction(prediction_id, req.image_url, req.user_id, mv, preds,
                       endpoint="v1/analyze", source=req.source)
    return build_analyze_response(preds, raw, dur, req.domain, prediction_id, mv)


@app.post("/v1/analyze/upload", response_model=AnalyzeResponse, dependencies=[Depends(require_api_key)])
async def analyze_upload(
    file: UploadFile = File(...),
    user_id: Optional[str] = Form(None),
    domain: str = Form("trash"),
    log: bool = Form(True),
    source: Optional[str] = Form(None),
):
    """Same as /v1/analyze but with a direct multipart image upload (no URL needed).

    Set log=false for live-camera PREVIEW frames: they are analyzed but NOT
    persisted, so the ~1/s preview stream never floods the training flywheel.
    Log only the final snapped photo (log=true, the default).

    `source` tags the context (e.g. 'product-scan'): logged rows with a
    non-litter source are kept in a SEPARATE training pool by the flywheel
    sync — supermarket shelf photos never mix into litter-detection training."""
    if domain != "trash":
        raise HTTPException(400, f"Unknown domain '{domain}'; supported: trash")
    img_bytes = await file.read()
    if not img_bytes:
        raise HTTPException(400, "empty upload")
    if len(img_bytes) > 10 * 1024 * 1024:
        raise HTTPException(400, "Image too large (> 10MB)")
    preds, raw, dur = run_inference(img_bytes)
    mv = model_version()
    prediction_id = str(uuid.uuid4())
    if log:
        log_prediction(prediction_id, f"upload:{file.filename or 'unnamed'}", user_id, mv, preds,
                       endpoint="v1/analyze/upload", source=source)
    return build_analyze_response(preds, raw, dur, domain, prediction_id, mv)


# Debug endpoints — only mounted when ALAMI_DEBUG=1 (never in production)
if DEBUG:
    @app.post("/debug-predict")
    def debug_predict(req: dict):  # raw dict, no validation
        print("🔍 RAW REQUEST BODY:", req)
        return {"received": req}

    @app.get("/debug-fetch")
    async def debug_fetch(url: str = "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/bus.jpg"):
        try:
            import requests
            print(f"🔍 Testing URL: {url}")
            response = requests.get(url, timeout=10, verify=not INSECURE_FETCH)
            return {
                "status_code": response.status_code,
                "content_type": response.headers.get("content-type"),
                "content_length": len(response.content),
                "success": True
            }
        except Exception as e:
            print(f"🔍 ERROR: {e}")
            return {"error": str(e), "success": False}

@app.post("/feedback", response_model=FeedbackResponse, dependencies=[Depends(require_api_key)])
def feedback(req: FeedbackRequest):
    # Minimal validation
    if not req.prediction_id:
        raise HTTPException(400, "prediction_id required")

    if not qc_weight_ok(req.corrected_weight_kg):
        raise HTTPException(400, "corrected_weight_kg out of bounds [0.001, 20.0]")

    # Quick QC: if no correction of any kind is provided → no-op. added_items
    # (the AI missed everything) and reasons (e.g. "unsure_skip") are meaningful
    # on their own, so they count as changes even without a relabel/weight.
    if (req.corrected_type is None and req.corrected_weight_kg is None
            and not req.corrected_items and not req.added_items and not req.reasons):
        return FeedbackResponse(ok=True, message="No changes supplied; feedback ignored.")

    corrected_items = [c.model_dump() for c in req.corrected_items] if req.corrected_items else None
    added_items = [a.model_dump() for a in req.added_items] if req.added_items else None
    reasons = list(req.reasons) if req.reasons else None

    # Persist to JSONL
    entry = {
        "prediction_id": req.prediction_id,
        "ts": now_iso(),
        "corrected_type": req.corrected_type,
        "corrected_weight_kg": req.corrected_weight_kg,
        "notes": req.notes,
        "source": req.source,
        "corrected_items": corrected_items,
        "added_items": added_items,
        "reasons": reasons,
    }
    append_jsonl(FB_LOG, entry)

    # Supabase update (best-effort)
    sb_update_feedback(
        prediction_id=req.prediction_id,
        corrected_type=req.corrected_type,
        corrected_weight_kg=req.corrected_weight_kg,
        corrected_items=corrected_items,
        source=req.source,
        notes=req.notes,
        added_items=added_items,
        reasons=reasons,
    )

    return FeedbackResponse(ok=True, message="Feedback recorded.")