Blame for db.py

"""Supabase-backed data layer for Transcriber.

All persistence goes through this module so main.py never touches the DB driver
directly. The backend authenticates with the Supabase SERVICE ROLE key (bypasses
RLS) and must only ever run on a trusted machine.
"""
import json
import os
import uuid
from datetime import datetime, timezone
from typing import Optional

from supabase import create_client, Client

TABLE = "saved_reels"
FAILURE_STATUSES = frozenset({"failed", "cookies_expired", "unsupported_format"})
_client: Optional[Client] = None


def get_client() -> Client:
    """Lazily create the Supabase client so importing db.py never needs creds."""
    global _client
    if _client is None:
        url = os.getenv("SUPABASE_URL")
        key = os.getenv("SUPABASE_SERVICE_KEY")
        if not url or not key:
            raise RuntimeError("SUPABASE_URL and SUPABASE_SERVICE_KEY must be set.")
        _client = create_client(url, key)
    return _client


def _chunked(seq: list, size: int):
    """Yield successive size-length slices of a sequence."""
    for i in range(0, len(seq), size):
        yield seq[i:i + size]


def _select_newest_first(columns: str, limit: Optional[int] = None, refine=None) -> list:
    """Select rows newest-first, paging past PostgREST's per-request row cap.

    refine, if given, adds filters to each page's query."""
    rows = []
    while limit is None or len(rows) < limit:
        last = len(rows) + 999 if limit is None else limit - 1
        query = get_client().table(TABLE).select(columns)
        if refine:
            query = refine(query)
        page = (query.order("created_at", desc=True).order("id")
                .range(len(rows), last).execute().data or [])
        if not page:
            break
        rows.extend(page)
    return rows


def row_to_record(r: dict) -> dict:
    """Map a raw DB row to the API record shape the frontend expects."""
    ej = r.get("extracted_json")
    if isinstance(ej, str):
        try:
            ej = json.loads(ej)
        except (json.JSONDecodeError, TypeError):
            ej = {}
    return {
        "id": r["id"],
        "url": r.get("url"),
        "title": r.get("title"),
        "raw_transcript": r.get("raw_transcript"),
        "post_caption": r.get("post_caption"),
        "extracted_json": ej,
        "created_at": r.get("created_at"),
        "cluster": r.get("cluster") or "Unclustered",
        "status": r.get("status") or "done",
        "error": r.get("error"),
    }


def get_reel_by_url(url: str) -> Optional[dict]:
    res = get_client().table(TABLE).select("*").eq("url", url).limit(1).execute()
    rows = res.data or []
    return row_to_record(rows[0]) if rows else None


def insert_reel(url, title, raw_transcript, post_caption, extracted_json,
                status="done", source=None, embedding=None, cluster=None) -> dict:
    row = {
        "id": str(uuid.uuid4()),
        "url": url,
        "title": title,
        "raw_transcript": raw_transcript,
        "post_caption": post_caption,
        "extracted_json": extracted_json,  # dict -> jsonb
        "created_at": datetime.now(timezone.utc).isoformat(),
        "status": status,
        "source": source,
    }
    if embedding is not None:
        row["embedding"] = embedding  # list -> pgvector
    if cluster is not None:
        row["cluster"] = cluster
    get_client().table(TABLE).insert(row).execute()
    return row_to_record(row)


def statuses_for_urls(urls: list) -> list:
    """Return [{url, status, error}] for whichever of the given URLs exist."""
    rows = []
    for chunk in _chunked(urls, 100):
        res = get_client().table(TABLE).select("url", "status", "error").in_("url", chunk).execute()
        rows.extend(res.data or [])
    return rows


def insert_pending_reels(reels: list, source: str) -> None:
    """Enqueue parsed [{url, title, caption}] reels as pending rows for the worker."""
    rows = [
        {
            "id": str(uuid.uuid4()),
            "url": r["url"],
            "title": r["title"] or f"Reel ({r['url'].split('/reel/')[-1].replace('/', '')})",
            "raw_transcript": "",
            "post_caption": r["caption"],
            "extracted_json": {},
            "created_at": datetime.now(timezone.utc).isoformat(),
            "status": "pending",
            "source": source,
        }
        for r in reels
    ]
    for chunk in _chunked(rows, 50):
        get_client().table(TABLE).insert(chunk).execute()


def reels_for_clustering() -> list:
    """Every row as {id, extracted_json, embedding}, with embedding as a list of floats or None."""
    rows = []
    for r in _select_newest_first("id, extracted_json, embedding"):
        embedding = r.get("embedding")
        if isinstance(embedding, str):  # PostgREST returns pgvector as text: "[0.1,0.2,...]"
            embedding = json.loads(embedding)
        rows.append({"id": r["id"], "extracted_json": row_to_record(r)["extracted_json"], "embedding": embedding})
    return rows


