"""Unified RAG API — 3 knobs: prompt, llm, tools.

Provides chat (with agent tool loop), document ingestion, and retrieval
through a single class.

Usage::

    from animuz_core import RAG, MCP, qdrant_retriever, tool

    # Local tools
    rag = RAG(
        prompt=http_fetcher,
        llm="gpt-4o-mini",
        tools=[qdrant_retriever(host="localhost", port=6333, collection="animuz")],
    )
    output = await rag.chat(team_id, assistant_id, messages)

    # MCP tools
    rag = RAG(prompt=ddb_fetcher, llm="gpt-4o-mini", tools=MCP(url=mcp_url))
    output = await rag.chat(team_id, assistant_id, messages, user_context=ctx)

    # Plain chat — no tools
    rag = RAG(prompt=my_fetcher, llm="gpt-4o-mini")
    output = await rag.chat(team_id, assistant_id, messages)

    # Ingest + retrieve (requires a qdrant_retriever tool)
    await rag.add_doc("docs/intro.md", user_chat_id="team1|asst1")
    texts, points = await rag.retrieve("what is this?", user_chat_id="team1|asst1")
"""

import inspect
import logging
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Union

from animuz_core.tooling import ToolSpec, RetrieverToolSpec, build_openai_tool

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# MCP config
# ---------------------------------------------------------------------------

@dataclass
class MCP:
    """Configuration holder for an MCP tool server."""

    url: str
    api_key: Optional[str] = None


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def build_prompt(assistant: dict) -> str:
    """Assemble a system prompt from an assistant config dict.

    Args:
        assistant: Must contain ``"prompt"`` key. Optional: ``"customPrompt"``.

    Returns:
        Assembled system prompt string including today's date.

    Raises:
        ValueError: If ``"prompt"`` key is missing.
    """
    if "prompt" not in assistant:
        raise ValueError("Assistant config missing 'prompt' key")

    prompt = assistant["prompt"]
    if "customPrompt" in assistant:
        prompt += "\n" + assistant["customPrompt"]
    prompt += f"\nToday is {datetime.now().strftime('%A, %d %B %Y')}."
    return prompt


def _resolve_provider(model: str) -> str:
    """Detect LLM provider from model name."""
    if model.startswith("claude"):
        return "anthropic"
    return "openai"


def _build_tool_dict(tools: List[ToolSpec]) -> Dict[str, tuple]:
    """Convert a list of ToolSpec into the ``{name: (OpenAITool, fn)}`` dict
    expected by ``OpenAIAgentClientResponses``."""
    return {spec.name: (build_openai_tool(spec), spec.fn) for spec in tools}


def _find_retriever(tools: Optional[Union[List[ToolSpec], "MCP"]]) -> Optional[RetrieverToolSpec]:
    """Return the first RetrieverToolSpec in the tools list, or None."""
    if not isinstance(tools, list):
        return None
    for t in tools:
        if isinstance(t, RetrieverToolSpec):
            return t
    return None


# ---------------------------------------------------------------------------
# RAG
# ---------------------------------------------------------------------------

