import inspect
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional

from animuz_core.genai.tools import InputSchema, Property, OpenAITool, AnthropicTool


@dataclass
class ToolSpec:
    name: str
    description: str
    parameters: Dict[str, Any]
    fn: Callable[..., Any]


@dataclass
class RetrieverToolSpec(ToolSpec):
    """ToolSpec that also carries its embedding and DB clients.

    Returned by ``qdrant_retriever()``. ``RAG`` detects this subclass to
    enable ``add_doc()`` and ``retrieve()`` without extra configuration.
    """

    embedding_client: Any = None
    db_client: Any = None


def _infer_parameters(fn: Callable[..., Any]) -> Dict[str, Any]:
    signature = inspect.signature(fn)
    properties: Dict[str, Dict[str, str]] = {}
    required = []

    type_map = {
        str: "string",
        int: "integer",
        float: "number",
        bool: "boolean",
    }

    for param in signature.parameters.values():
        if param.kind in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL):
            continue

        annotation = param.annotation
        json_type = type_map.get(annotation, "string")
        properties[param.name] = {
            "type": json_type,
            "description": "",
        }

        if param.default is inspect.Parameter.empty:
            required.append(param.name)

    return {
        "type": "object",
        "properties": properties,
        "required": required,
    }


def tool(
    fn: Optional[Callable[..., Any]] = None,
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
    parameters: Optional[Dict[str, Any]] = None,
):
    def wrap(func: Callable[..., Any]) -> ToolSpec:
        return ToolSpec(
            name=name or func.__name__,
            description=description or (func.__doc__ or ""),
            parameters=parameters or _infer_parameters(func),
            fn=func,
        )

    return wrap if fn is None else wrap(fn)


def _schema_to_input_schema(schema: Dict[str, Any]) -> InputSchema:
    properties = {
        key: Property(type=value.get("type", "string"), description=value.get("description", ""))
        for key, value in schema.get("properties", {}).items()
    }
    return InputSchema(
        type=schema.get("type", "object"),
        properties=properties,
        required=schema.get("required", []),
    )


def build_openai_tool(spec: ToolSpec) -> OpenAITool:
    return OpenAITool(
        name=spec.name,
        description=spec.description,
        parameters=_schema_to_input_schema(spec.parameters),
    )


def build_anthropic_tool(spec: ToolSpec) -> AnthropicTool:
    return AnthropicTool(
        name=spec.name,
        description=spec.description,
        input_schema=_schema_to_input_schema(spec.parameters),
    )


def qdrant_retriever(
    host: str = "localhost",
    port: int = 6333,
    collection: str = "animuz",
    *,
    embedding_client=None,
    api_key: Optional[str] = None,
) -> ToolSpec:
    """Factory that returns a ToolSpec named ``retriever`` backed by Qdrant hybrid search.

    The returned tool declares only a ``query`` parameter (what the LLM sees).
    ``user_chat_id`` is injected at runtime by ``BaseAgent.process_tool_use``.

    Args:
        host: Qdrant host.
        port: Qdrant port.
        collection: Qdrant collection name.
        embedding_client: Override the default ``ModalEmbeddingClient``.
        api_key: Qdrant API key (also reads ``QDRANT_API_KEY`` env var).
    """
    import os as _os
    from animuz_core.embedding.modal_embedding_client import ModalEmbeddingClient
    from animuz_core.vectordb.utils import MyQdrantClient

    _emb = embedding_client or ModalEmbeddingClient()
    _db = MyQdrantClient(host=host, port=port, collection_name=collection)

    async def query_db(query: str, user_chat_id: str, hybrid: bool = True, top_k: int = 3):
        emb_res = await _emb.get_embedding(query)
        indices = emb_res["sparse_embedding"].keys()
        values = emb_res["sparse_embedding"].values()
        search_results = await _db.hybrid_search(
            emb_res["embedding"], indices, values,
            top_k=top_k, user_chat_id=user_chat_id,
        )
        return [res.payload["text"] for res in search_results.points], search_results.points

    return RetrieverToolSpec(
        name="retriever",
        description="Retrieve documents related to a given query from a vector database",
        parameters={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The user query to be used for retrieval. The query must be comprehensive and in Thai",
                },
            },
            "required": ["query"],
        },
        fn=query_db,
        embedding_client=_emb,
        db_client=_db,
    )
