#!/usr/bin/env python3
"""codebase-index — lokaler semantischer Index ueber einen Git-Code-Baum.

F1-Bau nach Konzept (handoffs/codebase-verstaendnis-lokal-konzept.md),
Auflagen Dieter 43102:
- Streaming in vorallokierte float32-Matrix (memmap), KEINE Vektoren-Liste
  (Run-3-Befund: Liste trieb RAM-Spitze auf 10,5 GB).
- Inkrementell: Content-Hash je Datei im Manifest, Delta statt Vollindex.
- Query mit Frische-Auskunft + Selbstauskunft in JEDER Antwort.
- Query-Fehlschlag-Log lokal (query-log.jsonl) — der Moat.

Aufruf:
  python codebase_index.py build  <repo>
  python codebase_index.py update <repo>
  python codebase_index.py query  <repo> "<frage>" [--top 5]
  python codebase_index.py sample <repo>

Index-Ort: ~/.blun/codebase-index/<workspace-hash>/ (sichtbar, mit Rueckweg).
"""
import hashlib
import json
import os
import subprocess
import sys
import time
from datetime import datetime, timezone

import numpy as np

CODE_EXT = (
    ".ts", ".tsx", ".js", ".mjs", ".cjs", ".py", ".html", ".css",
    ".json", ".md", ".toml", ".yml", ".yaml", ".sh", ".ps1", ".sql",
)
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 100
EMBED_DIM = 384
BATCH = 256
MODEL_NAME = "BAAI/bge-small-en-v1.5"
JINA_CODE_MODEL_NAME = "jinaai/jina-embeddings-v2-base-code"
KNOWN_MODELS = {
    MODEL_NAME: 384,
    JINA_CODE_MODEL_NAME: 768,
}
MODEL_ALIASES = {
    "bge": MODEL_NAME,
    "jina-code": JINA_CODE_MODEL_NAME,
}
AUTO_MODEL_ORDER = (JINA_CODE_MODEL_NAME, MODEL_NAME)

# Iteration 2 (Dieter 43135): Rausch-Filter — Backup-Kopien und Vendor-
# Buelle machten 66% aller Index-Rows aus und dominieren Top-5.
EXCLUDE_PREFIXES = ("backup-", "vendor/")

# Aktives Modell + Dim zur Laufzeit (wird per --model ueberschrieben).
_active_model_name = MODEL_NAME
_active_dim = EMBED_DIM

# Stichprobe V3 (Dieter 43141): Fragen PARAPHRASIERT aus gelesenem
# Dateiinhalt — keine wörtlichen Bezeichner/Funktionsnamen/Kommentare aus
# den Dateien (misst semantische Suche, nicht Wortgleichheit).
# api-chat.js: routet Chat-Nachrichten ans Modell, Fallback-Kette, Stream.
# terminal-manager.js: startet Shell-Prozesse mit Sandbox-/Approval-Modi.
# renderer-view-shell.js: wechselt Login-/Willkommens-/Chat-Ansicht.
SAMPLE_QUERIES = [
    ("how are conversation messages routed to the right model with a fallback chain", "api-chat"),
    ("starting sandboxed shell processes with approval modes and output history", "terminal-manager"),
    ("switching between sign-in screen, welcome page and main conversation view", "renderer-view-shell"),
]

SELSTAUSKUNFT = (
    "HINWEIS: semantischer Index, kann falsch liegen — Fundstelle vor "
    "Verwendung in der Datei verifizieren."
)


def canonical_repo_path(repo: str) -> str:
    return os.path.normcase(os.path.realpath(os.path.abspath(repo)))


def index_root() -> str:
    return os.path.join(os.path.expanduser("~"), ".blun", "codebase-index")


def index_dir(repo: str) -> str:
    raw = canonical_repo_path(repo) + "|" + _active_model_name
    key = hashlib.sha256(raw.encode()).hexdigest()[:12]
    return os.path.join(index_root(), key)


def normalize_model_selector(selector: str) -> str:
    if selector == "auto":
        return selector
    model = MODEL_ALIASES.get(selector, selector)
    if model not in KNOWN_MODELS:
        allowed = ", ".join(("auto", *MODEL_ALIASES, *KNOWN_MODELS))
        raise ValueError(f"Unsupported index model {selector!r}; choose one of: {allowed}")
    return model