def reels_missing_embedding() -> list:
    """Rows with no search embedding yet, as [{id, extracted_json}]."""
    return _select_newest_first("id, extracted_json", refine=lambda q: q.is_("embedding", "null"))


def set_embedding(reel_id: str, embedding: list) -> None:
    get_client().table(TABLE).update({"embedding": embedding}).eq("id", reel_id).execute()


def nearest_cluster(embedding: list, min_similarity: float) -> Optional[str]:
    """Cluster of the most similar already-clustered reel at or above min_similarity, else None."""
    c = get_client()
    matches = c.rpc("match_reels", {"query_embedding": embedding, "match_threshold": min_similarity,
                                    "match_count": 10}).execute().data or []
    ids = [m["id"] for m in matches]  # closest first
    if not ids:
        return None
    rows = c.table(TABLE).select("id, cluster").in_("id", ids).execute().data or []
    clusters = {r["id"]: r["cluster"] for r in rows if r.get("cluster")}
    return next((clusters[i] for i in ids if i in clusters), None)


def set_cluster(reel_ids: list, cluster: Optional[str]) -> None:
    for chunk in _chunked(reel_ids, 100):
        get_client().table(TABLE).update({"cluster": cluster}).in_("id", chunk).execute()


def delete_reel(reel_id: str) -> bool:
    """Delete a reel by id. Returns True if a row was removed, False if none matched."""
    res = get_client().table(TABLE).delete().eq("id", reel_id).execute()
    return bool(res.data)


def claim_next_pending() -> Optional[dict]:
    """Atomically claim the oldest pending reel: flip it to 'processing' and
    return the raw row. Returns None if the queue is empty or another worker
    won the claim."""
    c = get_client()
    now = datetime.now(timezone.utc).isoformat()
    res = (c.table(TABLE).select("*")
           .eq("status", "pending")
           .or_(f"next_attempt_at.is.null,next_attempt_at.lte.{now}")
           .order("created_at").limit(1).execute())
    rows = res.data or []
    if not rows:
        return None
    row = rows[0]
    attempt_count = int(row.get("attempt_count") or 0) + 1
    payload = {
        "status": "processing",
        "processing_started_at": now,
        "next_attempt_at": None,
        "attempt_count": attempt_count,
    }
    upd = (c.table(TABLE).update(payload)
           .eq("id", row["id"]).eq("status", "pending").execute())
    if not upd.data:
        return None  # lost the race
    return {**row, **payload}


def update_reel_result(reel_id: str, title: str, raw_transcript, post_caption,
                       extracted_json, status: str = "done", embedding=None, cluster=None) -> None:
    """Write pipeline results back onto an existing (claimed) row."""
    payload = {
        "title": title,
        "raw_transcript": raw_transcript,
        "post_caption": post_caption,
        "extracted_json": extracted_json,
        "status": status,
        "error": None,
        "processing_started_at": None,
        "next_attempt_at": None,
    }
    if embedding is not None:
        payload["embedding"] = embedding
    if cluster is not None:
        payload["cluster"] = cluster
    get_client().table(TABLE).update(payload).eq("id", reel_id).execute()


def mark_failed(reel_id: str, error, status: str = "failed") -> None:
    """Mark a claimed row with a failure status, recording a truncated error message."""
    get_client().table(TABLE).update({
        "status": status,
        "error": str(error)[:500],
        "processing_started_at": None,
        "next_attempt_at": None,
    }).eq("id", reel_id).execute()


def retry_reel(reel_id: str) -> bool:
    """Requeue a terminal failure and reset its retry budget."""
    res = get_client().table(TABLE).update({
        "status": "pending",
        "error": None,
        "processing_started_at": None,
        "next_attempt_at": None,
        "attempt_count": 0,
    }).eq("id", reel_id).in_("status", list(FAILURE_STATUSES)).execute()
    return bool(res.data)

def recover_stale_processing(stale_before: datetime) -> int:
    """Return abandoned claims older than stale_before to the queue."""
    res = (get_client().table(TABLE).update({
        "status": "pending",
        "error": "Recovered stale processing claim",
        "processing_started_at": None,
    }).eq("status", "processing")
      .lt("processing_started_at", stale_before.astimezone(timezone.utc).isoformat())
      .execute())
    return len(res.data or [])


def schedule_retry(reel_id: str, error: str, next_attempt_at: datetime) -> None:
    """Release a claim for a future retry while preserving its attempt count."""
    get_client().table(TABLE).update({
        "status": "pending",
        "error": str(error)[:500],
        "processing_started_at": None,
        "next_attempt_at": next_attempt_at.astimezone(timezone.utc).isoformat(),
    }).eq("id", reel_id).execute()