"""RAG Knowledge Base — augment LLM decisions with historical facts.
Indexes: strategy performance history, market research, trade patterns."""

from __future__ import annotations
import json, hashlib, time
from pathlib import Path
from typing import Any, Optional

class SimpleRAG:
    """Lightweight RAG — no external vector DB. Uses keyword+embedding hybrid."""

    def __init__(self, index_dir: Optional[Path] = None):
        self._index: dict[str, list[dict]] = {}
        self._dir = index_dir or Path("data/rag_index")
        self._dir.mkdir(parents=True, exist_ok=True)

    def index_document(self, category: str, doc_id: str, content: str, metadata: dict = None) -> None:
        """Index a document under a category."""
        chunks = self._chunk(content, max_chars=500)
        entries = []
        for i, chunk in enumerate(chunks):
            entries.append({"doc_id": doc_id, "chunk_idx": i, "content": chunk,
                          "hash": hashlib.md5(chunk.encode()).hexdigest()[:8],
                          "meta": metadata or {}, "indexed_at": time.time()})
        if category not in self._index:
            self._index[category] = []
        self._index[category].extend(entries)
        self._save(category)

    def query(self, category: str, query_text: str, top_k: int = 5) -> list[dict]:
        """Simple keyword-based retrieval. Top-K by term overlap score."""
        self._load(category)
        entries = self._index.get(category, [])
        if not entries: return []
        query_terms = set(query_text.lower().split())
        scored = []
        for e in entries:
            content_terms = set(e["content"].lower().split())
            overlap = len(query_terms & content_terms)
            if overlap > 0:
                scored.append((overlap, e))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [e for _, e in scored[:top_k]]

    def build_prompt_context(self, query_text: str, categories: list[str] = None, top_k: int = 5) -> str:
        """Build a RAG context block for LLM prompt injection."""
        cats = categories or list(self._index.keys())
        all_results = []
        for cat in cats:
            results = self.query(cat, query_text, top_k)
            for r in results:
                all_results.append(f"[{cat}] {r['content'][:300]}")
        if not all_results: return ""
        return "## Historical Context (RAG)\n" + "\n".join(all_results[:top_k])

    @staticmethod
    def _chunk(text: str, max_chars: int = 500) -> list[str]:
        sentences = text.replace("\n", " ").split("。")
        chunks, current = [], ""
        for s in sentences:
            if len(current) + len(s) > max_chars and current:
                chunks.append(current.strip())
                current = s
            else:
                current += "。" + s if current else s
        if current.strip(): chunks.append(current.strip())
        return chunks

    def _save(self, category: str) -> None:
        f = self._dir / f"{category}.json"
        f.write_text(json.dumps(self._index.get(category, []), ensure_ascii=False, default=str), encoding="utf-8")

    def _load(self, category: str) -> None:
        f = self._dir / f"{category}.json"
        if f.exists():
            try:
                self._index[category] = json.loads(f.read_text(encoding="utf-8"))
            except Exception:
                pass

_rag: Optional[SimpleRAG] = None
def get_rag() -> SimpleRAG:
    global _rag
    if _rag is None: _rag = SimpleRAG()
    return _rag
