#!/usr/bin/env python3
"""aria-brain-embed — One-shot + incremental Embedding für brain_chunks (KAR-120).

Schreibt Brain-Files in chunks → embedded sie via OpenAI text-embedding-3-large
(3072 dim) → upsertet sie in `brain_chunks` (Supabase aria-DB pgvector).

Embedding-Modell-Wahl (Sektion 12.3 der Gibson-Analyse):
- text-embedding-3-large: 3072 dim, ~$0.13 per 1M Token, general-purpose
- One-shot full-brain (200 Files, 30 MB): geschätzt ~$0.13

Voraussetzungen:
- OPENAI_API_KEY in env oder /root/.aria-secrets/openai.env
- pgvector + brain_chunks Table vorhanden (siehe /tmp/aria-pgvector-schema.sql, KAR-120-Setup)

CLI:
    python3 aria-brain-embed.py status      # zeige welche Files embedded sind
    python3 aria-brain-embed.py backfill    # full re-embed alles
    python3 aria-brain-embed.py incremental # nur neu/geändert (default)
    python3 aria-brain-embed.py file <path> # einzelnes File
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import time
import urllib.request
from pathlib import Path

# --- Config -----------------------------------------------------------------
BRAIN_ROOT = Path("/root/aria/brain")
SUPABASE_URL = os.environ.get("SUPABASE_URL", "").rstrip("/")
SUPABASE_SERVICE_ROLE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_EMBED_URL = "https://api.openai.com/v1/embeddings"
EMBED_MODEL = "text-embedding-3-large"
CHUNK_TARGET_CHARS = 2000  # ca. 500-700 Token
CHUNK_OVERLAP = 200

if not OPENAI_API_KEY:
    for env_file in ("/root/.aria-secrets/openai.env", "/root/aria/.env"):
        if os.path.exists(env_file):
            for line in Path(env_file).read_text().splitlines():
                if line.startswith("OPENAI_API_KEY="):
                    OPENAI_API_KEY = line.split("=", 1)[1].strip().strip('"').strip("'")
                    break
        if OPENAI_API_KEY:
            break

REST_BRAIN_CHUNKS = f"{SUPABASE_URL}/rest/v1/brain_chunks" if SUPABASE_URL else ""


# --- Helpers ----------------------------------------------------------------

def _supabase_headers() -> dict:
    return {
        "apikey": SUPABASE_SERVICE_ROLE_KEY,
        "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
        "Content-Type": "application/json",
        "Prefer": "resolution=merge-duplicates",
    }


def chunk_text(text: str, target: int = CHUNK_TARGET_CHARS, overlap: int = CHUNK_OVERLAP) -> list[str]:
    """Simple character-window chunking. Brain-Files sind meist <30KB,
    1-15 chunks pro Datei."""
    text = text.strip()
    if len(text) <= target:
        return [text] if text else []
    chunks = []
    pos = 0
    while pos < len(text):
        end = min(pos + target, len(text))
        # Versuche an Newline zu brechen für saubere Chunks
        if end < len(text):
            nl = text.rfind("\n", pos, end)
            if nl > pos + target // 2:
                end = nl
        chunks.append(text[pos:end].strip())
        if end >= len(text):
            break
        pos = max(end - overlap, pos + 1)
    return [c for c in chunks if c]


def embed_batch(texts: list[str], retries: int = 5) -> list[list[float]]:
    """OpenAI text-embedding-3-large, batch. Exponential-backoff bei 429."""
    if not OPENAI_API_KEY:
        raise SystemExit(
            "ERROR: OPENAI_API_KEY missing. Place in /root/.aria-secrets/openai.env "
            "as OPENAI_API_KEY=sk-... (chmod 600)."
        )
    payload = json.dumps({"model": EMBED_MODEL, "input": texts}).encode()
    delay = 1.0
    last_err = None
    for attempt in range(retries + 1):
        try:
            req = urllib.request.Request(
                OPENAI_EMBED_URL, data=payload,
                headers={
                    "Authorization": f"Bearer {OPENAI_API_KEY}",
                    "Content-Type": "application/json",
                }, method="POST",
            )
            with urllib.request.urlopen(req, timeout=60) as r:
                data = json.loads(r.read())
            return [item["embedding"] for item in data["data"]]
        except urllib.error.HTTPError as exc:
            last_err = exc
            if exc.code == 429 and attempt < retries:
                sys.stderr.write(f"  [429] retry {attempt + 1}/{retries} in {delay:.1f}s\n")
                time.sleep(delay)
                delay = min(delay * 2, 30)
                continue
            raise
    raise last_err or RuntimeError("embed_batch exhausted retries")


def upsert_chunk(file_path: str, chunk_idx: int, content: str, embedding: list[float]) -> bool:
    """Upsert via Supabase REST. Unique(file_path, chunk_idx) — resolution=merge."""
    if not REST_BRAIN_CHUNKS:
        raise SystemExit("ERROR: SUPABASE_URL missing")
    row = {
        "file_path": file_path,
        "chunk_idx": chunk_idx,
        "content": content,
        "embedding": embedding,
    }
    req = urllib.request.Request(
        REST_BRAIN_CHUNKS, data=json.dumps(row).encode(),
        headers={**_supabase_headers(), "Prefer": "resolution=merge-duplicates,return=minimal"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as r:
        return 200 <= r.status < 300


def brain_files(only_under: Path = BRAIN_ROOT) -> list[Path]:
    """Liste alle .md-Files unter BRAIN_ROOT, ohne .git/.archive."""
    out = []
    for md in only_under.rglob("*.md"):
        if any(part.startswith(".") or part in ("__pycache__", "Archiv") for part in md.relative_to(BRAIN_ROOT).parts):
            continue
        out.append(md)
    return sorted(out)


def db_files_status() -> dict[str, dict]:
    """Was ist in brain_chunks_files (welche Files schon embedded)."""
    if not REST_BRAIN_CHUNKS:
        return {}
    url = f"{SUPABASE_URL}/rest/v1/brain_chunks_files?select=*&limit=10000"
    req = urllib.request.Request(url, headers=_supabase_headers(), method="GET")
    try:
        with urllib.request.urlopen(req, timeout=15) as r:
            rows = json.loads(r.read())
        return {row["file_path"]: row for row in rows}
    except Exception as exc:  # noqa: BLE001
        sys.stderr.write(f"[embed] status query failed: {exc}\n")
        return {}


# --- Commands ---------------------------------------------------------------

def cmd_status(args) -> int:
    files = brain_files()
    db_idx = db_files_status()
    print(f"Brain-Files: {len(files)} total")
    print(f"DB-Indexed:  {len(db_idx)} files")
    print(f"Missing:     {len(files) - len(db_idx)} files not embedded yet")
    if args.verbose:
        for f in files:
            rel = str(f.relative_to(BRAIN_ROOT))
            status = "✓" if rel in db_idx else "—"
            print(f"  {status} {rel}")
    return 0


def cmd_file(args) -> int:
    f = Path(args.path)
    if not f.is_absolute():
        f = BRAIN_ROOT / f
    if not f.exists():
        print(f"ERROR: {f} not found", file=sys.stderr)
        return 1
    rel = str(f.relative_to(BRAIN_ROOT))
    content = f.read_text(errors="replace")
    chunks = chunk_text(content)
    if not chunks:
        print(f"{rel}: empty, skipped")
        return 0
    print(f"{rel}: {len(chunks)} chunks", file=sys.stderr)
    embeddings = embed_batch(chunks)
    for i, (text, emb) in enumerate(zip(chunks, embeddings)):
        ok = upsert_chunk(rel, i, text, emb)
        print(f"  chunk {i}: {'OK' if ok else 'FAIL'}", file=sys.stderr)
    return 0


def cmd_backfill(args) -> int:
    files = brain_files()
    print(f"Embedding {len(files)} brain-files via {EMBED_MODEL}...", file=sys.stderr)
    total_chunks = 0
    total_cost_est = 0.0
    for f in files:
        rel = str(f.relative_to(BRAIN_ROOT))
        content = f.read_text(errors="replace")
        chunks = chunk_text(content)
        if not chunks:
            continue
        try:
            embeddings = embed_batch(chunks)
        except Exception as exc:
            print(f"  {rel} EMBED-FAIL: {exc}", file=sys.stderr)
            continue
        for i, (text, emb) in enumerate(zip(chunks, embeddings)):
            try:
                upsert_chunk(rel, i, text, emb)
            except Exception as exc:
                print(f"  {rel}#{i} UPSERT-FAIL: {exc}", file=sys.stderr)
        total_chunks += len(chunks)
        # Grobe Cost-Schätzung: text-embedding-3-large = $0.13/M tokens, ~ 4 chars/token
        approx_tokens = sum(len(c) for c in chunks) / 4
        total_cost_est += approx_tokens * 0.13 / 1_000_000
        print(f"  {rel}: {len(chunks)} chunks", file=sys.stderr)
        time.sleep(0.1)
    print(f"\nDone. {total_chunks} chunks embedded across {len(files)} files. "
          f"Estimated cost: ${total_cost_est:.4f}", file=sys.stderr)
    return 0


def cmd_incremental(args) -> int:
    files = brain_files()
    db_idx = db_files_status()
    pending = [f for f in files if str(f.relative_to(BRAIN_ROOT)) not in db_idx]
    print(f"Incremental: {len(pending)} new/missing files (of {len(files)} total)", file=sys.stderr)
    if not pending:
        return 0
    for f in pending:
        rel = str(f.relative_to(BRAIN_ROOT))
        content = f.read_text(errors="replace")
        chunks = chunk_text(content)
        if not chunks:
            continue
        try:
            embeddings = embed_batch(chunks)
        except Exception as exc:
            print(f"  {rel} EMBED-FAIL: {exc}", file=sys.stderr)
            continue
        for i, (text, emb) in enumerate(zip(chunks, embeddings)):
            upsert_chunk(rel, i, text, emb)
        print(f"  {rel}: {len(chunks)} chunks", file=sys.stderr)
        time.sleep(0.1)
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Aria Brain Embedding (KAR-120)")
    sub = parser.add_subparsers(dest="cmd", required=True)

    p_st = sub.add_parser("status", help="show indexed vs missing files")
    p_st.add_argument("-v", "--verbose", action="store_true")
    p_st.set_defaults(func=cmd_status)

    p_bf = sub.add_parser("backfill", help="full re-embed all brain files")
    p_bf.set_defaults(func=cmd_backfill)

    p_inc = sub.add_parser("incremental", help="embed only new/changed files")
    p_inc.set_defaults(func=cmd_incremental)

    p_f = sub.add_parser("file", help="embed a single file")
    p_f.add_argument("path", help="absolute or BRAIN_ROOT-relative")
    p_f.set_defaults(func=cmd_file)

    args = parser.parse_args()
    return args.func(args)


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