#!/usr/bin/env -S uv run --script --quiet
# /// script
# requires-python = ">=3.12"
# dependencies = ["httpx", "pydantic-settings"]
# ///
"""
Backfill embeddings for pub-search documents.

Usage:
    ./scripts/backfill-embeddings              # process all documents missing embeddings
    ./scripts/backfill-embeddings --limit 10   # process 10 documents
    ./scripts/backfill-embeddings --dry-run    # show what would be processed
"""

import argparse
import json
import os
import sys

import httpx
from pydantic_settings import BaseSettings, SettingsConfigDict


class Settings(BaseSettings):
    model_config = SettingsConfigDict(
        env_file=os.environ.get("ENV_FILE", ".env"), extra="ignore"
    )

    turso_url: str
    turso_token: str
    voyage_api_key: str

    @property
    def turso_host(self) -> str:
        """Strip libsql:// prefix if present."""
        url = self.turso_url
        if url.startswith("libsql://"):
            url = url[len("libsql://"):]
        return url


def turso_query(settings: Settings, sql: str, args: list | None = None) -> list[dict]:
    """Execute a query against Turso and return rows."""
    stmt = {"sql": sql}
    if args:
        stmt["args"] = [{"type": "text", "value": str(a)} for a in args]

    response = httpx.post(
        f"https://{settings.turso_host}/v2/pipeline",
        headers={
            "Authorization": f"Bearer {settings.turso_token}",
            "Content-Type": "application/json",
        },
        json={"requests": [{"type": "execute", "stmt": stmt}, {"type": "close"}]},
        timeout=30,
    )
    response.raise_for_status()
    data = response.json()

    result = data["results"][0]
    if result["type"] == "error":
        raise Exception(f"Turso error: {result['error']}")

    cols = [c["name"] for c in result["response"]["result"]["cols"]]
    rows = result["response"]["result"]["rows"]

    def extract_value(cell):
        if cell is None:
            return None
        if isinstance(cell, dict):
            return cell.get("value")
        # cell might be the value directly in some formats
        return cell

    return [dict(zip(cols, [extract_value(cell) for cell in row])) for row in rows]


def turso_exec(settings: Settings, sql: str, args: list | None = None, retries: int = 3) -> None:
    """Execute a statement against Turso with retry logic."""
    turso_batch_exec(settings, [(sql, args)], retries)


def turso_batch_exec(settings: Settings, statements: list[tuple[str, list | None]], retries: int = 3) -> None:
    """Execute multiple statements in a single pipeline request."""
    import time

    requests = []
    for sql, args in statements:
        stmt = {"sql": sql}
        if args:
            stmt["args"] = [{"type": "text", "value": str(a)} for a in args]
        requests.append({"type": "execute", "stmt": stmt})
    requests.append({"type": "close"})

    for attempt in range(retries):
        try:
            response = httpx.post(
                f"https://{settings.turso_host}/v2/pipeline",
                headers={
                    "Authorization": f"Bearer {settings.turso_token}",
                    "Content-Type": "application/json",
                },
                json={"requests": requests},
                timeout=120,
            )
            response.raise_for_status()
            data = response.json()
            for i, result in enumerate(data["results"][:-1]):  # skip the close result
                if result["type"] == "error":
                    raise Exception(f"Turso error on statement {i}: {result['error']}")
            return
        except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.ConnectError) as e:
            if attempt < retries - 1:
                wait = 2 ** (attempt + 1)
                print(f"  {type(e).__name__}, retrying in {wait}s...")
                time.sleep(wait)
            else:
                raise


def voyage_embed(settings: Settings, texts: list[str]) -> list[list[float]]:
    """Generate embeddings using Voyage AI."""
    response = httpx.post(
        "https://api.voyageai.com/v1/embeddings",
        headers={
            "Authorization": f"Bearer {settings.voyage_api_key}",
            "Content-Type": "application/json",
        },
        json={
            "input": texts,
            "model": "voyage-4-lite",
            "output_dimension": 1024,
            "input_type": "document",
        },
        timeout=60,
    )
    response.raise_for_status()
    data = response.json()
    return [item["embedding"] for item in data["data"]]


def main():
    parser = argparse.ArgumentParser(description="Backfill embeddings for pub-search")
    parser.add_argument("--limit", type=int, default=0, help="max documents to process (0 = all)")
    parser.add_argument("--batch-size", type=int, default=20, help="documents per Voyage API call")
    parser.add_argument("--dry-run", action="store_true", help="show what would be processed")
    args = parser.parse_args()

    try:
        settings = Settings()  # type: ignore
    except Exception as e:
        print(f"error loading settings: {e}", file=sys.stderr)
        print("required env vars: TURSO_URL, TURSO_TOKEN, VOYAGE_API_KEY", file=sys.stderr)
        sys.exit(1)

    # check if embedding column exists, add if not
    try:
        turso_query(settings, "SELECT embedding FROM documents LIMIT 1")
    except Exception as e:
        if "no such column" in str(e).lower():
            print("adding embedding column...")
            turso_exec(settings, "ALTER TABLE documents ADD COLUMN embedding F32_BLOB(1024)")
            print("done")
        else:
            raise

    # get documents needing embeddings
    limit_clause = f"LIMIT {args.limit}" if args.limit > 0 else ""
    docs = turso_query(
        settings,
        f"SELECT uri, title, content FROM documents WHERE embedding IS NULL {limit_clause}",
    )

    if not docs:
        print("no documents need embeddings")
        return

    print(f"found {len(docs)} documents needing embeddings")

    if args.dry_run:
        for doc in docs[:10]:
            print(f"  - {doc['uri']}: {doc['title'][:50]}...")
        if len(docs) > 10:
            print(f"  ... and {len(docs) - 10} more")
        return

    # process in batches with concurrency
    from concurrent.futures import ThreadPoolExecutor, as_completed

    def process_batch(batch_info):
        batch_num, batch = batch_info
        texts = [f"{doc['title']} {doc['content']}" for doc in batch]
        embeddings = voyage_embed(settings, texts)
        statements = []
        for doc, embedding in zip(batch, embeddings):
            embedding_json = json.dumps(embedding)
            statements.append((
                "UPDATE documents SET embedding = vector32(?) WHERE uri = ?",
                [embedding_json, doc["uri"]],
            ))
        turso_batch_exec(settings, statements)
        return batch_num, len(batch)

    batches = [(i // args.batch_size + 1, docs[i : i + args.batch_size])
               for i in range(0, len(docs), args.batch_size)]

    processed = 0
    workers = min(8, len(batches))  # more workers now that index is dropped
    print(f"processing {len(batches)} batches with {workers} workers...")

    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {executor.submit(process_batch, b): b[0] for b in batches}
        for future in as_completed(futures):
            batch_num, count = future.result()
            processed += count
            print(f"batch {batch_num} done ({processed}/{len(docs)})", flush=True)

    print(f"done! processed {processed} documents")


if __name__ == "__main__":
    main()
