"""Index build, cache, search, and delta operations for rule embeddings."""
from __future__ import annotations
import json, logging, time
from pulse.core.brain_utils import now_iso as _now_iso, read_jsonl_entries
from ._constants import EMBED_INDEX_FILE, BRAIN_FILES
from .backends import embed_texts, cosine_similarity

log = logging.getLogger("pulse.rule_embeddings")

_INDEX_CACHE: dict | None = None
_INDEX_CACHE_MTIME: float = 0.0


def _load_brain_rules() -> list[dict]:
    rules, _sqlite_loaded = [], set()
    try:
        from pulse.core.brain_db import get_conn
        conn = get_conn()
        for source in ("lessons", "failures", "patterns"):
            try:
                jsonl_name = f"auto_{source}.jsonl" if source != "failures" else "auto_fails.jsonl"
                rows = conn.execute(
                    f"SELECT content, title FROM {source} WHERE validity != 'archived' LIMIT 20000"
                ).fetchall()
                if rows:
                    for row in rows:
                        text = (row[0] or row[1] or "").strip()
                        if len(text) >= 25:
                            rules.append({"text": text, "source": source, "file": jsonl_name})
                    _sqlite_loaded.add(source)
            except Exception:
                pass
    except Exception:
        pass
    for source, jsonl_path in BRAIN_FILES.items():
        if source in _sqlite_loaded:
            continue
        for entry in read_jsonl_entries(jsonl_path):
            text = (entry.get("content") or entry.get("title") or entry.get("description", "")).strip()
            if len(text) >= 25:
                rules.append({"text": text, "source": source, "file": str(jsonl_path.name)})
    return rules


def build_rule_index(force: bool = False, verbose: bool = False) -> dict:
    """Build or refresh the vector embedding index for all brain rules."""
    start = time.time()
    if not force and EMBED_INDEX_FILE.exists():
        try:
            existing = json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8"))
            brain_mtime = max((f.stat().st_mtime for f in BRAIN_FILES.values() if f.exists()), default=0)
            if EMBED_INDEX_FILE.stat().st_mtime > brain_mtime and existing.get("version") == 2:
                if verbose:
                    log.info("Index up-to-date (%d rules)", len(existing.get("rules", [])))
                return {"status": "up_to_date", "rules": len(existing.get("rules", []))}
        except Exception:
            pass
    rules = _load_brain_rules()
    if not rules:
        return {"status": "empty", "rules": 0}
    all_vectors: list = []
    for i in range(0, len(rules), 50):
        all_vectors.extend(embed_texts([r["text"] for r in rules[i:i + 50]]))
    for i, rule in enumerate(rules):
        rule["vector"] = all_vectors[i]
    index_data = {"version": 2, "model": "rule_embeddings", "built_at": _now_iso(),
                  "rule_count": len(rules), "dimensions": len(all_vectors[0]) if all_vectors else 0, "rules": rules}
    EMBED_INDEX_FILE.parent.mkdir(parents=True, exist_ok=True)
    EMBED_INDEX_FILE.write_text(json.dumps(index_data), encoding="utf-8")
    global _INDEX_CACHE, _INDEX_CACHE_MTIME
    _INDEX_CACHE = index_data
    _INDEX_CACHE_MTIME = EMBED_INDEX_FILE.stat().st_mtime
    elapsed = round(time.time() - start, 2)
    if verbose:
        log.info("Index built: %d rules in %ss", len(rules), elapsed)
    return {"status": "built", "rules": len(rules), "dimensions": index_data["dimensions"], "elapsed": elapsed}


def _load_index_cached() -> dict | None:
    global _INDEX_CACHE, _INDEX_CACHE_MTIME
    if not EMBED_INDEX_FILE.exists():
        return None
    try:
        mtime = EMBED_INDEX_FILE.stat().st_mtime
    except OSError:
        return _INDEX_CACHE
    if _INDEX_CACHE is not None and mtime == _INDEX_CACHE_MTIME:
        return _INDEX_CACHE
    try:
        data = json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8"))
        _INDEX_CACHE, _INDEX_CACHE_MTIME = data, mtime
        return data
    except Exception:
        return _INDEX_CACHE


def search_similar(query_text: str, top_k: int = 5) -> list[dict]:
    """Find the top-k most similar brain rules to a query text."""
    data = _load_index_cached()
    if not data or data.get("version") != 2 or not data.get("rules"):
        return []
    query_vec = embed_texts([query_text])[0]
    scored = [
        {"text": r["text"], "source": r["source"], "similarity": round(cosine_similarity(query_vec, r["vector"]), 4)}
        for r in data["rules"] if r.get("vector")
    ]
    scored.sort(key=lambda x: -x["similarity"])
    return scored[:top_k]


def get_new_rules() -> list[dict]:
    """Return brain rules not yet in the embedding index."""
    current = _load_brain_rules()
    if not EMBED_INDEX_FILE.exists():
        return current
    try:
        indexed = {r["text"] for r in json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8")).get("rules", [])}
    except Exception:
        return current
    return [r for r in current if r["text"] not in indexed]
