"""
Poma ↔ Qdrant helper utils.

Pure helpers to:
- convert POMA chunk_data (.poma, dict, or PomaResult) into Qdrant PointStruct using models.Document
- convert Qdrant results into POMA cheatsheets
"""

from __future__ import annotations

import os
import uuid
import warnings
from collections import defaultdict
from typing import Any

from qdrant_client.http import models as qmodels

from poma.models import PomaResult
from poma.utils import unpack_poma_archive
from poma.retrieval import _get_relevant_chunks_for_ids, generate_cheatsheets

__all__ = [
    "chunk_uuid_string",
    "prepare_points_from_chunk_data",
    "points_from_chunk_data",
    "results_to_cheatsheet_inputs",
    "cheatsheets_from_results",
    "ensure_collection",
]


ChunkDataInput = dict[str, Any] | PomaResult | str | os.PathLike[str]


def chunk_uuid_string(file_id: str, chunkset_index: int) -> str:
    """
    Deterministic RFC4122 UUID (v5) from (file_id, chunkset_index).
    Same inputs -> same UUID. Different index -> different UUID.
    """
    if chunkset_index < 0:
        raise ValueError("chunkset_index must be non-negative")
    name = f"{file_id}#{chunkset_index}"
    return str(uuid.uuid5(uuid.NAMESPACE_URL, name))


def _chunk_data_from_input(
    chunk_data: ChunkDataInput,
) -> dict[str, Any]:
    """Accept chunk_data dict, PomaResult, or path to a .poma file; return chunk_data dict."""
    if isinstance(chunk_data, PomaResult):
        return _normalize_chunk_data_dict(chunk_data.to_dict())
    if isinstance(chunk_data, (str, os.PathLike)):
        path = os.fspath(chunk_data)
        if not path.lower().endswith(".poma"):
            raise ValueError(
                "Path must point to a .poma file; got {!r}".format(path)
            )
        return _normalize_chunk_data_dict(
            unpack_poma_archive(poma_archive_path=chunk_data)
        )
    return _normalize_chunk_data_dict(chunk_data)


def _get_file_id_fallback(chunksets: list[dict], chunks: list[dict]) -> str:
    file_id = None
    if chunksets:
        file_id = chunksets[0].get("file_id") or chunksets[0].get("tag")
    if not file_id and chunks:
        file_id = chunks[0].get("file_id") or chunks[0].get("tag")
    return file_id or "unknown"


def _normalize_chunk_data_dict(chunk_data: dict[str, Any]) -> dict[str, Any]:
    chunks = chunk_data.get("chunks", []) or []
    chunksets = chunk_data.get("chunksets", []) or []
    fallback_file_id = _get_file_id_fallback(chunksets, chunks)

    normalized_chunks: list[dict[str, Any]] = []
    for chunk in chunks:
        if not isinstance(chunk, dict):
            continue
        normalized = dict(chunk)
        file_id = (
            normalized.get("file_id")
            or normalized.get("tag")
            or fallback_file_id
        )
        if file_id:
            normalized["file_id"] = file_id
        normalized_chunks.append(normalized)

    normalized_chunksets: list[dict[str, Any]] = []
    for chunkset in chunksets:
        if not isinstance(chunkset, dict):
            continue
        normalized = dict(chunkset)
        file_id = (
            normalized.get("file_id")
            or normalized.get("tag")
            or fallback_file_id
        )
        if file_id:
            normalized["file_id"] = file_id
        if "contents" not in normalized and "to_embed" in normalized:
            normalized["contents"] = normalized.get("to_embed", "")
        if "to_embed" not in normalized and "contents" in normalized:
            normalized["to_embed"] = normalized.get("contents", "")
        raw_chunks = normalized.get("chunks", [])
        if isinstance(raw_chunks, tuple):
            normalized["chunks"] = list(raw_chunks)
        normalized_chunksets.append(normalized)

    normalized_chunk_data = dict(chunk_data)
    normalized_chunk_data["chunks"] = normalized_chunks
    normalized_chunk_data["chunksets"] = normalized_chunksets
    return normalized_chunk_data


def _build_chunk_index_lookup(
    chunks: list[dict[str, Any]],
) -> dict[str, dict[int, dict[str, Any]]]:
    lookup: dict[str, dict[int, dict[str, Any]]] = defaultdict(dict)
    for chunk in chunks:
        chunk_index = chunk.get("chunk_index")
        file_id = chunk.get("file_id") or "unknown"
        if isinstance(chunk_index, int):
            lookup[file_id][chunk_index] = chunk
    return dict(lookup)


def _build_doc_chunks(
    chunks: list[dict[str, Any]],
) -> dict[str, list[dict[str, Any]]]:
    doc_chunks: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for chunk in chunks:
        file_id = chunk.get("file_id") or "unknown"
        doc_chunks[file_id].append(chunk)
    return dict(doc_chunks)


