#!/usr/bin/env python3
"""
Aria Brain Search

Hybrid search über das Aria Obsidian Vault. Kombiniert:
- BM25 Keyword Ranking (pure Python, keine Dependencies)
- Recency Boost (neuere Einträge gewinnen bei gleicher Relevanz)
- Frontmatter Boost (Type-Match in Frontmatter zählt extra)
- Path Match (Dateiname/Pfad Treffer zählen extra)

Output: Top-N Treffer mit Pfad, Score, und kurzem Kontext.

Usage:
  aria-brain-search "Müller GmbH"
  aria-brain-search "DSGVO" --top 10
  aria-brain-search "kadi-v2" --type fact

CLI ist direkt nutzbar, aber primär als Tool für Aria gedacht:
Aria ruft via Bash auf und bekommt strukturierten Output.
"""

import os
import re
import sys
import json
import math
import argparse
from pathlib import Path
from collections import Counter, defaultdict
from datetime import datetime

BRAIN_DIR = Path(os.path.expandvars("$HOME/aria/brain"))
INDEX_FILE = Path("/root/.aria-brain-index.json")

# German stopwords (kurz, manuell)
STOPWORDS = {
    "der", "die", "das", "ein", "eine", "und", "oder", "aber", "wenn",
    "ist", "sind", "war", "waren", "hat", "haben", "wird", "werden",
    "auf", "in", "zu", "von", "mit", "bei", "für", "über", "unter",
    "den", "dem", "des", "im", "am", "zur", "zum", "vom", "beim",
    "ich", "du", "er", "sie", "es", "wir", "ihr", "mich", "dich",
    "nicht", "kein", "keine", "nur", "auch", "noch", "schon", "mal",
    "dass", "weil", "wie", "was", "wann", "wo", "wer", "warum",
    "the", "is", "and", "or", "of", "to", "a", "an", "for", "in",
}


def tokenize(text: str) -> list[str]:
    """Lowercase, alphanumerisch, ohne Stopwords."""
    text = text.lower()
    # Replace hyphens and underscores with space (für "Aria-Tester" -> "aria tester")
    text = re.sub(r'[-_]', ' ', text)
    tokens = re.findall(r'[a-zäöüß0-9]+', text)
    return [t for t in tokens if t not in STOPWORDS and len(t) > 1]


def parse_frontmatter(text: str) -> tuple[dict, str]:
    """Trennt YAML Frontmatter vom Body."""
    if not text.startswith("---"):
        return {}, text
    parts = text.split("---", 2)
    if len(parts) < 3:
        return {}, text
    fm_text = parts[1]
    body = parts[2]
    fm = {}
    for line in fm_text.strip().split("\n"):
        if ":" in line:
            k, v = line.split(":", 1)
            fm[k.strip()] = v.strip()
    return fm, body


def index_brain() -> dict:
    """Liest alle Markdown-Files und baut einen BM25 Index."""
    docs = []
    df = Counter()  # Document frequency per term
    n_docs = 0

    for md_file in BRAIN_DIR.rglob("*.md"):
        # Skip archive
        if "06-Daily/weekly" in str(md_file) or "06-Daily/monthly" in str(md_file):
            pass  # we DO want to index these
        if "/.git/" in str(md_file):
            continue

        try:
            text = md_file.read_text(encoding="utf-8")
        except (OSError, UnicodeDecodeError):
            continue

        fm, body = parse_frontmatter(text)
        tokens = tokenize(body + " " + " ".join(fm.values()))
        if not tokens:
            continue

        token_freq = Counter(tokens)
        rel_path = str(md_file.relative_to(BRAIN_DIR))

        # Mtime for recency
        try:
            mtime = md_file.stat().st_mtime
        except OSError:
            mtime = 0

        docs.append({
            "path": rel_path,
            "name": md_file.stem,
            "type": fm.get("type", ""),
            "description": fm.get("description", ""),
            "superseded_by": fm.get("superseded_by", ""),  # KAR-756 conflict-surfacing
            "status": fm.get("status", ""),
            "token_freq": dict(token_freq),
            "doc_len": len(tokens),
            "mtime": mtime,
            "preview": body.strip()[:200].replace("\n", " "),
        })

        for term in set(tokens):
            df[term] += 1
        n_docs += 1

    # Compute average doc length
    avg_dl = sum(d["doc_len"] for d in docs) / max(1, len(docs))

    index = {
        "docs": docs,
        "df": dict(df),
        "n_docs": n_docs,
        "avg_dl": avg_dl,
        "indexed_at": datetime.now().isoformat(),
    }
    return index


