"""
Phase 61: Simple Vector Store for Semantic Search.

TF-IDF with cosine similarity — no external dependencies.
Documents are stored as JSON files; vocabulary and IDF are rebuilt on demand.

Storage:
  state/broker/vectors/
    {doc_id}.json — individual document with text, metadata, and TF vector
    vocab.json    — cached vocabulary + IDF weights
"""

import json
import math
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from brynq.server import BROKER_DIR, _write_json_sync

# ── Paths ──────────────────────────────────────────────────────────────────

VECTORS_DIR = BROKER_DIR / "vectors"
VECTORS_DIR.mkdir(parents=True, exist_ok=True)

VOCAB_PATH = VECTORS_DIR / "vocab.json"


# ── Tokenization & TF-IDF ─────────────────────────────────────────────────

def _tokenize(text: str) -> List[str]:
    """Lowercase word tokens."""
    return re.findall(r"[a-z0-9_]+", text.lower())


def _compute_tf(tokens: List[str]) -> Dict[str, float]:
    """Term frequency: count / total."""
    if not tokens:
        return {}
    counts: Dict[str, int] = {}
    for t in tokens:
        counts[t] = counts.get(t, 0) + 1
    total = len(tokens)
    return {t: c / total for t, c in counts.items()}


def _compute_idf(docs_tf: List[Dict[str, float]]) -> Dict[str, float]:
    """Inverse document frequency across all documents."""
    n = len(docs_tf)
    if n == 0:
        return {}
    doc_freq: Dict[str, int] = {}
    for tf in docs_tf:
        for term in tf:
            doc_freq[term] = doc_freq.get(term, 0) + 1
    return {t: math.log((n + 1) / (df + 1)) + 1.0 for t, df in doc_freq.items()}


def _tfidf_vector(tf: Dict[str, float], idf: Dict[str, float]) -> Dict[str, float]:
    """Build a TF-IDF vector."""
    return {t: tf_val * idf.get(t, 1.0) for t, tf_val in tf.items()}


def _cosine_similarity(a: Dict[str, float], b: Dict[str, float]) -> float:
    """Cosine similarity between two sparse vectors."""
    if not a or not b:
        return 0.0
    common = set(a) & set(b)
    if not common:
        return 0.0
    dot = sum(a[k] * b[k] for k in common)
    mag_a = math.sqrt(sum(v * v for v in a.values()))
    mag_b = math.sqrt(sum(v * v for v in b.values()))
    if mag_a == 0 or mag_b == 0:
        return 0.0
    return dot / (mag_a * mag_b)


# ── Disk I/O ──────────────────────────────────────────────────────────────

def _doc_path(doc_id: str) -> Path:
    return VECTORS_DIR / f"{doc_id}.json"


def _load_doc(doc_id: str) -> Optional[dict]:
    p = _doc_path(doc_id)
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_doc(doc: dict) -> None:
    _write_json_sync(_doc_path(doc["doc_id"]), doc)


def _load_all_docs() -> List[dict]:
    docs = []
    for f in VECTORS_DIR.glob("*.json"):
        if f.name == "vocab.json":
            continue
        try:
            docs.append(json.loads(f.read_text("utf-8")))
        except (json.JSONDecodeError, OSError):
            continue
    return docs


def _rebuild_idf() -> Dict[str, float]:
    """Rebuild IDF from all stored docs and cache to disk."""
    docs = _load_all_docs()
    all_tf = [d.get("tf", {}) for d in docs]
    idf = _compute_idf(all_tf)
    vocab = {"idf": idf, "doc_count": len(docs), "vocab_size": len(idf)}
    _write_json_sync(VOCAB_PATH, vocab)
    return idf


def _get_idf() -> Dict[str, float]:
    """Load cached IDF or rebuild."""
    if VOCAB_PATH.exists():
        try:
            data = json.loads(VOCAB_PATH.read_text("utf-8"))
            return data.get("idf", {})
        except (json.JSONDecodeError, OSError):
            pass
    return _rebuild_idf()


# ── Public API ─────────────────────────────────────────────────────────────

def add_document(doc_id: str, text: str, metadata: Optional[Dict[str, Any]] = None) -> str:
    """Index a document. Returns the doc_id."""
    if not doc_id or not doc_id.strip():
        raise ValueError("doc_id cannot be empty")
    if not text or not text.strip():
        raise ValueError("text cannot be empty")

    tokens = _tokenize(text)
    tf = _compute_tf(tokens)
    now = datetime.now(timezone.utc).isoformat()

    doc = {
        "doc_id": doc_id,
        "text": text,
        "metadata": metadata or {},
        "tf": tf,
        "token_count": len(tokens),
        "created_at": now,
        "updated_at": now,
    }
    _save_doc(doc)
    # Invalidate cached IDF
    if VOCAB_PATH.exists():
        VOCAB_PATH.unlink()
    return doc_id


def search(query: str, limit: int = 5) -> List[dict]:
    """Find documents similar to query by cosine similarity on TF-IDF vectors."""
    query_tokens = _tokenize(query)
    if not query_tokens:
        return []

    docs = _load_all_docs()
    if not docs:
        return []

    idf = _get_idf()
    query_tf = _compute_tf(query_tokens)
    query_vec = _tfidf_vector(query_tf, idf)

    scored = []
    for doc in docs:
        doc_vec = _tfidf_vector(doc.get("tf", {}), idf)
        sim = _cosine_similarity(query_vec, doc_vec)
        if sim > 0:
            scored.append({
                "doc_id": doc["doc_id"],
                "text": doc["text"],
                "metadata": doc.get("metadata", {}),
                "score": round(sim, 6),
            })

    scored.sort(key=lambda x: x["score"], reverse=True)
    return scored[:limit]


def delete_document(doc_id: str) -> bool:
    """Delete a document. Returns True if it existed."""
    p = _doc_path(doc_id)
    if not p.exists():
        return False
    p.unlink()
    if VOCAB_PATH.exists():
        VOCAB_PATH.unlink()
    return True


def get_document(doc_id: str) -> Optional[dict]:
    """Retrieve a document by ID."""
    doc = _load_doc(doc_id)
    if doc is None:
        return None
    return {
        "doc_id": doc["doc_id"],
        "text": doc["text"],
        "metadata": doc.get("metadata", {}),
        "token_count": doc.get("token_count", 0),
        "created_at": doc.get("created_at"),
        "updated_at": doc.get("updated_at"),
    }


def get_stats() -> dict:
    """Document count, vocab size, etc."""
    docs = _load_all_docs()
    all_terms: set = set()
    total_tokens = 0
    for d in docs:
        all_terms.update(d.get("tf", {}).keys())
        total_tokens += d.get("token_count", 0)
    return {
        "doc_count": len(docs),
        "vocab_size": len(all_terms),
        "total_tokens": total_tokens,
        "avg_doc_length": round(total_tokens / len(docs), 1) if docs else 0,
    }