def find_existing_indexes(repo: str) -> list[tuple[str, dict]]:
    root = index_root()
    if not os.path.isdir(root):
        return []
    canonical_repo = canonical_repo_path(repo)
    canonical_root = os.path.realpath(root)
    candidates = []
    for entry in os.scandir(root):
        if not entry.is_dir(follow_symlinks=False):
            continue
        directory = os.path.realpath(entry.path)
        try:
            if os.path.commonpath((canonical_root, directory)) != canonical_root:
                continue
        except ValueError:
            continue
        manifest_path = os.path.join(directory, "manifest.json")
        vectors_path = os.path.join(directory, "vectors.npy")
        if not os.path.isfile(manifest_path) or not os.path.isfile(vectors_path):
            continue
        try:
            manifest = load_manifest(directory)
        except (OSError, ValueError, TypeError):
            continue
        model = manifest.get("model")
        if model not in KNOWN_MODELS or manifest.get("dim") != KNOWN_MODELS[model]:
            continue
        manifest_repo = manifest.get("repo")
        if not isinstance(manifest_repo, str):
            continue
        if canonical_repo_path(manifest_repo) != canonical_repo:
            continue
        candidates.append((directory, manifest))
    return sorted(candidates, key=lambda item: item[0])


def resolve_existing_index(repo: str, selector: str = "auto") -> tuple[str, dict]:
    requested = normalize_model_selector(selector)
    candidates = find_existing_indexes(repo)
    model_order = AUTO_MODEL_ORDER if requested == "auto" else (requested,)
    for model in model_order:
        matches = [item for item in candidates if item[1]["model"] == model]
        if matches:
            return max(
                matches,
                key=lambda item: (str(item[1].get("built_at", "")), item[0]),
            )
    available = sorted({item[1]["model"] for item in candidates})
    suffix = f"; available for this repository: {', '.join(available)}" if available else ""
    raise FileNotFoundError(
        f"No compatible codebase index for {os.path.abspath(repo)!r} and selector {selector!r}{suffix}"
    )


def activate_model(model: str, dim: int | None = None) -> None:
    global _active_model_name, _active_dim
    normalized = normalize_model_selector(model)
    if normalized == "auto":
        raise ValueError("auto can select an existing index only; it cannot build a new one")
    expected_dim = KNOWN_MODELS[normalized]
    if dim is not None and dim != expected_dim:
        raise ValueError(
            f"Index dimension {dim} does not match {normalized} ({expected_dim})"
        )
    _active_model_name = normalized
    _active_dim = expected_dim


def git_files(repo: str) -> list[str]:
    out = subprocess.run(
        ["git", "ls-files"], cwd=repo, capture_output=True, text=True, check=True
    ).stdout.splitlines()
    return [
        f for f in out
        if f.lower().endswith(CODE_EXT) and not f.startswith(EXCLUDE_PREFIXES)
    ]


def git_head(repo: str) -> str:
    return subprocess.run(
        ["git", "rev-parse", "--short", "HEAD"],
        cwd=repo, capture_output=True, text=True, check=True,
    ).stdout.strip()


def sha_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def chunk_text(text: str) -> list[str]:
    step = CHUNK_SIZE - CHUNK_OVERLAP
    return [
        text[i : i + CHUNK_SIZE]
        for i in range(0, len(text), step)
        if text[i : i + CHUNK_SIZE].strip()
    ]


def load_model():
    from fastembed import TextEmbedding

    return TextEmbedding(_active_model_name)


def collect_chunks(repo: str, files: list[str]):
    """Liest Dateien, liefert (texts, meta, file_spans, read_errors)."""
    texts, meta, spans, read_errors = [], [], {}, 0
    for rel in files:
        try:
            with open(os.path.join(repo, rel), encoding="utf-8", errors="replace") as fh:
                chunks = chunk_text(fh.read())
        except OSError:
            read_errors += 1
            continue
        start = len(texts)
        for idx, chunk in enumerate(chunks):
            texts.append(chunk)
            meta.append(f"{rel}#{idx}")
        spans[rel] = [start, len(texts)]
    return texts, meta, spans, read_errors


def embed_into(model, texts: list[str], matrix, offset: int) -> None:
    """Streamt Embeddings batchweise DIREKT in die (vorallokierte) Matrix."""
    for i in range(0, len(texts), BATCH):
        vecs = list(model.embed(texts[i : i + BATCH]))
        matrix[offset + i : offset + i + len(vecs)] = np.asarray(
            vecs, dtype=np.float32
        )
        if i % (BATCH * 8) == 0:
            print(f"PROGRESS embedded={offset + i}", flush=True)


def write_matrix(dirpath: str, total: int):
    return np.lib.format.open_memmap(
        os.path.join(dirpath, "vectors.npy"),
        mode="w+", dtype=np.float32, shape=(total, _active_dim),
    )


