#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
sync_supabase_feedback.py — pull predictions + human corrections from Supabase
into the JSONL format that ingest_feedback.py consumes.

Why: the HF Space's local feedback_logs/*.jsonl are ephemeral (lost on every
restart). Supabase (`trash_predictions`) is the durable store — this script
closes the gap between "app users correct labels" and "training dataset".

Reads   : Supabase table trash_predictions (service role)
Writes  : <out>/predictions.jsonl  — {prediction_id, ts, image_url, user_id,
                                      model_version, predictions:[{xyxy,cls,conf,label,raw_label}]}
          <out>/feedback.jsonl     — {prediction_id, ts, corrected_type,
                                      corrected_weight_kg, notes, source, corrected_items}

Rows without a per-object `predictions` JSONB array cannot yield YOLO boxes;
they are counted and skipped for predictions.jsonl (apply the migration in
docs/FLYWHEEL.md so the server stores full boxes).

Usage:
  SUPABASE_URL=... SUPABASE_SERVICE_ROLE_KEY=... \
  python ml/scripts/sync_supabase_feedback.py --out feedback_logs_sync
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger("sync_supabase_feedback")

PAGE_SIZE = 1000

# Rows whose `source` is in this set feed LITTER-detection training. Anything
# else (e.g. "product-scan" = supermarket shelf photos) goes to a separate
# pool file — kept for other purposes (brand radar!) but NEVER mixed into the
# litter training set, where it would hurt detection quality.
LITTER_SOURCES = {None, "", "alami-mobile"}

# Known test/dummy rows (e2e wiring checks) that must never become training
# data. One id or unique id-prefix per line, '#' comments.
DEFAULT_EXCLUDE_FILE = Path(__file__).resolve().parents[1] / "flywheel" / "exclude_predictions.txt"


def load_exclusions(path: Optional[Path]) -> Tuple[str, ...]:
    """Read id/prefix exclusion list; missing file -> empty (never fails)."""
    if path is None or not path.is_file():
        return ()
    prefixes = []
    for line in path.read_text(encoding="utf-8").splitlines():
        entry = line.split("#", 1)[0].strip()
        if entry:
            prefixes.append(entry)
    return tuple(prefixes)