def _chunk_details_for_chunkset(
    chunk_indices: list[int],
    *,
    file_id: str,
    doc_chunks_by_file_id: dict[str, list[dict[str, Any]]],
    chunk_lookup_by_file_id: dict[str, dict[int, dict[str, Any]]],
) -> list[dict[str, Any]]:
    if not chunk_indices:
        return []

    chunk_lookup = chunk_lookup_by_file_id.get(file_id, {})
    direct_chunks = [
        chunk_lookup[chunk_index]
        for chunk_index in chunk_indices
        if chunk_index in chunk_lookup
    ]
    if not direct_chunks:
        return []

    doc_chunks = doc_chunks_by_file_id.get(file_id, [])
    if not doc_chunks:
        return direct_chunks

    try:
        contextual_chunks = _get_relevant_chunks_for_ids(chunk_indices, doc_chunks)
    except (KeyError, TypeError, ValueError):
        return direct_chunks
    return contextual_chunks or direct_chunks


def prepare_points_from_chunk_data(
    chunk_data: ChunkDataInput,
    *,
    store_chunk_details: bool = True,
) -> tuple[list[str], list[str], list[dict]]:
    """
    Returns: (ids, documents, payloads)
    - ids: deterministic UUID strings (file_id + chunkset_index)
    - documents: chunkset.contents (this is what you embed)
    - payloads: metadata including "text"; "chunk_details" only if store_chunk_details=True
    """
    chunk_data = _chunk_data_from_input(chunk_data)
    chunksets: list[dict] = chunk_data.get("chunksets", [])
    chunks: list[dict] = chunk_data.get("chunks", [])
    if not chunksets:
        return [], [], []

    file_id = _get_file_id_fallback(chunksets, chunks)
    chunk_lookup_by_file_id = _build_chunk_index_lookup(chunks)
    doc_chunks_by_file_id = _build_doc_chunks(chunks)

    ids: list[str] = []
    documents: list[str] = []
    payloads: list[dict] = []

    for idx, chunkset in enumerate(chunksets):
        to_embed = chunkset.get("to_embed") or chunkset.get("contents", "")
        if not to_embed:
            continue

        chunkset_index = chunkset.get("chunkset_index")
        if chunkset_index is None:
            chunkset_index = idx

        chunkset_file_id = chunkset.get("file_id") or file_id
        point_id = chunk_uuid_string(chunkset_file_id, int(chunkset_index))

        raw_chunks = chunkset.get("chunks", [])
        chunk_indices = list(raw_chunks) if isinstance(raw_chunks, (list, tuple)) else []

        payload: dict = {
            "chunkset_index": int(chunkset_index),
            "chunks": chunk_indices,
            "file_id": chunkset_file_id,
            "text": to_embed,
        }
        if store_chunk_details:
            payload["chunk_details"] = _chunk_details_for_chunkset(
                chunk_indices,
                file_id=chunkset_file_id,
                doc_chunks_by_file_id=doc_chunks_by_file_id,
                chunk_lookup_by_file_id=chunk_lookup_by_file_id,
            )

        ids.append(point_id)
        documents.append(to_embed)
        payloads.append(payload)

    return ids, documents, payloads


def _validate_model_name(model: str | None, label: str) -> str:
    if not isinstance(model, str) or not model.strip():
        raise ValueError(f"{label} must be a non-empty string")
    return model


def points_from_chunk_data(
    chunk_data: ChunkDataInput,
    *,
    dense_model: str,
    sparse_model: str | None = None,
    dense_name: str = "dense",
    sparse_name: str = "sparse",
    dense_options: dict | None = None,
    sparse_options: dict | None = None,
    store_chunk_details: bool = True,
) -> list[qmodels.PointStruct]:
    """
    Build Qdrant PointStruct list using models.Document for dense/sparse vectors.
    """
    dense_model = _validate_model_name(dense_model, "dense_model")
    if sparse_model is not None:
        sparse_model = _validate_model_name(sparse_model, "sparse_model")

    ids, documents, payloads = prepare_points_from_chunk_data(
        chunk_data,
        store_chunk_details=store_chunk_details,
    )
    if not documents:
        return []

    points: list[qmodels.PointStruct] = []
    for pid, doc, payload in zip(ids, documents, payloads, strict=True):
        vector: dict[str, qmodels.Document] = {
            dense_name: qmodels.Document(
                text=doc,
                model=dense_model,
                options=dense_options,
            )
        }
        if sparse_model is not None:
            vector[sparse_name] = qmodels.Document(
                text=doc,
                model=sparse_model,
                options=sparse_options,
            )
        points.append(qmodels.PointStruct(id=pid, vector=vector, payload=payload))
    return points