def save_manifest(dirpath: str, repo: str, files_hashes: dict, spans: dict,
                  meta: list[str], build_s: float) -> None:
    manifest = {
        "repo": os.path.abspath(repo),
        "head": git_head(repo),
        "built_at": datetime.now(timezone.utc).isoformat(),
        "model": _active_model_name,
        "dim": _active_dim,
        "chunk_size": CHUNK_SIZE,
        "chunk_overlap": CHUNK_OVERLAP,
        "chunks": len(meta),
        "build_s": round(build_s, 2),
        "files": files_hashes,
        "spans": spans,
        "meta": meta,
    }
    with open(os.path.join(dirpath, "manifest.json"), "w", encoding="utf-8") as fh:
        json.dump(manifest, fh)


def load_manifest(dirpath: str) -> dict:
    with open(os.path.join(dirpath, "manifest.json"), encoding="utf-8") as fh:
        return json.load(fh)


def cmd_build(repo: str) -> None:
    t0 = time.perf_counter()
    files = git_files(repo)
    print(f"BUILD files={len(files)} repo={repo}")
    texts, meta, spans, read_errors = collect_chunks(repo, files)
    print(f"CHUNKS total={len(texts)} read_errors={read_errors}")

    dirpath = index_dir(repo)
    os.makedirs(dirpath, exist_ok=True)
    matrix = write_matrix(dirpath, len(texts))
    model = load_model()
    t_embed = time.perf_counter()
    embed_into(model, texts, matrix, 0)
    matrix.flush()
    build_s = time.perf_counter() - t0
    embed_s = time.perf_counter() - t_embed

    hashes = {rel: sha_file(os.path.join(repo, rel)) for rel in spans}
    save_manifest(dirpath, repo, hashes, spans, meta, build_s)
    mb = len(texts) * _active_dim * 4 / 1e6
    print(
        f"DONE build_s={build_s:.1f} embed_s={embed_s:.1f} chunks={len(texts)} "
        f"index_mb={mb:.1f} dir={dirpath}",
        flush=True,
    )


def cmd_update(repo: str) -> None:
    """Delta: geaenderte/neue Dateien neu einbetten, geloeschte entfernen."""
    t0 = time.perf_counter()
    dirpath = index_dir(repo)
    man = load_manifest(dirpath)
    old_hashes: dict = man["files"]
    old_spans: dict = man["spans"]
    old_meta: list[str] = man["meta"]
    old_mat = np.load(os.path.join(dirpath, "vectors.npy"))

    files = git_files(repo)
    texts_new, meta_new, spans_new, read_errors = collect_chunks(repo, files)
    new_hashes = {rel: sha_file(os.path.join(repo, rel)) for rel in spans_new}

    changed = {r for r, h in new_hashes.items() if old_hashes.get(r) != h}
    deleted = set(old_hashes) - set(new_hashes)
    if not changed and not deleted:
        print(f"UPDATE noop delta_s={time.perf_counter() - t0:.1f} head={git_head(repo)}")
        return

    # Behaltene Dateien: weder geaendert noch geloescht.
    kept_rels = [r for r in old_spans if r not in changed and r not in deleted]

    changed_chunks: dict[str, list[str]] = {}
    for rel in sorted(changed):
        with open(os.path.join(repo, rel), encoding="utf-8", errors="replace") as fh:
            changed_chunks[rel] = chunk_text(fh.read())

    total = sum(old_spans[r][1] - old_spans[r][0] for r in kept_rels) + sum(
        len(c) for c in changed_chunks.values()
    )
    matrix = write_matrix(dirpath, total)
    model = load_model()

    spans_out, meta_out, cursor = {}, [], 0
    for rel in kept_rels:
        a, b = old_spans[rel]
        matrix[cursor : cursor + (b - a)] = old_mat[a:b]
        spans_out[rel] = [cursor, cursor + (b - a)]
        meta_out.extend(old_meta[a:b])
        cursor += b - a
    for rel in sorted(changed_chunks):
        chunks = changed_chunks[rel]
        embed_into(model, chunks, matrix, cursor)
        spans_out[rel] = [cursor, cursor + len(chunks)]
        meta_out.extend(f"{rel}#{i}" for i in range(len(chunks)))
        cursor += len(chunks)
    matrix.flush()

    save_manifest(dirpath, repo, new_hashes, spans_out, meta_out,
                  time.perf_counter() - t0)
    print(
        f"UPDATE changed={len(changed)} deleted={len(deleted)} "
        f"delta_s={time.perf_counter() - t0:.1f} chunks={total}",
        flush=True,
    )


