"""
Phase 62: RAG Pipeline — Query knowledge before acting.

Retrieves relevant context from registered data sources (channels,
knowledge base, files) and formats prompts with injected context.

Storage:
  state/broker/rag/sources.json — registered data sources
"""

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, _read_jsonl_sync

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

RAG_DIR = BROKER_DIR / "rag"
RAG_DIR.mkdir(parents=True, exist_ok=True)

SOURCES_PATH = RAG_DIR / "sources.json"

VALID_SOURCE_TYPES = ("channel", "knowledge_base", "file")


# ── Tokenization & Scoring ────────────────────────────────────────────────

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]:
    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: List[List[str]]) -> Dict[str, float]:
    n = len(docs)
    if n == 0:
        return {}
    df: Dict[str, int] = {}
    for doc_tokens in docs:
        for t in set(doc_tokens):
            df[t] = df.get(t, 0) + 1
    return {t: math.log((n + 1) / (d + 1)) + 1.0 for t, d in df.items()}


def _cosine_sim(a: Dict[str, float], b: Dict[str, float]) -> float:
    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)
    ma = math.sqrt(sum(v * v for v in a.values()))
    mb = math.sqrt(sum(v * v for v in b.values()))
    if ma == 0 or mb == 0:
        return 0.0
    return dot / (ma * mb)


# ── Source management helpers ──────────────────────────────────────────────

def _load_sources() -> Dict[str, dict]:
    if not SOURCES_PATH.exists():
        return {}
    try:
        return json.loads(SOURCES_PATH.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}


def _save_sources(sources: Dict[str, dict]) -> None:
    _write_json_sync(SOURCES_PATH, sources)


# ── Chunk retrieval from sources ───────────────────────────────────────────

def _chunks_from_channel(channel_name: str) -> List[dict]:
    """Load messages from a broker channel as chunks."""
    channels_dir = BROKER_DIR / "channels"
    path = channels_dir / f"{channel_name}.jsonl"
    if not path.exists():
        return []
    messages = _read_jsonl_sync(path)
    chunks = []
    for msg in messages:
        text = msg.get("message", msg.get("body", ""))
        if text:
            chunks.append({
                "source": f"channel:{channel_name}",
                "content": text,
                "metadata": {"ts": msg.get("ts"), "sender": msg.get("session_name")},
            })
    return chunks


def _chunks_from_knowledge() -> List[dict]:
    """Load entries from the knowledge base as chunks."""
    chunks = []
    try:
        from brynq.knowledge_mcp import KNOWLEDGE_DIR
        for f in KNOWLEDGE_DIR.glob("*.json"):
            if f.name == "index.jsonl":
                continue
            try:
                entry = json.loads(f.read_text("utf-8"))
                if entry.get("archived"):
                    continue
                text = f"{entry.get('topic', '')} — {entry.get('content', '')}"
                chunks.append({
                    "source": f"knowledge:{entry.get('knowledge_id', '')}",
                    "content": text,
                    "metadata": {"category": entry.get("category"),
                                 "confidence": entry.get("confidence")},
                })
            except (json.JSONDecodeError, OSError):
                continue
    except ImportError:
        pass
    return chunks


def _chunks_from_file(file_path: str) -> List[dict]:
    """Load a text file as chunks (one per paragraph)."""
    p = Path(file_path)
    if not p.exists():
        return []
    try:
        text = p.read_text("utf-8", errors="replace")
    except OSError:
        return []

    paragraphs = [para.strip() for para in text.split("\n\n") if para.strip()]
    return [
        {"source": f"file:{file_path}", "content": para, "metadata": {"line": i}}
        for i, para in enumerate(paragraphs)
    ]


def _gather_chunks(sources: Optional[List[str]] = None) -> List[dict]:
    """Gather chunks from all registered sources (or a subset)."""
    all_sources = _load_sources()
    if sources:
        all_sources = {k: v for k, v in all_sources.items() if k in sources}

    chunks: List[dict] = []
    for name, src in all_sources.items():
        stype = src.get("source_type", "")
        config = src.get("path_or_config", "")

        if stype == "channel":
            chunks.extend(_chunks_from_channel(config))
        elif stype == "knowledge_base":
            chunks.extend(_chunks_from_knowledge())
        elif stype == "file":
            chunks.extend(_chunks_from_file(config))

    # If no sources registered, still try knowledge base
    if not all_sources:
        chunks.extend(_chunks_from_knowledge())

    return chunks


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

def query_context(question: str, sources: Optional[List[str]] = None,
                  max_chunks: int = 5) -> List[dict]:
    """Retrieve relevant context chunks for a question.

    Args:
        question: Natural language question.
        sources: Optional list of source names to search. None = all.
        max_chunks: Maximum number of chunks to return.

    Returns:
        Ranked list of context chunks with scores.
    """
    q_tokens = _tokenize(question)
    if not q_tokens:
        return []

    chunks = _gather_chunks(sources)
    if not chunks:
        return []

    # TF-IDF scoring
    corpus = [_tokenize(c["content"]) for c in chunks]
    idf = _compute_idf(corpus)
    q_tf = _compute_tf(q_tokens)
    q_vec = {t: q_tf[t] * idf.get(t, 1.0) for t in q_tf}

    scored = []
    for chunk, doc_tokens in zip(chunks, corpus):
        d_tf = _compute_tf(doc_tokens)
        d_vec = {t: d_tf[t] * idf.get(t, 1.0) for t in d_tf}
        score = _cosine_sim(q_vec, d_vec)
        if score > 0:
            result = dict(chunk)
            result["score"] = round(score, 6)
            scored.append(result)

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


def build_prompt(question: str, context_chunks: List[dict]) -> str:
    """Format a prompt with retrieved context for an LLM.

    Args:
        question: The user's question.
        context_chunks: Chunks from query_context().

    Returns:
        Formatted prompt string.
    """
    parts = ["Use the following context to answer the question.\n"]

    if context_chunks:
        parts.append("## Context\n")
        for i, chunk in enumerate(context_chunks, 1):
            source = chunk.get("source", "unknown")
            score = chunk.get("score", 0)
            content = chunk.get("content", "")
            parts.append(f"### [{i}] {source} (relevance: {score:.3f})")
            parts.append(content)
            parts.append("")
    else:
        parts.append("No relevant context found.\n")

    parts.append(f"## Question\n{question}")
    return "\n".join(parts)


def add_source(name: str, source_type: str, path_or_config: str) -> dict:
    """Register a data source for RAG retrieval.

    Args:
        name: Unique name for this source.
        source_type: "channel", "knowledge_base", or "file".
        path_or_config: Channel name, path, or config string.

    Returns:
        Source record.
    """
    if not name or not name.strip():
        raise ValueError("Source name cannot be empty")
    if source_type not in VALID_SOURCE_TYPES:
        raise ValueError(f"Invalid source_type '{source_type}'. "
                         f"Must be one of: {VALID_SOURCE_TYPES}")

    sources = _load_sources()
    now = datetime.now(timezone.utc).isoformat()
    sources[name] = {
        "name": name,
        "source_type": source_type,
        "path_or_config": path_or_config,
        "added_at": now,
    }
    _save_sources(sources)
    return sources[name]


def list_sources() -> List[dict]:
    """List all registered data sources."""
    sources = _load_sources()
    return list(sources.values())


def remove_source(name: str) -> bool:
    """Remove a registered data source. Returns True if found."""
    sources = _load_sources()
    if name not in sources:
        return False
    del sources[name]
    _save_sources(sources)
    return True
