#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Rule Embeddings — Brain V6 Phase 1

Vectorizes brain rules (lessons, failures, patterns) via Ollama local embeddings.
Stores a lightweight JSON index for fast cosine similarity search.
Used by the RAG Sleep Cycle to find top-k similar rules for surgical merging.

Usage: python -m pulse.brain.rule_embeddings [--build|--search "query"|--status]
"""
from __future__ import annotations

import json
import logging
import math
import os
import time
import urllib.request
from pathlib import Path

from pulse.core.brain_utils import (
    HIPPOCAMPUS_DIR, STATE_DIR, now_iso as _now_iso,
)

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

EMBED_INDEX_FILE = STATE_DIR / "rule_embeddings.json"
EMBED_MODEL = os.environ.get("PULSE_EMBED_MODEL", "nomic-embed-text")
OLLAMA_ENDPOINT = os.environ.get("OLLAMA_ENDPOINT", "http://localhost:11434")

BRAIN_FILES = {
    "lessons": HIPPOCAMPUS_DIR / "Lessons" / "auto_lessons.md",
    "failures": HIPPOCAMPUS_DIR / "Failures" / "auto_fails.md",
    "patterns": HIPPOCAMPUS_DIR / "Patterns" / "auto_patterns.md",
}


def _embed_ollama(texts: list[str], model: str = EMBED_MODEL) -> list[list[float]]:
    """Batch-embed texts via Ollama /api/embed endpoint."""
    url = f"{OLLAMA_ENDPOINT}/api/embed"
    payload = json.dumps({"model": model, "input": texts}).encode()
    req = urllib.request.Request(url, data=payload, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read().decode())
    embeddings = data.get("embeddings", [])
    if len(embeddings) != len(texts):
        raise ValueError(f"Expected {len(texts)} embeddings, got {len(embeddings)}")
    return embeddings


def _embed_sentence_transformers(texts: list[str]) -> list[list[float]]:
    """Fallback: embed via sentence-transformers (if installed)."""
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer("all-MiniLM-L6-v2")
    vectors = model.encode(texts, show_progress_bar=False)
    return [v.tolist() for v in vectors]


def embed_texts(texts: list[str]) -> list[list[float]]:
    """Embed texts using best available backend: sentence-transformers (CPU-friendly) > Ollama."""
    if not texts:
        return []
    try:
        return _embed_sentence_transformers(texts)
    except (ImportError, Exception) as e:
        log.debug("sentence-transformers failed (%s), trying Ollama", e)
    try:
        return _embed_ollama(texts)
    except Exception as e:
        raise RuntimeError(f"No embedding backend available: {e}. pip install sentence-transformers")


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Pure-Python cosine similarity between two vectors."""
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