def freshness(repo: str, man: dict) -> str:
    current = git_head(repo)
    same = "== HEAD" if current == man["head"] else f"Index {man['head']} != HEAD {current} — 'update' faellig"
    return f"Index-Stand: {man['head']} ({man['built_at']}), {man['chunks']} Chunks | {same}"


def cmd_query(repo: str, question: str, top: int, selector: str = "auto") -> int:
    dirpath, man = resolve_existing_index(repo, selector)
    activate_model(man["model"], man["dim"])
    mat = np.load(os.path.join(dirpath, "vectors.npy"))
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    mat_n = mat / np.maximum(norms, 1e-12)

    model = load_model()
    t0 = time.perf_counter()
    qv = np.asarray(list(model.embed([question])), dtype=np.float32)[0]
    qv = qv / max(np.linalg.norm(qv), 1e-12)
    scores = mat_n @ qv
    idx = np.argsort(scores)[::-1][:top]
    latency_ms = (time.perf_counter() - t0) * 1000

    hits = [(man["meta"][i], round(float(scores[i]), 4)) for i in idx]
    print(SELSTAUSKUNFT)
    print(f"Index-Modell: {man['model']}")
    print(freshness(repo, man))
    for m, s in hits:
        print(f"  {s:.4f}  {m}")
    print(f"query_ms={latency_ms:.1f}")

    log_path = os.path.join(dirpath, "query-log.jsonl")
    with open(log_path, "a", encoding="utf-8") as fh:
        fh.write(json.dumps({
            "ts": datetime.now(timezone.utc).isoformat(),
            "query": question,
            "top": hits,
            "query_ms": round(latency_ms, 1),
            "index_head": man["head"],
        }) + "\n")
    return 0


def cmd_sample(repo: str, selector: str = "auto") -> int:
    """Gate: Stichprobe aus dem Subjekt-Baum, Ziel >= 2/3 in Top-5."""
    dirpath, man = resolve_existing_index(repo, selector)
    activate_model(man["model"], man["dim"])
    mat = np.load(os.path.join(dirpath, "vectors.npy"))
    norms = np.linalg.norm(mat, axis=1, keepdims=True)
    mat_n = mat / np.maximum(norms, 1e-12)
    model = load_model()

    hits_count = 0
    for question, expected in SAMPLE_QUERIES:
        qv = np.asarray(list(model.embed([question])), dtype=np.float32)[0]
        qv = qv / max(np.linalg.norm(qv), 1e-12)
        scores = mat_n @ qv
        idx = np.argsort(scores)[::-1][:5]
        hits = [(man["meta"][i], round(float(scores[i]), 4)) for i in idx]
        ok = any(expected in m for m, _ in hits)
        hits_count += ok
        print(f"SAMPLE {'HIT ' if ok else 'MISS'} {question!r} expected={expected}")
        for m, s in hits:
            print(f"    {s:.4f}  {m}")
    verdict = "PASS" if hits_count >= 2 else "FAIL"
    print(f"SAMPLE_RESULT {hits_count}/3 gate=2/3 -> {verdict}")
    return 0 if hits_count >= 2 else 1


def main() -> int:
    global _active_model_name, _active_dim
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    cmd, repo = sys.argv[1], sys.argv[2]
    selector = sys.argv[sys.argv.index("--model") + 1] if "--model" in sys.argv else (
        "auto" if cmd in ("query", "sample") else MODEL_NAME
    )
    if cmd in ("build", "update"):
        try:
            activate_model(selector)
        except ValueError as error:
            print(f"MODEL_ERROR {error}")
            return 2
        from fastembed import TextEmbedding

        dims = {m["model"]: m.get("dim") for m in TextEmbedding.list_supported_models()}
        if _active_model_name not in dims:
            print(f"MODEL_UNKNOWN {_active_model_name} — nicht im fastembed-Angebot")
            return 2
        _active_dim = int(dims[_active_model_name])
        print(f"MODEL {_active_model_name} dim={_active_dim}")
    if cmd == "build":
        cmd_build(repo)
        return 0
    if cmd == "update":
        cmd_update(repo)
        return 0
    if cmd == "query":
        top = 5
        if "--top" in sys.argv:
            top = int(sys.argv[sys.argv.index("--top") + 1])
        try:
            return cmd_query(repo, sys.argv[3], top, selector)
        except (FileNotFoundError, ValueError) as error:
            print(f"INDEX_ERROR {error}")
            return 2
    if cmd == "sample":
        try:
            return cmd_sample(repo, selector)
        except (FileNotFoundError, ValueError) as error:
            print(f"INDEX_ERROR {error}")
            return 2
    print(f"unknown command: {cmd}")
    return 2


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