"""KnowledgeBase — orchestrator for RAG pipeline."""
from __future__ import annotations

from typing import Any, Awaitable, Callable

from logos_adk.context import get_run_context
from logos_adk.rag import Chunk, Document, DocumentLoadError
from logos_adk.rag.loaders import (
    CSVLoader,
    DocumentLoader,
    HTMLLoader,
    JSONLoader,
    MarkdownLoader,
    TextLoader,
    URLLoader,
)
from logos_adk.rag.splitter import RecursiveCharacterSplitter, TextSplitter
from logos_adk.rag.vector_store import VectorStore


class KnowledgeBase:
    """Orchestrates the RAG pipeline: load -> split -> embed -> store."""

    def __init__(
        self,
        store: VectorStore,
        embed_fn: Callable[[str], Awaitable[list[float]]],
        splitter: TextSplitter | None = None,
        loaders: list[DocumentLoader] | None = None,
        retrieval_learning: Any | None = None,
    ) -> None:
        self.store = store
        self.embed_fn = embed_fn
        self.splitter = splitter or RecursiveCharacterSplitter()
        self.loaders = loaders or [TextLoader(), MarkdownLoader(), CSVLoader(), JSONLoader(), HTMLLoader(), URLLoader()]
        self.retrieval_learning = retrieval_learning

    async def ingest(self, source: str | Document | list[Document]) -> list[str]:
        """Load, split, embed, store. Returns document IDs."""
        if isinstance(source, list):
            ids: list[str] = []
            for doc in source:
                ids.extend(await self._ingest_document(doc))
            return ids

        if isinstance(source, Document):
            return await self._ingest_document(source)

        return await self._ingest_file(source)

    async def ingest_text(self, text: str, metadata: dict[str, Any] | None = None) -> str:
        """Shortcut: index a text string. Returns document ID."""
        doc = Document(content=text, metadata=metadata or {})
        ids = await self._ingest_document(doc)
        return ids[0]

    async def search(
        self,
        query: str,
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        """Embed query and search vector store."""
        effective_query = query
        effective_limit = limit
        effective_filter = filter
        context = get_run_context()
        if self.retrieval_learning is not None:
            effective_query = self.retrieval_learning.rewrite_query(
                query,
                workspace_id=getattr(context, "workspace_id", None),
                task_type=getattr(context, "task_type", None),
                user_segment=getattr(context, "user_segment", None),
            )
            settings = self.retrieval_learning.resolve_settings(
                limit=limit,
                filter=filter,
                workspace_id=getattr(context, "workspace_id", None),
                task_type=getattr(context, "task_type", None),
                user_segment=getattr(context, "user_segment", None),
            )
            effective_limit = settings.limit
            effective_filter = settings.filter
        query_vector = await self.embed_fn(effective_query)
        chunks = await self.store.search(query_vector, limit=effective_limit, filter=effective_filter)
        if self.retrieval_learning is not None:
            chunks = self.retrieval_learning.rerank(
                chunks,
                workspace_id=getattr(context, "workspace_id", None),
            )
            self.retrieval_learning.record_retrieval(
                trace_id=getattr(context, "trace_id", None),
                session_id=getattr(context, "session_id", None),
                workspace_id=getattr(context, "workspace_id", None),
                task_type=getattr(context, "task_type", None),
                user_segment=getattr(context, "user_segment", None),
                prompt_name=getattr(context, "prompt_name", None),
                prompt_version=getattr(context, "prompt_version", None),
                original_query=query,
                effective_query=effective_query,
                limit=effective_limit,
                filter=effective_filter,
                chunks=chunks,
            )
        return chunks

    async def delete(self, document_id: str) -> None:
        """Remove a document and all its chunks."""
        await self.store.delete(document_id)

    async def clear(self) -> None:
        """Clear all indexed data."""
        await self.store.clear()

    async def _ingest_file(self, path: str) -> list[str]:
        """Load a file using the appropriate loader."""
        for loader in self.loaders:
            if loader.supports(path):
                docs = await loader.load(path)
                ids: list[str] = []
                for doc in docs:
                    ids.extend(await self._ingest_document(doc))
                return ids
        raise DocumentLoadError(f"No loader supports file: {path}")

    async def _ingest_document(self, doc: Document) -> list[str]:
        """Split, embed, and store a single document."""
        chunks = self.splitter.split(doc)
        if not chunks:
            chunks = [Chunk(
                content=doc.content,
                metadata=doc.metadata,
                document_id=doc.id,
            )]

        source_key = (
            doc.metadata.get("source_key")
            or doc.metadata.get("source")
            or doc.metadata.get("url")
            or doc.metadata.get("path")
            or doc.id
        )
        for chunk in chunks:
            chunk.metadata = dict(chunk.metadata or {})
            chunk.metadata.setdefault("source_key", source_key)
            chunk.metadata.setdefault("document_id", doc.id)
            chunk.metadata.setdefault("chunk_size_chars", len(chunk.content))
            chunk.metadata.setdefault("chunk_count", len(chunks))
            chunk.metadata.setdefault("source", doc.metadata.get("source"))
            chunk.embedding = await self.embed_fn(chunk.content)

        await self.store.add(chunks)
        if self.retrieval_learning is not None:
            metadata = dict(doc.metadata or {})
            self.retrieval_learning.register_source_document(
                workspace_id=metadata.get("workspace_id"),
                source_key=str(source_key),
                document_id=doc.id,
                content=doc.content,
                metadata=metadata,
                source_ref=metadata.get("source") or metadata.get("url") or metadata.get("path"),
            )
        return [doc.id]