def row_to_prediction_entry(row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """Transform one trash_predictions row into a predictions.jsonl entry.

    Returns None when the row carries no usable per-object boxes.
    """
    preds = row.get("predictions")
    if isinstance(preds, str):
        try:
            preds = json.loads(preds)
        except Exception:
            preds = None
    if not isinstance(preds, list) or not preds:
        return None
    boxes = []
    for b in preds:
        if not isinstance(b, dict):
            continue
        xyxy = b.get("xyxy")
        if not isinstance(xyxy, list) or len(xyxy) < 4:
            continue
        boxes.append({
            "xyxy": [float(x) for x in xyxy[:4]],
            "cls": int(b.get("cls", 0)),
            "conf": float(b.get("conf", 0.0)),
            "label": str(b.get("label", "")),
            "raw_label": str(b.get("raw_label", b.get("label", ""))),
        })
    if not boxes:
        return None
    if not row.get("image_url") or not row.get("prediction_id"):
        return None
    return {
        "prediction_id": str(row["prediction_id"]),
        "ts": row.get("created_at"),
        "image_url": row["image_url"],
        "user_id": row.get("user_id"),
        "model_version": row.get("model_version"),
        "predictions": boxes,
    }


def row_to_feedback_entry(row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """Transform one corrected trash_predictions row into a feedback.jsonl entry."""
    if not row.get("corrected_at") or not row.get("prediction_id"):
        return None
    if row.get("corrected_type") is None and row.get("corrected_weight_kg") is None \
            and not row.get("corrected_items") and not row.get("added_items"):
        return None

    def _maybe_json(v: Any) -> Any:
        """jsonb columns arrive as objects; some drivers hand back a JSON string."""
        if isinstance(v, str):
            try:
                return json.loads(v)
            except Exception:
                return None
        return v

    return {
        "prediction_id": str(row["prediction_id"]),
        "ts": row.get("corrected_at"),
        "corrected_type": row.get("corrected_type"),
        "corrected_weight_kg": row.get("corrected_weight_kg"),
        "notes": row.get("notes"),
        "source": row.get("feedback_source"),
        "corrected_items": _maybe_json(row.get("corrected_items")),
        # v2 (#141): recall signal (objects the AI missed) + failure chips
        "added_items": _maybe_json(row.get("added_items")),
        "reasons": _maybe_json(row.get("feedback_reasons")),
    }


def fetch_all_rows(url: str, key: str, table: str) -> List[Dict[str, Any]]:
    from supabase import create_client
    sb = create_client(url, key)
    rows: List[Dict[str, Any]] = []
    offset = 0
    while True:
        resp = (sb.table(table)
                .select("*")
                .order("created_at", desc=False)
                .range(offset, offset + PAGE_SIZE - 1)
                .execute())
        page = resp.data or []
        rows.extend(page)
        if len(page) < PAGE_SIZE:
            break
        offset += PAGE_SIZE
    return rows


def write_jsonl(path: Path, entries: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for e in entries:
            f.write(json.dumps(e, ensure_ascii=False) + "\n")


def transform_rows(rows: List[Dict[str, Any]],
                   exclude: Tuple[str, ...] = ()) -> Tuple[List[Dict], List[Dict], List[Dict], Dict[str, Any]]:
    """Pure transform: rows -> (litter predictions, other-source predictions,
    feedback entries, stats). Pool separation happens HERE: non-litter sources
    (product-scan etc.) never reach the litter training set. Rows whose
    prediction_id matches an `exclude` prefix (known test/dummy rows) are
    dropped from EVERY output."""
    pred_entries: List[Dict] = []
    other_entries: List[Dict] = []
    fb_entries: List[Dict] = []
    no_boxes = 0
    excluded = 0
    other_by_source: Dict[str, int] = {}
    for row in rows:
        pid = str(row.get("prediction_id") or "")
        if pid and any(pid.startswith(x) for x in exclude):
            excluded += 1
            continue
        p = row_to_prediction_entry(row)
        src = row.get("source")
        if p is not None:
            if src in LITTER_SOURCES:
                pred_entries.append(p)
            else:
                p["source"] = src
                other_entries.append(p)
                other_by_source[str(src)] = other_by_source.get(str(src), 0) + 1
        elif row.get("prediction_id") and row.get("image_url"):
            no_boxes += 1
        fb = row_to_feedback_entry(row)
        if fb is not None and src in LITTER_SOURCES:
            fb_entries.append(fb)
    stats = {
        "rows_total": len(rows),
        "rows_excluded_testdata": excluded,
        "predictions_with_boxes": len(pred_entries),
        "predictions_other_sources": len(other_entries),
        "other_sources_breakdown": other_by_source,
        "predictions_without_boxes": no_boxes,
        "feedback_entries": len(fb_entries),
    }
    return pred_entries, other_entries, fb_entries, stats


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Sync trash_predictions from Supabase to ingest-ready JSONL.")
    ap.add_argument("--out", default="feedback_logs_sync", help="output directory")
    ap.add_argument("--table", default=os.environ.get("SB_TABLE_PREDICTIONS", "trash_predictions"))
    ap.add_argument("--exclude-file", default=str(DEFAULT_EXCLUDE_FILE),
                    help="id/prefix blocklist for known test rows (default: ml/flywheel/exclude_predictions.txt)")
    args = ap.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    url = os.environ.get("SUPABASE_URL")
    key = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
    if not url or not key:
        logger.error("SUPABASE_URL / SUPABASE_SERVICE_ROLE_KEY not set — nothing to sync.")
        return 2

    exclude = load_exclusions(Path(args.exclude_file))
    rows = fetch_all_rows(url, key, args.table)
    pred_entries, other_entries, fb_entries, stats = transform_rows(rows, exclude=exclude)
    if stats["rows_excluded_testdata"]:
        logger.info("Excluded %d known test/dummy row(s) via %s.",
                    stats["rows_excluded_testdata"], args.exclude_file)

    out = Path(args.out)
    write_jsonl(out / "predictions.jsonl", pred_entries)
    write_jsonl(out / "predictions_other_sources.jsonl", other_entries)
    write_jsonl(out / "feedback.jsonl", fb_entries)
    (out / "sync_stats.json").write_text(json.dumps(stats, indent=2), encoding="utf-8")

    logger.info("Synced %d rows: %d litter predictions, %d other-source (separate pool: %s), "
                "%d without boxes (need migration), %d feedback entries.",
                stats["rows_total"], stats["predictions_with_boxes"],
                stats["predictions_other_sources"], stats["other_sources_breakdown"],
                stats["predictions_without_boxes"], stats["feedback_entries"])
    if stats["predictions_without_boxes"] > 0 and stats["predictions_with_boxes"] == 0:
        logger.warning("No rows carry per-object boxes — apply the 'predictions jsonb' migration "
                       "(docs/FLYWHEEL.md) so new predictions become trainable.")
    return 0


if __name__ == "__main__":
    sys.exit(main())