class RAG:
    """Unified RAG interface with 3 knobs: **prompt**, **llm**, **tools**.

    Args:
        prompt: Callable ``(team_id, assistant_id) -> dict`` returning an
            assistant config dict. May be sync or async.
        llm: Default model name string. Provider is auto-detected from name.
        tools: ``list[ToolSpec]`` for local tools, ``MCP(...)`` for MCP
            server, or ``None`` for plain chat.
    """

    def __init__(
        self,
        prompt: Callable,
        *,
        llm: str = "gpt-4o-mini",
        tools: Optional[Union[List[ToolSpec], MCP]] = None,
        base_url: Optional[str] = None,
    ):
        self._fetcher = prompt
        self._default_model = llm
        self._tools = tools
        self._base_url = base_url

        # Pre-build local tool dict for the agent
        self._tool_dict: Optional[Dict[str, tuple]] = None
        if isinstance(tools, list):
            self._tool_dict = _build_tool_dict(tools)

        # Extract retriever clients for add_doc / retrieve
        self._retriever = _find_retriever(tools)

    # ------------------------------------------------------------------
    # Chat
    # ------------------------------------------------------------------

    async def _fetch_assistant(self, team_id: str, assistant_id: str) -> dict:
        """Call the user-provided fetcher, handling both sync and async."""
        result = self._fetcher(team_id, assistant_id)
        if inspect.isawaitable(result):
            result = await result
        return result

    async def chat(
        self,
        team_id: str,
        assistant_id: str,
        messages: List[Dict[str, str]],
        *,
        user_context: Optional[Dict] = None,
        profiler: Any = None,
    ) -> list:
        """Run the chat pipeline: fetch prompt → create agent → call agent → clean output.

        Args:
            team_id: Tenant / team identifier.
            assistant_id: Assistant identifier.
            messages: Conversation messages ``[{"role": ..., "content": ...}]``.
            user_context: Optional user context dict (passed to MCP tools).
            profiler: Optional profiler instance.

        Returns:
            Frontend-ready list of message dicts (via ``clean_output``).
        """
        from animuz_core.streaming.message_mappers import clean_output, map_to_responses_api

        # 1. Fetch assistant config
        assistant = await self._fetch_assistant(team_id, assistant_id)

        # 2. Build system prompt
        system_prompt = build_prompt(assistant)

        # 3. Resolve model (assistant override or default)
        model = assistant.get("model") or self._default_model

        # 4. Detect provider
        provider = _resolve_provider(model)
        if provider == "anthropic":
            raise NotImplementedError("Anthropic agent support is not yet implemented")

        # 5. Create agent
        from animuz_core.genai.openai_responses_client import OpenAIAgentClientResponses

        user_chat_id = f"{team_id}|{assistant_id}"

        agent_kwargs: Dict[str, Any] = {
            "user_chat_id": user_chat_id,
            "model": model,
            "profiler": profiler,
            "base_url": self._base_url,
        }

        if isinstance(self._tools, MCP):
            agent_kwargs["mcp_url"] = self._tools.url
            agent_kwargs["mcp_api_key"] = self._tools.api_key
            agent_kwargs["user_context"] = user_context
        elif self._tool_dict is not None:
            agent_kwargs["tools"] = self._tool_dict

        agent = OpenAIAgentClientResponses(**agent_kwargs)

        # 6. Map messages
        mapped = [map_to_responses_api(m) for m in messages]

        # 7. Get reply
        result = await agent.get_reply_frontend(mapped, system_prompt)

        # 8. Clean output
        return clean_output(result)

    # ------------------------------------------------------------------
    # Ingest
    # ------------------------------------------------------------------

    async def add_doc(self, path: str, user_chat_id: str) -> None:
        """Ingest a document: parse → embed → upload to vector DB.

        Requires a ``qdrant_retriever()`` tool in the tools list.

        Args:
            path: File path to ingest.
            user_chat_id: Tenant ID for multi-tenant isolation in the vector DB.

        Raises:
            RuntimeError: If no retriever tool is configured.
        """
        if not self._retriever:
            raise RuntimeError(
                "add_doc() requires a qdrant_retriever() tool. "
                "Pass tools=[qdrant_retriever(...)] when creating the RAG instance."
            )

        from animuz_core.pipelines.simple_rag import SimpleRAG
        from animuz_core.genai.base import BaseLLM

        # Create a minimal pipeline just for ingest (no LLM needed)
        class _NullLLM(BaseLLM):
            def __init__(self):
                self.clear_history()
            async def get_reply(self, prompt: str):
                raise RuntimeError("LLM not needed for ingest")

        pipeline = SimpleRAG(
            embedding_client=self._retriever.embedding_client,
            db_client=self._retriever.db_client,
            LLM=_NullLLM(),
        )
        await pipeline.add_doc(path, user_chat_id=user_chat_id)

    # ------------------------------------------------------------------
    # Retrieve
    # ------------------------------------------------------------------

    async def retrieve(self, query: str, user_chat_id: str, top_k: int = 3):
        """Retrieve documents from the vector DB (embed query → hybrid search).

        Requires a ``qdrant_retriever()`` tool in the tools list.

        Args:
            query: Search query string.
            user_chat_id: Tenant ID for multi-tenant filtering.
            top_k: Number of results to return.

        Returns:
            Tuple of (list of text strings, list of scored points).

        Raises:
            RuntimeError: If no retriever tool is configured.
        """
        if not self._retriever:
            raise RuntimeError(
                "retrieve() requires a qdrant_retriever() tool. "
                "Pass tools=[qdrant_retriever(...)] when creating the RAG instance."
            )

        return await self._retriever.fn(query, user_chat_id=user_chat_id, top_k=top_k)
