"""
cogniq.rag — Main RAG class.

The simplest possible API with production-grade internals.

    from cogniq import RAG
    r = RAG()
    r += "your document"
    print(r("your question"))
"""

from __future__ import annotations
import os, time, logging, hashlib
from pathlib import Path
from functools import lru_cache
from typing import List, Dict, Any, Optional, Union, Tuple

import numpy as np

from .embedders import BaseEmbedder, AutoEmbedder, TFIDFEmbedder
from .stores import BaseStore, MemoryStore, AutoStore
from .chunkers import BaseChunker, SmartChunker, AutoChunker
from .graph import KnowledgeGraph
from .plugins import PluginChain, Plugin

log = logging.getLogger("cogniq")


# ─────────────────────────────────────────────
# RESULT
# ─────────────────────────────────────────────

class SearchResult:
    """Rich result object with dict-like access."""

    def __init__(self, text: str, score: float, meta: Dict[str, Any]):
        self.text = text
        self.score = score
        self.meta = meta
        self.source = meta.get("source", "unknown")

    def __repr__(self):
        src = f" | src={self.source}" if self.source != "unknown" else ""
        return f"SearchResult(score={self.score:.3f}{src})\n  {self.text[:120]}..."

    def __getitem__(self, key):
        return self.meta[key]

    def get(self, key, default=None):
        return self.meta.get(key, default)


class AskResult:
    """Result of rag.ask() — contains answer + source chunks."""

    def __init__(self, answer: str, sources: List[SearchResult], query: str, latency_ms: float):
        self.answer = answer
        self.sources = sources
        self.query = query
        self.latency_ms = latency_ms

    def __str__(self):
        return self.answer

    def __repr__(self):
        return (f"AskResult(\n"
                f"  query={self.query!r}\n"
                f"  answer={self.answer[:200]!r}\n"
                f"  sources={len(self.sources)}, latency={self.latency_ms:.1f}ms\n)")

    def show(self):
        """Pretty print result with sources."""
        print(f"\n{'═'*60}")
        print(f"Q: {self.query}")
        print(f"{'─'*60}")
        print(f"A: {self.answer}")
        print(f"{'─'*60}")
        for i, s in enumerate(self.sources, 1):
            print(f"[{i}] score={s.score:.3f} | {s.source}")
            print(f"    {s.text[:120]}...")
        print(f"{'═'*60}")
        print(f"⏱  {self.latency_ms:.1f}ms")


# ─────────────────────────────────────────────
# MAIN RAG CLASS
# ─────────────────────────────────────────────

