import asyncio
import os
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple

from animuz_core.genai.anthropic_client import AnthropicAgentClient, AnthropicClient
from animuz_core.genai.openai_client import OpenAIAgentClient
from animuz_core.genai.openai_llm_client import OpenAILLMClient
from animuz_core.genai.ollama_client import OLlamaClient
from animuz_core.genai.tools import AnthropicRetrieverTool, OpenAIRetrieverTool
from animuz_core.embedding.embedding_client import EmbeddingClient
from animuz_core.ingest.azure_doc_ai_client import AzureDocAiClient
from animuz_core.ingest.custom_unstructured_client import MyUnstructuredClient
from animuz_core.pipelines.simple_rag import SimpleRAG
from animuz_core.utils import format_sources
from animuz_core.vectordb.utils import MyQdrantClient
from animuz_core.tooling import ToolSpec, build_anthropic_tool, build_openai_tool
from animuz_core.genai.base import BaseLLM


@dataclass
class RAGConfig:
    llm_provider: str = "openai"
    llm_model: Optional[str] = None
    embedding_backend: str = "http"
    embedding_host: Optional[str] = None
    embedding_port: Optional[int] = None
    qdrant_host: Optional[str] = None
    qdrant_port: Optional[int] = None
    qdrant_collection: str = "animuz"
    system_prompt: Optional[str] = None

    @classmethod
    def from_env(cls) -> "RAGConfig":
        return cls(
            llm_provider=os.getenv("ANIMUZ_LLM_PROVIDER", "openai"),
            llm_model=os.getenv("ANIMUZ_LLM_MODEL"),
            embedding_backend=os.getenv("ANIMUZ_EMBEDDING_BACKEND", "http"),
            embedding_host=os.getenv("EMBEDDING_HOST"),
            embedding_port=_int_env("EMBEDDING_PORT"),
            qdrant_host=os.getenv("QDRANT_HOST", "localhost"),
            qdrant_port=_int_env("QDRANT_PORT"),
            qdrant_collection=os.getenv("QDRANT_COLLECTION_NAME", "animuz"),
            system_prompt=os.getenv("ANIMUZ_SYSTEM_PROMPT"),
        )

    def with_defaults(self) -> "RAGConfig":
        if not self.embedding_host:
            self.embedding_host = "localhost"
        if not self.embedding_port:
            self.embedding_port = 12081
        if not self.qdrant_host:
            self.qdrant_host = "localhost"
        if not self.qdrant_port:
            self.qdrant_port = 6333
        if not self.llm_model:
            if self.llm_provider == "anthropic":
                self.llm_model = "claude-3-haiku-20240307"
            elif self.llm_provider == "ollama":
                self.llm_model = "llama3"
            else:
                self.llm_model = "gpt-4o"
        return self

    def validate(self, use_tools: bool) -> None:
        errors = []

        if self.embedding_backend == "http":
            if not self.embedding_host or not self.embedding_port:
                errors.append("Missing EMBEDDING_HOST or EMBEDDING_PORT")

        if not self.qdrant_host or not self.qdrant_port:
            errors.append("Missing QDRANT_HOST or QDRANT_PORT")

        if self.llm_provider == "openai" and not os.getenv("OPENAI_API_KEY"):
            errors.append("Missing OPENAI_API_KEY")
        if use_tools and self.llm_provider == "anthropic" and not os.getenv("ANTHROPIC_API_KEY"):
            errors.append("Missing ANTHROPIC_API_KEY")

        if errors:
            raise ValueError("; ".join(errors))


