"""
Poma ↔ Qdrant integration.

Simple, stateful wrapper around QdrantClient:
- define collection + model settings once in the class
- then mainly call upsert_poma_points(...) and get_cheatsheets(...)
"""

from __future__ import annotations

import os
from typing import Any, Awaitable, Callable

from qdrant_client import QdrantClient
from qdrant_client.http import models as qmodels

from .qdrant_poma_utils import (
    cheatsheets_from_results,
    ensure_collection,
    points_from_chunk_data,
)

__all__ = [
    "PomaQdrant",
]

_UNSET = object()


class PomaQdrant(QdrantClient):
    """
    Qdrant client with POMA defaults stored on the instance.

    Typical flow:
    1) configure this class once (collection/model/vector names)
    2) call upsert_poma_points(chunk_data)
    3) call get_cheatsheets(query="...")
    """

    def __init__(
        self,
        location: str | None = None,
        url: str | None = None,
        port: int | None = 6333,
        grpc_port: int = 6334,
        prefer_grpc: bool = False,
        https: bool | None = None,
        api_key: str | None = None,
        prefix: str | None = None,
        timeout: int | None = None,
        host: str | None = None,
        path: str | None = None,
        force_disable_check_same_thread: bool = False,
        grpc_options: dict[str, Any] | None = None,
        auth_token_provider: Callable[[], str]
        | Callable[[], Awaitable[str]]
        | None = None,
        cloud_inference: bool = False,
        local_inference_batch_size: int | None = None,
        check_compatibility: bool = True,
        pool_size: int | None = None,
        *,
        collection_name: str | None = None,
        dense_model: str | None = None,
        sparse_model: str | None = "Qdrant/bm25",
        dense_name: str = "dense",
        sparse_name: str = "sparse",
        dense_options: dict | None = None,
        sparse_options: dict | None = None,
        dense_size: int | None = None,
        distance: qmodels.Distance = qmodels.Distance.COSINE,
        store_chunk_details: bool = True,
        auto_create_collection: bool = False,
        **kwargs: Any,
    ):
        super().__init__(
            location=location,
            url=url,
            port=port,
            grpc_port=grpc_port,
            prefer_grpc=prefer_grpc,
            https=https,
            api_key=api_key,
            prefix=prefix,
            timeout=timeout,
            host=host,
            path=path,
            force_disable_check_same_thread=force_disable_check_same_thread,
            grpc_options=grpc_options,
            auth_token_provider=auth_token_provider,
            cloud_inference=cloud_inference,
            local_inference_batch_size=local_inference_batch_size,
            check_compatibility=check_compatibility,
            pool_size=pool_size,
            **kwargs,
        )
        if auto_create_collection and dense_size is None:
            raise ValueError(
                "dense_size must be set when auto_create_collection=True"
            )

        self.collection_name = collection_name
        self.dense_model = dense_model
        self.sparse_model = sparse_model
        self.dense_name = dense_name
        self.sparse_name = sparse_name
        self.dense_options = dense_options
        self.sparse_options = sparse_options
        self.dense_size = int(dense_size) if dense_size is not None else None
        self.distance = distance
        self.store_chunk_details = bool(store_chunk_details)
        self.auto_create_collection = bool(auto_create_collection)

    def _resolve_collection_name(self, collection_name: str | None) -> str:
        name = collection_name or self.collection_name
        if not name:
            raise ValueError(
                "collection_name is required (set it in __init__ or pass per call)"
            )
        return name

    def _resolve_dense_model(self, dense_model: str | None) -> str:
        model = dense_model or self.dense_model
        if not isinstance(model, str) or not model.strip():
            raise ValueError(
                "dense_model is required (set it in __init__ or pass per call)"
            )
        return model

    def _resolve_sparse_model(self, sparse_model: str | None | object) -> str | None:
        if sparse_model is _UNSET:
            return self.sparse_model
        return sparse_model

    def _ensure_collection_if_needed(
        self,
        collection_name: str,
        *,
        dense_name: str,
        sparse_name: str,
        sparse_model: str | None,
    ) -> None:
        if not self.auto_create_collection:
            return
        if hasattr(self, "collection_exists") and self.collection_exists(collection_name):
            return
        if self.dense_size is None:
            raise ValueError(
                "dense_size must be set to auto-create collections"
            )
        ensure_collection(
            self,
            collection_name,
            dense=(dense_name, self.dense_size),
            sparse=sparse_name if sparse_model else None,
            distance=self.distance,
        )

    @staticmethod
    def _merge_with_overrides(
        base_kwargs: dict[str, Any],
        override_kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Merge call kwargs with user overrides.
        Keys present in `override_kwargs` win over `base_kwargs`.
        """
        merged = dict(base_kwargs)
        merged.update(override_kwargs)
        return merged

    def upsert_poma_points(
        self,
        chunk_data: dict | str | os.PathLike[str],
        *,
        collection_name: str | None = None,
        dense_model: str | None = None,
        sparse_model: str | None | object = _UNSET,
        dense_name: str | None = None,
        sparse_name: str | None = None,
        dense_options: dict | None = None,
        sparse_options: dict | None = None,
        store_chunk_details: bool | None = None,
        poma_batch_size: int | None = 100,
        batch_size: int | None = 100,
        **kwargs: Any,
    ) -> Any:
        """
        Build POMA points and upsert them to Qdrant.
        If batching is enabled, use Qdrant's native `upload_points` path:
        - `poma_batch_size` controls outer slicing in this wrapper
        - `batch_size` controls internal Qdrant upload batching
        If a key is present in both convenience params and `kwargs`,
        the `kwargs` value is used.
        """
        collection_name = self._resolve_collection_name(collection_name)
        dense_model = self._resolve_dense_model(dense_model)
        sparse_model = self._resolve_sparse_model(sparse_model)
        dense_name = dense_name or self.dense_name
        sparse_name = sparse_name or self.sparse_name
        dense_options = self.dense_options if dense_options is None else dense_options
        sparse_options = self.sparse_options if sparse_options is None else sparse_options
        store_chunk_details = (
            self.store_chunk_details
            if store_chunk_details is None
            else bool(store_chunk_details)
        )

        self._ensure_collection_if_needed(
            collection_name,
            dense_name=dense_name,
            sparse_name=sparse_name,
            sparse_model=sparse_model,
        )

        points = points_from_chunk_data(
            chunk_data,
            dense_model=dense_model,
            sparse_model=sparse_model,
            dense_name=dense_name,
            sparse_name=sparse_name,
            dense_options=dense_options,
            sparse_options=sparse_options,
            store_chunk_details=store_chunk_details,
        )
        if not points:
            return None
        operation_kwargs = self._merge_with_overrides(
            {
                "collection_name": collection_name,
                "points": points,
                "poma_batch_size": poma_batch_size,
                "batch_size": batch_size,
            },
            kwargs,
        )
        resolved_poma_batch_size = operation_kwargs.pop("poma_batch_size", None)
        resolved_batch_size = operation_kwargs.pop("batch_size", None)

        if resolved_poma_batch_size is None and resolved_batch_size is None:
            return self.upsert(**operation_kwargs)

        if resolved_poma_batch_size is not None:
            resolved_poma_batch_size = int(resolved_poma_batch_size)
            if resolved_poma_batch_size <= 0:
                raise ValueError("poma_batch_size must be a positive integer")

        if resolved_batch_size is not None:
            resolved_batch_size = int(resolved_batch_size)
            if resolved_batch_size <= 0:
                raise ValueError("batch_size must be a positive integer")

        if "ordering" in operation_kwargs:
            raise ValueError(
                "ordering is not supported when batching is enabled; "
                "upload_points() path is used."
            )

        if "wait" not in operation_kwargs:
            operation_kwargs["wait"] = True

        points_to_upload = operation_kwargs.get("points")
        if points_to_upload is None:
            return None

        if resolved_batch_size is not None:
            operation_kwargs["batch_size"] = resolved_batch_size

        if resolved_poma_batch_size is None:
            return self.upload_points(**operation_kwargs)

        points_list = (
            list(points_to_upload)
            if not isinstance(points_to_upload, (list, tuple))
            else list(points_to_upload)
        )
        if not points_list:
            return None

        last_result = None
        for idx in range(0, len(points_list), resolved_poma_batch_size):
            batch_kwargs = dict(operation_kwargs)
            batch_kwargs["points"] = points_list[idx : idx + resolved_poma_batch_size]
            last_result = self.upload_points(**batch_kwargs)
        return last_result

    def get_cheatsheets(
        self,
        *,
        query: str | None = None,
        results: Any | None = None,
        query_obj: Any | None = None,
        prefetch: Any | None = None,
        using: str | None = None,
        collection_name: str | None = None,
        chunk_data: dict | str | os.PathLike[str] | None = None,
        limit: int = 5,
        query_filter: qmodels.Filter | None = None,
        with_vectors: bool | list[str] = False,
        dense_model: str | None = None,
        sparse_model: str | None | object = _UNSET,
        dense_name: str | None = None,
        sparse_name: str | None = None,
        dense_options: dict | None = None,
        sparse_options: dict | None = None,
        prefetch_limit: int | None = None,
        **kwargs: Any,
    ) -> list[dict]:
        """
        Return POMA cheatsheets in one of three modes:
        - `results`: convert precomputed Qdrant results to cheatsheets
        - `query_obj`: passthrough query/prefetch to query_points(...)
        - `query`: convenience text query path (dense-only or default hybrid RRF)
        If a key is present in both convenience params and `kwargs`,
        the `kwargs` value is used for query_points(...).
        """
        if results is not None:
            if (
                query is not None
                or query_obj is not None
                or prefetch is not None
                or using is not None
            ):
                raise ValueError(
                    "When results is provided, query/query_obj/prefetch/using must not be set"
                )
        elif query_obj is not None:
            collection_name = self._resolve_collection_name(collection_name)
            limit = max(1, int(limit))
            query_points_kwargs = self._merge_with_overrides(
                {
                    "collection_name": collection_name,
                    "query": query_obj,
                    "prefetch": prefetch,
                    "using": using,
                    "query_filter": query_filter,
                    "limit": limit,
                    "with_payload": True,
                    "with_vectors": with_vectors,
                },
                kwargs,
            )
            results = self.query_points(**query_points_kwargs)
        elif query is not None:
            if not isinstance(query, str) or not query.strip():
                raise ValueError("query must be a non-empty string when provided")

            collection_name = self._resolve_collection_name(collection_name)
            dense_model = self._resolve_dense_model(dense_model)
            sparse_model = self._resolve_sparse_model(sparse_model)
            dense_name = dense_name or self.dense_name
            sparse_name = sparse_name or self.sparse_name
            dense_options = self.dense_options if dense_options is None else dense_options
            sparse_options = self.sparse_options if sparse_options is None else sparse_options
            limit = max(1, int(limit))
            if prefetch_limit is None:
                prefetch_limit = max(limit, 20)
            prefetch_limit = prefetch_limit or max(limit, 20)
            if int(prefetch_limit) <= 0:
                raise ValueError("prefetch_limit must be a positive integer")
            prefetch_limit = int(prefetch_limit)

            if sparse_model:
                query_points_kwargs = self._merge_with_overrides(
                    {
                        "collection_name": collection_name,
                        "query": qmodels.FusionQuery(fusion=qmodels.Fusion.RRF),
                        "prefetch": [
                            qmodels.Prefetch(
                                query=qmodels.Document(
                                    text=query,
                                    model=dense_model,
                                    options=dense_options,
                                ),
                                using=dense_name,
                                limit=prefetch_limit,
                            ),
                            qmodels.Prefetch(
                                query=qmodels.Document(
                                    text=query,
                                    model=sparse_model,
                                    options=sparse_options,
                                ),
                                using=sparse_name,
                                limit=prefetch_limit,
                            ),
                        ],
                        "query_filter": query_filter,
                        "limit": limit,
                        "with_payload": True,
                        "with_vectors": with_vectors,
                    },
                    kwargs,
                )
                results = self.query_points(**query_points_kwargs)
            else:
                query_points_kwargs = self._merge_with_overrides(
                    {
                        "collection_name": collection_name,
                        "query": qmodels.Document(
                            text=query,
                            model=dense_model,
                            options=dense_options,
                        ),
                        "using": dense_name,
                        "query_filter": query_filter,
                        "limit": limit,
                        "with_payload": True,
                        "with_vectors": with_vectors,
                    },
                    kwargs,
                )
                results = self.query_points(**query_points_kwargs)
        else:
            raise ValueError("One of results, query_obj, or query must be provided")

        return cheatsheets_from_results(results, chunk_data=chunk_data)
