"""
Brynq Runtime — File-based RAG (Retrieval-Augmented Generation).

A self-contained, zero-dependency RAG system for the CLI chat. Text files
are chunked by paragraph, scored with TF-IDF vectors, and retrieved via
cosine similarity.

Storage layout (under config.data_dir):
    rag/
        index.json          — file manifest: filepath -> list of chunk IDs
        chunks/
            {chunk_id}.json — individual chunk with text, TF vector, metadata

All pure stdlib Python: json, math, re, pathlib, hashlib, uuid.
"""

from __future__ import annotations

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


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

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


def _compute_tf(tokens: List[str]) -> Dict[str, float]:
    """Term frequency: count / total tokens."""
    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 a corpus of TF dicts."""
    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]:
    """Multiply TF values by IDF weights to produce 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 (dicts)."""
    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 helpers
# ---------------------------------------------------------------------------

def _rag_dir(data_dir: Path) -> Path:
    return data_dir / "rag"


def _index_path(data_dir: Path) -> Path:
    return _rag_dir(data_dir) / "index.json"


def _chunks_dir(data_dir: Path) -> Path:
    return _rag_dir(data_dir) / "chunks"


def _chunk_path(data_dir: Path, chunk_id: str) -> Path:
    return _chunks_dir(data_dir) / f"{chunk_id}.json"


def _ensure_dirs(data_dir: Path) -> None:
    _chunks_dir(data_dir).mkdir(parents=True, exist_ok=True)


def _load_index(data_dir: Path) -> dict:
    """Load the file index.

    Structure:
        {
            "files": {
                "/abs/path/to/file.txt": {
                    "filepath": "/abs/path/to/file.txt",
                    "hash": "<sha256 of file content>",
                    "chunk_ids": ["id1", "id2", ...],
                    "added_at": "<iso timestamp>"
                },
                ...
            }
        }
    """
    p = _index_path(data_dir)
    if not p.exists():
        return {"files": {}}
    try:
        data = json.loads(p.read_text("utf-8"))
        if "files" not in data:
            data["files"] = {}
        return data
    except (json.JSONDecodeError, OSError):
        return {"files": {}}


def _save_index(data_dir: Path, index: dict) -> None:
    _ensure_dirs(data_dir)
    p = _index_path(data_dir)
    p.write_text(json.dumps(index, indent=2, ensure_ascii=False), encoding="utf-8")


def _load_chunk(data_dir: Path, chunk_id: str) -> Optional[dict]:
    p = _chunk_path(data_dir, chunk_id)
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_chunk(data_dir: Path, chunk: dict) -> None:
    _ensure_dirs(data_dir)
    p = _chunk_path(data_dir, chunk["chunk_id"])
    p.write_text(json.dumps(chunk, indent=2, ensure_ascii=False), encoding="utf-8")


def _delete_chunk(data_dir: Path, chunk_id: str) -> None:
    p = _chunk_path(data_dir, chunk_id)
    if p.exists():
        p.unlink()


def _load_all_chunks(data_dir: Path) -> List[dict]:
    chunks: List[dict] = []
    d = _chunks_dir(data_dir)
    if not d.exists():
        return chunks
    for f in d.glob("*.json"):
        try:
            chunks.append(json.loads(f.read_text("utf-8")))
        except (json.JSONDecodeError, OSError):
            continue
    return chunks


def _file_hash(filepath: Path) -> str:
    """SHA-256 hex digest of file contents."""
    h = hashlib.sha256()
    h.update(filepath.read_bytes())
    return h.hexdigest()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def add_file(data_dir: Path, filepath: str) -> dict:
    """Read a text file, chunk by paragraphs, compute TF vectors, and store.

    If the file was previously indexed, its old chunks are removed and
    replaced with fresh ones (supports re-indexing after edits).

    Args:
        data_dir: Base data directory (e.g. ``~/.brynq``).
        filepath: Absolute or relative path to a text file.

    Returns:
        Summary dict with filepath, chunk count, and token count.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the file is empty or produces no chunks.
    """
    fp = Path(filepath).resolve()
    if not fp.exists():
        raise FileNotFoundError(f"File not found: {fp}")

    text = fp.read_text("utf-8", errors="replace")
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    if not paragraphs:
        raise ValueError(f"No content to index in: {fp}")

    # Remove previous index entry for this file (if any)
    remove_file(data_dir, str(fp))

    _ensure_dirs(data_dir)
    index = _load_index(data_dir)
    now = datetime.now(timezone.utc).isoformat()
    content_hash = _file_hash(fp)

    chunk_ids: List[str] = []
    total_tokens = 0

    for seq, para in enumerate(paragraphs):
        tokens = _tokenize(para)
        tf = _compute_tf(tokens)
        chunk_id = uuid.uuid4().hex[:12]

        chunk = {
            "chunk_id": chunk_id,
            "filepath": str(fp),
            "seq": seq,
            "text": para,
            "tf": tf,
            "token_count": len(tokens),
            "created_at": now,
        }
        _save_chunk(data_dir, chunk)
        chunk_ids.append(chunk_id)
        total_tokens += len(tokens)

    index["files"][str(fp)] = {
        "filepath": str(fp),
        "hash": content_hash,
        "chunk_ids": chunk_ids,
        "chunk_count": len(chunk_ids),
        "total_tokens": total_tokens,
        "added_at": now,
    }
    _save_index(data_dir, index)

    return {
        "filepath": str(fp),
        "chunks": len(chunk_ids),
        "tokens": total_tokens,
    }


def remove_file(data_dir: Path, filepath: str) -> bool:
    """Remove an indexed file and all its chunks.

    Args:
        data_dir: Base data directory.
        filepath: Path to the file (as it was indexed).

    Returns:
        True if the file was found and removed, False otherwise.
    """
    fp = str(Path(filepath).resolve())
    index = _load_index(data_dir)

    if fp not in index["files"]:
        return False

    entry = index["files"].pop(fp)
    for chunk_id in entry.get("chunk_ids", []):
        _delete_chunk(data_dir, chunk_id)

    _save_index(data_dir, index)
    return True


def list_files(data_dir: Path) -> List[dict]:
    """List all indexed files.

    Args:
        data_dir: Base data directory.

    Returns:
        List of file summaries (filepath, chunk_count, total_tokens, added_at).
    """
    index = _load_index(data_dir)
    results = []
    for entry in index["files"].values():
        results.append({
            "filepath": entry["filepath"],
            "chunks": entry.get("chunk_count", len(entry.get("chunk_ids", []))),
            "tokens": entry.get("total_tokens", 0),
            "added_at": entry.get("added_at", ""),
        })
    results.sort(key=lambda x: x["added_at"], reverse=True)
    return results


def query(data_dir: Path, question: str, top_k: int = 3) -> List[dict]:
    """Search indexed chunks for the most relevant matches to a question.

    Uses TF-IDF cosine similarity. IDF is computed on the fly across all
    stored chunks so the vocabulary stays fresh without a separate rebuild
    step.

    Args:
        data_dir: Base data directory.
        question: Natural language query.
        top_k: Maximum number of chunks to return.

    Returns:
        List of chunk dicts with ``text``, ``filepath``, ``score``, and
        ``seq`` (paragraph index within the source file), ranked by
        descending relevance.
    """
    q_tokens = _tokenize(question)
    if not q_tokens:
        return []

    chunks = _load_all_chunks(data_dir)
    if not chunks:
        return []

    # Build IDF across the full corpus
    all_tf = [c.get("tf", {}) for c in chunks]
    idf = _compute_idf(all_tf)

    # Query vector
    q_tf = _compute_tf(q_tokens)
    q_vec = _tfidf_vector(q_tf, idf)

    scored: List[dict] = []
    for chunk in chunks:
        d_vec = _tfidf_vector(chunk.get("tf", {}), idf)
        score = _cosine_similarity(q_vec, d_vec)
        if score > 0:
            scored.append({
                "text": chunk["text"],
                "filepath": chunk.get("filepath", ""),
                "seq": chunk.get("seq", 0),
                "score": round(score, 6),
            })

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


def build_context(chunks: List[dict]) -> str:
    """Format retrieved chunks as a context block for LLM prompt injection.

    Args:
        chunks: List of chunk dicts as returned by :func:`query`.

    Returns:
        A formatted string ready to prepend to an LLM prompt. Returns an
        empty string if no chunks are provided.
    """
    if not chunks:
        return ""

    parts = ["[Retrieved context from indexed files]\n"]
    for i, chunk in enumerate(chunks, 1):
        filepath = chunk.get("filepath", "unknown")
        score = chunk.get("score", 0)
        text = chunk.get("text", "")
        parts.append(f"--- [{i}] {filepath} (relevance: {score:.3f}) ---")
        parts.append(text)
        parts.append("")

    parts.append("[End of retrieved context]\n")
    return "\n".join(parts)