class RAG:
    def __init__(
        self,
        config: RAGConfig,
        tools: Optional[Iterable[ToolSpec]] = None,
        use_tools: Optional[bool] = None,
        llm: Optional[BaseLLM] = None,
        agent: Optional[Any] = None,
        unstructured_client: Optional[MyUnstructuredClient] = None,
        azure_doc_ai_client: Optional[AzureDocAiClient] = None,
    ) -> None:
        self.config = config.with_defaults()
        self.tools = list(tools or [])
        if use_tools is None:
            if self.tools:
                self.use_tools = True
            elif self.config.llm_provider in ("openai", "anthropic"):
                self.use_tools = True
            else:
                self.use_tools = False
        else:
            self.use_tools = use_tools

        self.config.validate(self.use_tools)

        self.embedding_client = EmbeddingClient(
            host=self.config.embedding_host,
            port=int(self.config.embedding_port),
        )
        self.db_client = MyQdrantClient(
            host=self.config.qdrant_host,
            port=int(self.config.qdrant_port),
            collection_name=self.config.qdrant_collection,
        )

        self._pipeline = SimpleRAG(
            embedding_client=self.embedding_client,
            db_client=self.db_client,
            LLM=_NullLLM(),
            unstructured_client=unstructured_client,
            azure_doc_ai_client=azure_doc_ai_client,
        )

        self._agent = None
        self._llm = None
        if self.use_tools:
            self._agent = agent or self._build_agent(self.tools)
        else:
            self._llm = llm or self._build_llm()

    @classmethod
    def quickstart(cls) -> "RAG":
        return cls(RAGConfig.from_env().with_defaults())

    async def add_doc(self, path: str, user_chat_id: str) -> None:
        await self._pipeline.add_doc(path, user_chat_id=user_chat_id)

    async def retrieve(self, query: str, user_chat_id: str) -> Tuple[List[str], List[Dict[str, Any]]]:
        return await self._pipeline.query_db(query, user_chat_id)

    async def chat(self, query: str, user_chat_id: str):
        if self.use_tools:
            messages = [{"role": "user", "content": query}]
            return await self._agent.get_reply_frontend(messages, user_chat_id=user_chat_id)

        search_results, _ = await self._pipeline.query_db(query, user_chat_id)
        prompt = (
            "\n Use the following context to answer the user query: "
            f"{format_sources(search_results)}. "
            "Be sure to include citations in your response (in the form of "
            '"Source x: [answer]" where x is an index). '\
            "\n\n User query: "
            f"{query}"
        )
        return await self._llm.get_reply(prompt)

    def _build_llm(self) -> BaseLLM:
        if self.config.llm_provider == "anthropic":
            return AnthropicClient(system_prompt=self.config.system_prompt)
        if self.config.llm_provider == "ollama":
            return OLlamaClient(system_prompt=self.config.system_prompt)
        if self.config.llm_provider == "openai":
            return OpenAILLMClient(
                system_prompt=self.config.system_prompt,
                model=self.config.llm_model or "gpt-4o",
            )
        raise ValueError("Unknown LLM provider")

    def _build_agent(self, tools: Iterable[ToolSpec]):
        if self.config.llm_provider == "openai":
            agent = OpenAIAgentClient()
            agent.add_tools(self._build_openai_tools(tools))
            return agent
        if self.config.llm_provider == "anthropic":
            agent = AnthropicAgentClient(system_prompt=self.config.system_prompt)
            agent.add_tools(self._build_anthropic_tools(tools))
            return agent
        raise ValueError("Tool use is not supported for this provider")

    def _build_openai_tools(self, tools: Iterable[ToolSpec]) -> Dict[str, Tuple[Any, Callable]]:
        tool_map = {"retriever": (OpenAIRetrieverTool.create_tool(), self._retriever)}
        self._add_custom_tools(tool_map, tools, build_openai_tool)
        return tool_map

    def _build_anthropic_tools(self, tools: Iterable[ToolSpec]) -> Dict[str, Tuple[Any, Callable]]:
        tool_map = {"retriever": (AnthropicRetrieverTool.create_tool(), self._retriever)}
        self._add_custom_tools(tool_map, tools, build_anthropic_tool)
        return tool_map

    def _add_custom_tools(
        self,
        tool_map: Dict[str, Tuple[Any, Callable]],
        tools: Iterable[ToolSpec],
        build_tool: Callable[[ToolSpec], Any],
    ) -> None:
        for spec in tools:
            if spec.name == "retriever":
                raise ValueError("Tool name 'retriever' is reserved")
            tool_map[spec.name] = (build_tool(spec), _wrap_tool_fn(spec.fn))

    async def _retriever(self, query: str, user_chat_id: str):
        return await self._pipeline.query_db(query, user_chat_id)


class _NullLLM(BaseLLM):
    def __init__(self) -> None:
        self.clear_history()

    async def get_reply(self, prompt: str):
        raise RuntimeError("LLM not configured for this operation")


def _wrap_tool_fn(fn: Callable[..., Any]) -> Callable[..., Any]:
    if asyncio.iscoroutinefunction(fn):
        return fn

    async def _async_wrapper(*args, **kwargs):
        return fn(*args, **kwargs)

    return _async_wrapper


def _int_env(name: str) -> Optional[int]:
    value = os.getenv(name)
    if not value:
        return None
    try:
        return int(value)
    except ValueError:
        return None