def _load_brain_rules() -> list[dict]:
    """Parse all brain files into a flat list of rules with metadata.

    Handles both formats:
      V4+ (## heading sections with paragraph body, separated by ---)
      Legacy (- bullet point lists)
    Skips YAML frontmatter, metadata lines (**Source:**, **Salience:**), and short entries.
    """
    rules = []
    for source, filepath in BRAIN_FILES.items():
        if not filepath.exists():
            continue
        text = filepath.read_text(encoding="utf-8", errors="replace")
        lines = text.splitlines()
        # Skip YAML frontmatter
        start = 0
        if lines and lines[0].strip() == "---":
            for i in range(1, len(lines)):
                if lines[i].strip() == "---":
                    start = i + 1
                    break
        # Parse ## heading sections
        current_body: list[str] = []
        in_section = False
        for line in lines[start:]:
            stripped = line.strip()
            if stripped.startswith("## "):
                if in_section and current_body:
                    rule_text = " ".join(current_body).strip()
                    if len(rule_text) >= 25:
                        rules.append({"text": rule_text, "source": source, "file": str(filepath.name)})
                current_body = []
                in_section = True
                continue
            if not in_section:
                # Legacy bullet format fallback
                if stripped.startswith("- ") and len(stripped) >= 25:
                    rules.append({"text": stripped[2:].strip(), "source": source, "file": str(filepath.name)})
                continue
            if stripped == "---" or stripped.startswith("# ") or stripped.startswith("> "):
                if current_body:
                    rule_text = " ".join(current_body).strip()
                    if len(rule_text) >= 25:
                        rules.append({"text": rule_text, "source": source, "file": str(filepath.name)})
                    current_body = []
                in_section = stripped.startswith("## ")
                continue
            if stripped.startswith("**") and ("Source:" in stripped or "Salience:" in stripped or "Date:" in stripped or "Consolidation" in stripped):
                continue
            if stripped:
                current_body.append(stripped)
        if in_section and current_body:
            rule_text = " ".join(current_body).strip()
            if len(rule_text) >= 25:
                rules.append({"text": rule_text, "source": source, "file": str(filepath.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()

    # Check if rebuild needed
    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,
            )
            index_mtime = EMBED_INDEX_FILE.stat().st_mtime
            if index_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:
        if verbose:
            log.info("No brain rules found — skipping index build")
        return {"status": "empty", "rules": 0}

    if verbose:
        log.info("Embedding %d rules via %s...", len(rules), EMBED_MODEL)

    # Batch embed (chunk to avoid timeouts on large brains)
    all_vectors = []
    batch_size = 50
    for i in range(0, len(rules), batch_size):
        batch_texts = [r["text"] for r in rules[i:i + batch_size]]
        vectors = embed_texts(batch_texts)
        all_vectors.extend(vectors)

    for i, rule in enumerate(rules):
        rule["vector"] = all_vectors[i]

    index_data = {
        "version": 2,
        "model": EMBED_MODEL,
        "built_at": _now_iso(),
        "rule_count": len(rules),
        "dimensions": len(all_vectors[0]) if all_vectors else 0,
        "rules": rules,
    }
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    EMBED_INDEX_FILE.write_text(json.dumps(index_data), encoding="utf-8")

    elapsed = round(time.time() - start, 2)
    if verbose:
        log.info("Index built: %d rules, %d dims in %ss", len(rules), index_data["dimensions"], elapsed)
    return {"status": "built", "rules": len(rules), "dimensions": index_data["dimensions"], "elapsed": elapsed}


def search_similar(query_text: str, top_k: int = 5) -> list[dict]:
    """Find the top-k most similar existing rules to a query text."""
    if not EMBED_INDEX_FILE.exists():
        return []

    data = json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8"))
    if data.get("version") != 2 or not data.get("rules"):
        return []

    query_vec = embed_texts([query_text])[0]

    scored = []
    for rule in data["rules"]:
        vec = rule.get("vector")
        if not vec:
            continue
        sim = _cosine_similarity(query_vec, vec)
        scored.append({"text": rule["text"], "source": rule["source"], "similarity": round(sim, 4)})

    scored.sort(key=lambda x: -x["similarity"])
    return scored[:top_k]


def get_new_rules() -> list[dict]:
    """Return brain rules that are NOT yet in the embedding index."""
    current_rules = _load_brain_rules()
    if not EMBED_INDEX_FILE.exists():
        return current_rules

    try:
        data = json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8"))
        indexed_texts = {r["text"] for r in data.get("rules", [])}
    except Exception:
        return current_rules

    return [r for r in current_rules if r["text"] not in indexed_texts]


if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [EMBED] %(message)s")

    if "--build" in sys.argv:
        result = build_rule_index(force="--force" in sys.argv, verbose=True)
        print(json.dumps(result, indent=2))
    elif "--search" in sys.argv:
        idx = sys.argv.index("--search")
        query = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else ""
        if not query:
            print("Usage: --search \"query text\"")
            sys.exit(1)
        results = search_similar(query, top_k=5)
        for r in results:
            print(f"  [{r['similarity']:.3f}] ({r['source']}) {r['text'][:100]}")
    elif "--status" in sys.argv:
        if EMBED_INDEX_FILE.exists():
            data = json.loads(EMBED_INDEX_FILE.read_text(encoding="utf-8"))
            print(f"Rules: {data.get('rule_count', 0)}, Dims: {data.get('dimensions', 0)}")
            print(f"Model: {data.get('model', '?')}, Built: {data.get('built_at', '?')}")
        else:
            print("No index found. Run --build first.")
    else:
        print("Usage: python -m pulse.brain.rule_embeddings [--build|--search \"q\"|--status]")