def save_index(index: dict):
    INDEX_FILE.write_text(json.dumps(index), encoding="utf-8")


def load_index() -> dict:
    # KAR-768: empty/corrupt Index (z.B. 0-Byte nach abgebrochenem Write) darf
    # nicht crashen — None zurueckgeben triggert needs_reindex -> Selbstheilung.
    if not INDEX_FILE.exists():
        return None
    try:
        text = INDEX_FILE.read_text(encoding="utf-8")
        if not text.strip():
            return None
        return json.loads(text)
    except (json.JSONDecodeError, OSError):
        return None


def needs_reindex(index: dict) -> bool:
    """Re-index wenn älter als 1h oder wenn Brain neuer ist."""
    if not index:
        return True
    indexed_at = datetime.fromisoformat(index["indexed_at"]).timestamp()
    if (datetime.now().timestamp() - indexed_at) > 3600:
        return True
    # Quick check: any md file newer than index?
    for md_file in BRAIN_DIR.rglob("*.md"):
        try:
            if md_file.stat().st_mtime > indexed_at:
                return True
        except OSError:
            pass
    return False


def bm25_score(query_tokens: list[str], doc: dict, index: dict, k1: float = 1.5, b: float = 0.75) -> float:
    """BM25 Okapi Score für ein Dokument."""
    score = 0.0
    n = index["n_docs"]
    avg_dl = index["avg_dl"]
    dl = doc["doc_len"]

    for term in query_tokens:
        if term not in doc["token_freq"]:
            continue
        f = doc["token_freq"][term]
        df = index["df"].get(term, 0)
        if df == 0:
            continue
        idf = math.log((n - df + 0.5) / (df + 0.5) + 1.0)
        tf = (f * (k1 + 1)) / (f + k1 * (1 - b + b * dl / avg_dl))
        score += idf * tf

    return score


def search(query: str, top: int = 5, type_filter: str = None) -> list[dict]:
    """Hauptfunktion: lädt/baut Index und ranked alle Docs."""
    index = load_index()
    if needs_reindex(index):
        index = index_brain()
        save_index(index)

    query_tokens = tokenize(query)
    if not query_tokens:
        return []

    results = []
    now = datetime.now().timestamp()

    for doc in index["docs"]:
        if type_filter and doc["type"] != type_filter:
            continue

        score = bm25_score(query_tokens, doc, index)
        if score == 0:
            continue

        # Filename match boost
        name_lower = doc["name"].lower()
        for tok in query_tokens:
            if tok in name_lower:
                score *= 1.5

        # Path match boost
        path_lower = doc["path"].lower()
        for tok in query_tokens:
            if tok in path_lower:
                score *= 1.2

        # Recency boost: max 30% boost for very recent docs (last 7 days)
        age_days = (now - doc["mtime"]) / 86400
        if age_days < 7:
            score *= 1.3
        elif age_days < 30:
            score *= 1.1

        results.append({
            "path": doc["path"],
            "score": round(score, 3),
            "type": doc["type"],
            "superseded_by": doc.get("superseded_by", ""),  # KAR-756
            "status": doc.get("status", ""),
            "preview": doc["preview"],
        })

    results.sort(key=lambda r: r["score"], reverse=True)
    return results[:top]