class RAG:
    """
    CogniQ RAG — Hybrid retrieval with knowledge graph reasoning.

    Quick Start:
        r = RAG()
        r += "your document text"
        answer = r("your question")

    Full control:
        r = RAG(
            embedder="sentence",          # tfidf | sentence | openai | ollama | auto
            store="memory",               # memory | faiss
            chunker="smart",              # smart | sentence | fixed
            llm=my_llm_fn,                # optional: fn(question, context) → str
            plugins=["dedup", "min_score"],
            graph=True,                   # enable knowledge graph layer
        )

    Operator API:
        r += "text"          # add document
        r += ["t1", "t2"]   # add multiple
        r("question")        # ask
        r["keyword"]         # raw search
        len(r)               # document count
    """

    def __init__(
        self,
        embedder: Union[str, BaseEmbedder] = "auto",
        store: Union[str, BaseStore] = "memory",
        chunker: Union[str, BaseChunker] = "smart",
        llm=None,
        plugins: Optional[List] = None,
        graph: bool = True,
        top_k: int = 5,
        graph_weight: float = 0.2,
        embed_cache_size: int = 10_000,
        verbose: bool = False,
        **embedder_kwargs,
    ):
        if verbose:
            logging.basicConfig(level=logging.INFO)

        # Embedder
        if isinstance(embedder, str):
            self._embedder: BaseEmbedder = AutoEmbedder(prefer=embedder, **embedder_kwargs)
        else:
            self._embedder = embedder

        # Store
        if isinstance(store, str):
            self._store: BaseStore = AutoStore(backend=store)
        else:
            self._store = store

        # Chunker
        if isinstance(chunker, str):
            self._chunker: BaseChunker = AutoChunker(strategy=chunker)
        else:
            self._chunker = chunker

        # Graph
        self._graph = KnowledgeGraph() if graph else None
        self._graph_weight = graph_weight

        # LLM (optional — for generation)
        self._llm = llm

        # Plugin chain
        self._plugins = PluginChain.build(plugins or [])

        # Embedding cache (LRU)
        self._embed_cache_size = embed_cache_size
        self._embed_cache: Dict[str, np.ndarray] = {}

        # Incremental TF-IDF support
        self._tfidf_docs: List[str] = []

        self.default_top_k = top_k
        self._doc_count = 0

        log.info(f"CogniQ RAG initialized | {self._embedder} | {self._store} | {self._chunker}")

    # ─── EMBED (with caching) ─────────────────────────────────

    def _embed(self, text: str) -> np.ndarray:
        key = hashlib.blake2b(text.encode(), digest_size=8).hexdigest()
        if key in self._embed_cache:
            return self._embed_cache[key]

        # For TF-IDF: update IDF with new text
        if isinstance(self._embedder, TFIDFEmbedder):
            self._embedder.partial_fit(text)

        vec = self._embedder.embed(text)
        if len(self._embed_cache) < self._embed_cache_size:
            self._embed_cache[key] = vec
        return vec

    def _embed_batch(self, texts: List[str]) -> np.ndarray:
        cached, uncached, uncached_idx = {}, [], []
        for i, t in enumerate(texts):
            key = hashlib.blake2b(t.encode(), digest_size=8).hexdigest()
            if key in self._embed_cache:
                cached[i] = self._embed_cache[key]
            else:
                uncached.append(t)
                uncached_idx.append(i)

        if uncached:
            if isinstance(self._embedder, TFIDFEmbedder):
                for t in uncached:
                    self._embedder.partial_fit(t)
            new_vecs = self._embedder.embed_batch(uncached)
            for j, (i, t) in enumerate(zip(uncached_idx, uncached)):
                cached[i] = new_vecs[j]
                key = hashlib.blake2b(t.encode(), digest_size=8).hexdigest()
                if len(self._embed_cache) < self._embed_cache_size:
                    self._embed_cache[key] = new_vecs[j]

        return np.vstack([cached[i] for i in range(len(texts))])

    # ─── ADD ──────────────────────────────────────────────────

    def add(
        self,
        text: Union[str, List[str], Path],
        source: Optional[str] = None,
        **meta,
    ) -> "RAG":
        """
        Add document(s) to the RAG.

        Args:
            text: string, list of strings, or file path
            source: label for this document (filename, URL, etc.)
            **meta: any additional metadata to store

        Returns:
            self (chainable)
        """
        # Normalize input
        if isinstance(text, Path) or (isinstance(text, str) and os.path.isfile(text)):
            path = Path(text)
            source = source or path.name
            text = path.read_text(encoding="utf-8", errors="replace")
        
        if isinstance(text, list):
            for i, t in enumerate(text):
                self.add(t, source=source or f"doc_{i}", **meta)
            return self

        # Plugin: pre_add
        meta["source"] = source or "inline"
        text, meta = self._plugins.pre_add(text, meta)

        # Chunk
        chunks = self._chunker.chunk(text, source=meta.get("source"))
        chunks = self._plugins.post_chunk(chunks, meta)

        if not chunks:
            return self

        # Merge metadata into chunks
        for chunk in chunks:
            chunk.update({k: v for k, v in meta.items() if k not in chunk})

        # Batch embed
        texts_to_embed = [
            self._plugins.pre_embed(c["text"]) for c in chunks
        ]
        vectors = self._embed_batch(texts_to_embed)

        # Store
        self._store.add(vectors, chunks)

        # Graph
        if self._graph is not None:
            self._graph.ingest(text, source=meta.get("source"))

        self._doc_count += 1
        log.debug(f"Added {len(chunks)} chunks from '{meta.get('source')}'")
        return self

    def learn(self, text, source=None, **meta):
        """Alias for add()"""
        return self.add(text, source=source, **meta)

    # ─── SEARCH ───────────────────────────────────────────────

    def search(
        self,
        query: str,
        top_k: Optional[int] = None,
        return_scores: bool = True,
    ) -> List[SearchResult]:
        """
        Pure vector similarity search.

        Returns:
            List of SearchResult objects, sorted by relevance.
        """
        k = top_k or self.default_top_k
        q_vec = self._embed(query)
        raw = self._store.search(q_vec, top_k=k * 2)  # over-fetch for plugin filtering

        # Graph boost
        if self._graph is not None and raw:
            boosted = []
            for doc, score in raw:
                graph_score = self._graph.score_relevance(query, doc.get("text", ""))
                final_score = (1 - self._graph_weight) * score + self._graph_weight * graph_score
                boosted.append((doc, final_score))
            raw = sorted(boosted, key=lambda x: x[1], reverse=True)

        # Plugin: post_search
        raw = self._plugins.post_search(raw, query)

        return [
            SearchResult(
                text=doc.get("text", ""),
                score=score,
                meta=doc,
            )
            for doc, score in raw[:k]
        ]

    # ─── ASK ──────────────────────────────────────────────────

    def ask(
        self,
        question: str,
        top_k: Optional[int] = None,
        llm=None,
        system_prompt: Optional[str] = None,
        **kwargs,
    ) -> AskResult:
        """
        Ask a question. Returns AskResult with answer + sources.

        If an LLM function is provided (at init or here), it generates a full answer.
        Otherwise, returns the most relevant chunk as the answer.

        LLM function signature:
            def my_llm(question: str, context: str, **kwargs) -> str: ...
        """
        t0 = time.perf_counter()
        sources = self.search(question, top_k=top_k)
        context = "\n\n---\n\n".join(s.text for s in sources)

        _llm = llm or self._llm

        if _llm:
            answer = _llm(question=question, context=context,
                          system_prompt=system_prompt, **kwargs)
        else:
            # No LLM — return top result as answer
            answer = sources[0].text if sources else "No relevant documents found."

        answer = self._plugins.post_answer(answer, question, context)

        latency = (time.perf_counter() - t0) * 1000
        return AskResult(answer=answer, sources=sources, query=question, latency_ms=latency)

    # ─── OPERATOR API ─────────────────────────────────────────

    def __iadd__(self, text):
        """r += "text"  or  r += ["t1", "t2"]"""
        self.add(text)
        return self

    def __call__(self, question: str, **kwargs) -> AskResult:
        """r("question") → AskResult"""
        return self.ask(question, **kwargs)

    def __getitem__(self, query: str) -> List[SearchResult]:
        """r["keyword"] → List[SearchResult]"""
        return self.search(query)

    def __len__(self) -> int:
        return len(self._store)

    def __repr__(self):
        return (f"RAG(docs={len(self._store)}, embedder={self._embedder}, "
                f"graph={'on' if self._graph else 'off'})")

    # ─── LLM INTEGRATIONS ─────────────────────────────────────

    def use_openai(self, model: str = "gpt-4o-mini", api_key: Optional[str] = None, **kwargs):
        """Attach OpenAI as the answer generator."""
        import openai, os
        client = openai.OpenAI(api_key=api_key or os.environ["OPENAI_API_KEY"])

        def _llm(question, context, system_prompt=None, **kw):
            sys = system_prompt or (
                "You are a helpful assistant. Answer based ONLY on the provided context. "
                "If the answer is not in the context, say so."
            )
            resp = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": sys},
                    {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
                ],
                **kw,
            )
            return resp.choices[0].message.content

        self._llm = _llm
        return self

    def use_ollama(self, model: str = "llama3.2", base_url: str = "http://localhost:11434", **kwargs):
        """Attach Ollama as the answer generator."""
        import urllib.request, json

        def _llm(question, context, system_prompt=None, **kw):
            sys = system_prompt or "Answer only from the provided context."
            prompt = f"System: {sys}\nContext:\n{context}\n\nQuestion: {question}\nAnswer:"
            body = json.dumps({"model": model, "prompt": prompt, "stream": False}).encode()
            req = urllib.request.Request(
                f"{base_url.rstrip('/')}/api/generate",
                data=body, headers={"Content-Type": "application/json"}, method="POST"
            )
            with urllib.request.urlopen(req, timeout=60) as r:
                return json.loads(r.read())["response"]

        self._llm = _llm
        return self

    def use_anthropic(self, model: str = "claude-haiku-4-5-20251001", api_key: Optional[str] = None, **kwargs):
        """Attach Anthropic Claude as the answer generator."""
        import anthropic, os
        client = anthropic.Anthropic(api_key=api_key or os.environ["ANTHROPIC_API_KEY"])

        def _llm(question, context, system_prompt=None, **kw):
            sys = system_prompt or "Answer only from the provided context."
            resp = client.messages.create(
                model=model,
                max_tokens=1024,
                system=sys,
                messages=[{
                    "role": "user",
                    "content": f"Context:\n{context}\n\nQuestion: {question}",
                }],
            )
            return resp.content[0].text

        self._llm = _llm
        return self

    def use_llm(self, fn):
        """Attach any callable as LLM: fn(question, context, **kwargs) → str"""
        self._llm = fn
        return self

    # ─── PERSISTENCE ──────────────────────────────────────────

    def save(self, path: Union[str, Path]):
        """Save RAG state to disk."""
        import pickle
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        state = {
            "store": self._store,
            "graph": self._graph,
            "embedder": self._embedder,
            "embed_cache": self._embed_cache,
        }
        with open(path, "wb") as f:
            pickle.dump(state, f)
        log.info(f"Saved to {path}")

    @classmethod
    def load(cls, path: Union[str, Path], **kwargs) -> "RAG":
        """Load RAG from disk."""
        import pickle
        with open(path, "rb") as f:
            state = pickle.load(f)
        rag = cls.__new__(cls)
        rag._store = state["store"]
        rag._graph = state.get("graph")
        rag._embedder = state["embedder"]
        rag._embed_cache = state.get("embed_cache", {})
        rag._llm = None
        rag._plugins = PluginChain([])
        rag._graph_weight = 0.2
        rag.default_top_k = 5
        rag._doc_count = 0
        rag._chunker = SmartChunker()
        rag._embed_cache_size = 10_000
        log.info(f"Loaded from {path} ({len(rag._store)} docs)")
        return rag

    # ─── STATS ────────────────────────────────────────────────

    def info(self) -> Dict[str, Any]:
        d = {
            "documents": len(self._store),
            "embedder": str(self._embedder),
            "store": type(self._store).__name__,
            "chunker": type(self._chunker).__name__,
            "graph": self._graph.stats() if self._graph else None,
            "embed_cache_size": len(self._embed_cache),
        }
        return d