def results_to_cheatsheet_inputs(
    results: Any,
) -> tuple[list[dict], list[dict]]:
    """
    Extract relevant_chunksets and all_chunks from Qdrant results.

    Accepts:
    - list of ScoredPoint
    - QueryResponse with .points
    - list of dicts with payload/metadata keys
    """
    if results is None:
        return [], []
    hits = results
    if not isinstance(results, list):
        hits = getattr(results, "points", results)
    if hits is None:
        return [], []

    relevant_chunksets: list[dict] = []
    all_chunks: list[dict] = []
    seen_chunks: set[tuple[str, int]] = set()

    for h in hits:
        payload: dict | None
        if isinstance(h, dict):
            payload = h.get("payload") or h.get("metadata")
        else:
            payload = getattr(h, "payload", None)
        if not isinstance(payload, dict):
            payload = {}

        payload_file_id = payload.get("file_id") or payload.get("tag") or "unknown"
        relevant_chunksets.append(
            {
                "chunkset_index": payload.get("chunkset_index"),
                "chunks": payload.get("chunks", []),
                "file_id": payload_file_id,
            }
        )

        chunk_details = payload.get("chunk_details", []) or []
        for ch in chunk_details:
            if not isinstance(ch, dict):
                continue
            cidx = ch.get("chunk_index")
            chunk_file_id = ch.get("file_id") or ch.get("tag") or payload_file_id
            if isinstance(cidx, int):
                dedupe_key = (str(chunk_file_id), cidx)
                if dedupe_key in seen_chunks:
                    continue
                seen_chunks.add(dedupe_key)
                normalized_chunk = dict(ch)
                normalized_chunk["file_id"] = chunk_file_id
                all_chunks.append(normalized_chunk)

    return relevant_chunksets, all_chunks


def cheatsheets_from_results(
    results: Any,
    *,
    chunk_data: ChunkDataInput | None = None,
) -> list[dict]:
    """
    Convert Qdrant results into POMA cheatsheets.

    If chunk_data is provided, its chunks are used instead of payload chunk_details.
    """
    relevant_chunksets, all_chunks = results_to_cheatsheet_inputs(results)
    if chunk_data is not None:
        chunk_data = _chunk_data_from_input(chunk_data)
        all_chunks = chunk_data.get("chunks", []) or []
    return generate_cheatsheets(relevant_chunksets, all_chunks) or []


def ensure_collection(
    client,
    name: str,
    *,
    dense: tuple[str, int] | dict | None = None,
    sparse: str | None = None,
    distance: qmodels.Distance = qmodels.Distance.COSINE,
) -> None:
    """
    Idempotent collection creation helper for Qdrant.
    Uses QdrantClient collection_exists() when available, otherwise checks get_collections().
    """
    if hasattr(client, "collection_exists"):
        exists = client.collection_exists(name)
    else:
        exists = name in {c.name for c in client.get_collections().collections}
    if exists:
        try:
            info = client.get_collection(name)
            params = getattr(getattr(info, "config", None), "params", None)
            vectors = getattr(params, "vectors", None) if params is not None else None
            sparse_vectors = getattr(params, "sparse_vectors", None) if params is not None else None
            vectors_desc = _format_vectors_config(vectors)
            sparse_desc = _format_sparse_vectors_config(sparse_vectors)
            warnings.warn(
                f"Collection {name!r} already exists; using existing config. "
                f"vectors={vectors_desc}; sparse_vectors={sparse_desc}",
                RuntimeWarning,
                stacklevel=2,
            )
        except Exception:
            warnings.warn(
                f"Collection {name!r} already exists; using existing config.",
                RuntimeWarning,
                stacklevel=2,
            )
        return

    dense_name = None
    dense_size = None
    if dense is not None:
        if isinstance(dense, dict):
            dense_name = dense.get("name")
            dense_size = dense.get("size")
        else:
            dense_name, dense_size = dense

    vectors_config = {}
    if dense_name is not None and dense_size is not None:
        vectors_config = {
            dense_name: qmodels.VectorParams(
                size=int(dense_size),
                distance=distance,
            )
        }

    kwargs = {
        "collection_name": name,
        "vectors_config": vectors_config,
    }
    if sparse is not None:
        kwargs["sparse_vectors_config"] = {sparse: qmodels.SparseVectorParams()}

    client.create_collection(**kwargs)


def _format_vectors_config(vectors: object) -> str:
    if vectors is None:
        return "None"
    if isinstance(vectors, dict):
        parts = []
        for name, params in vectors.items():
            size = getattr(params, "size", None)
            distance = getattr(params, "distance", None)
            parts.append(f"{name}(size={size}, distance={distance})")
        return ", ".join(parts) if parts else "None"
    size = getattr(vectors, "size", None)
    distance = getattr(vectors, "distance", None)
    return f"default(size={size}, distance={distance})"


def _format_sparse_vectors_config(sparse_vectors: object) -> str:
    if not sparse_vectors:
        return "None"
    if isinstance(sparse_vectors, dict):
        return ", ".join(sorted(str(k) for k in sparse_vectors.keys()))
    return str(sparse_vectors)