def vector_search(query: str, top: int = 5) -> list[dict]:
    """KAR-120: pgvector semantic search via OpenAI text-embedding-3-large.
    Liefert chunks mit cosine-distance. Aufrufer kombiniert mit BM25 via RRF."""
    openai_key = os.environ.get("OPENAI_API_KEY", "")
    supabase_url = os.environ.get("SUPABASE_URL", "").rstrip("/")
    supabase_key = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")
    if not openai_key:
        # Lade aus secrets
        sec = Path("/root/.aria-secrets/openai.env")
        if sec.exists():
            for line in sec.read_text().splitlines():
                if line.startswith("OPENAI_API_KEY="):
                    openai_key = line.split("=", 1)[1].strip().strip('"').strip("'")
                    break
    if not openai_key or not supabase_url or not supabase_key:
        return []
    import urllib.request
    # Embed query
    payload = json.dumps({"model": "text-embedding-3-large", "input": query}).encode()
    req = urllib.request.Request(
        "https://api.openai.com/v1/embeddings", data=payload,
        headers={"Authorization": f"Bearer {openai_key}", "Content-Type": "application/json"},
        method="POST")
    try:
        with urllib.request.urlopen(req, timeout=15) as r:
            embedding = json.loads(r.read())["data"][0]["embedding"]
    except Exception as exc:
        print(f"[vector] embed-fail: {exc}", file=sys.stderr)
        return []
    # Postgres-RPC for cosine-similarity. Use PostgREST RPC function or
    # direct REST with order=embedding for now (pgvector supports operator on REST)
    # Easiest: use ORDER BY embedding <=> query_embedding via SQL through a
    # database function. We'll create one on the fly via the REST `rpc`-Endpoint.
    body = json.dumps({"q_embedding": embedding, "match_limit": top}).encode()
    req = urllib.request.Request(
        f"{supabase_url}/rest/v1/rpc/brain_chunks_vector_search", data=body,
        headers={"apikey": supabase_key, "Authorization": f"Bearer {supabase_key}",
                 "Content-Type": "application/json"},
        method="POST")
    try:
        with urllib.request.urlopen(req, timeout=15) as r:
            return json.loads(r.read())
    except Exception as exc:
        print(f"[vector] rpc-fail: {exc}", file=sys.stderr)
        return []


def hybrid_search(query: str, top: int = 5, type_filter: str = None) -> list[dict]:
    """KAR-120: BM25 + pgvector via RRF (Reciprocal Rank Fusion).
    k=60 (Standard-Wert).
    """
    bm25 = search(query, top=top * 3, type_filter=type_filter)
    vec = vector_search(query, top=top * 3)
    K = 60.0
    fused = {}
    for rank, r in enumerate(bm25):
        path = r["path"]
        fused.setdefault(path, {"path": path, "bm25_rank": rank + 1, "vec_rank": None,
                                  "score": 0, "preview": r["preview"], "type": r.get("type")})
        fused[path]["score"] += 1 / (K + rank + 1)
    for rank, r in enumerate(vec):
        path = r.get("file_path", "")
        if path not in fused:
            fused[path] = {"path": path, "bm25_rank": None, "vec_rank": rank + 1,
                            "score": 0, "preview": (r.get("content") or "")[:200], "type": None}
        else:
            fused[path]["vec_rank"] = rank + 1
        fused[path]["score"] += 1 / (K + rank + 1)
    ranked = sorted(fused.values(), key=lambda x: x["score"], reverse=True)
    return ranked[:top]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("query", help="Search query")
    parser.add_argument("--top", type=int, default=5, help="Number of results")
    parser.add_argument("--type", default=None, help="Filter by frontmatter type")
    parser.add_argument("--json", action="store_true", help="JSON output")
    parser.add_argument("--reindex", action="store_true", help="Force reindex")
    parser.add_argument("--hybrid", action="store_true",
                        help="BM25 + pgvector RRF (KAR-120). Falls Embeddings nicht da: degraded zu BM25.")
    args = parser.parse_args()

    if args.reindex:
        index = index_brain()
        save_index(index)
        print(f"Indexed {index['n_docs']} documents", file=sys.stderr)

    if args.hybrid:
        results = hybrid_search(args.query, top=args.top, type_filter=args.type)
    else:
        results = search(args.query, top=args.top, type_filter=args.type)

    if args.json:
        print(json.dumps(results, indent=2, ensure_ascii=False))
    else:
        if not results:
            print("Keine Treffer.")
            return
        for i, r in enumerate(results, 1):
            score_val = r.get("score", 0)
            extra = ""
            if "bm25_rank" in r or "vec_rank" in r:
                extra = f" (bm25 #{r.get('bm25_rank') or '-'}, vec #{r.get('vec_rank') or '-'})"
            print(f"{i}. [{score_val:.4f}{extra}] brain/{r['path']}")
            if r.get("type"):
                print(f"   type: {r['type']}")
            # KAR-756 Conflict-Surfacing: veraltete/abgeloeste Notes nicht still nutzen
            if r.get("superseded_by"):
                print(f"   ⚠ SUPERSEDED → nutze stattdessen: {r['superseded_by']}")
            elif r.get("status") in ("archiv", "historical", "dormant"):
                print(f"   ⚠ status: {r['status']} (veraltet — vor Nutzung prüfen)")
            print(f"   {r['preview']}")
            print()


if __name__ == "__main__":
    main()