from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch


def _normalize_id(value: Any) -> str:
    return str(value)


def load_vectors(
    pt_path: str | Path,
    vector_key: str | None = None,
    id_key: str | None = None,
) -> tuple[torch.Tensor, list[str]]:
    """Load vectors (and optional ids) from a .pt file.

    Supported payloads:
    - Tensor: shape (N, D), ids become ["0", "1", ...]
    - Dict with vectors and ids keys (or custom keys via args)
    """
    pt_path = Path(pt_path)
    payload = torch.load(pt_path, map_location="cpu")

    if isinstance(payload, torch.Tensor):
        vectors = payload
        ids = [str(i) for i in range(vectors.shape[0])]
        return vectors, ids

    if not isinstance(payload, dict):
        raise ValueError(
            "Unsupported .pt payload. Expected a Tensor or a dict containing vectors."
        )

    possible_vector_keys = [vector_key] if vector_key else [
        "vectors",
        "embeddings",
        "tensor",
        "data",
    ]
    chosen_vector_key = None
    for key in possible_vector_keys:
        if key and key in payload:
            chosen_vector_key = key
            break

    if chosen_vector_key is None:
        raise ValueError(
            "Could not find vectors in dict payload. Pass --vector-key or store vectors under one "
            "of: vectors, embeddings, tensor, data."
        )

    vectors = payload[chosen_vector_key]
    if not isinstance(vectors, torch.Tensor):
        raise ValueError(f"Value at key '{chosen_vector_key}' is not a torch.Tensor.")

    possible_id_keys = [id_key] if id_key else ["ids", "keys", "labels", "names"]
    chosen_id_key = None
    for key in possible_id_keys:
        if key and key in payload:
            chosen_id_key = key
            break

    if chosen_id_key is None:
        ids = [str(i) for i in range(vectors.shape[0])]
    else:
        raw_ids = payload[chosen_id_key]
        if not isinstance(raw_ids, (list, tuple)):
            raise ValueError(f"Value at key '{chosen_id_key}' must be a list or tuple.")
        ids = [_normalize_id(x) for x in raw_ids]

    if vectors.ndim != 2:
        raise ValueError(f"Expected vectors shape (N, D), got {tuple(vectors.shape)}")
    if vectors.shape[0] != len(ids):
        raise ValueError(
            f"Vector row count ({vectors.shape[0]}) does not match ids count ({len(ids)})."
        )

    return vectors, ids


def load_metadata_records(metadata_path: str | Path) -> list[dict[str, Any]]:
    """Load metadata from .json (object or list) or .jsonl."""
    metadata_path = Path(metadata_path)
    suffix = metadata_path.suffix.lower()

    if suffix == ".jsonl":
        records: list[dict[str, Any]] = []
        with metadata_path.open("r", encoding="utf-8") as f:
            for line_number, line in enumerate(f, start=1):
                line = line.strip()
                if not line:
                    continue
                item = json.loads(line)
                if not isinstance(item, dict):
                    raise ValueError(
                        f"Expected object per line in jsonl, got {type(item)} at line {line_number}."
                    )
                records.append(item)
        return records

    if suffix == ".json":
        with metadata_path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, list):
            if not all(isinstance(x, dict) for x in data):
                raise ValueError("JSON list metadata must contain objects.")
            return data
        if isinstance(data, dict):
            # Allow id -> object mapping style
            if all(isinstance(v, dict) for v in data.values()):
                normalized: list[dict[str, Any]] = []
                for k, v in data.items():
                    row = dict(v)
                    row.setdefault("id", k)
                    normalized.append(row)
                return normalized
            return [data]
        raise ValueError("Unsupported JSON metadata format.")

    raise ValueError("Metadata file must have .json or .jsonl extension.")


def build_metadata_index(
    records: list[dict[str, Any]],
    metadata_key: str,
) -> dict[str, dict[str, Any]]:
    index: dict[str, dict[str, Any]] = {}
    for row in records:
        if metadata_key not in row:
            continue
        index[_normalize_id(row[metadata_key])] = row
    return index


def resolve_image_paths(
    ids: list[str],
    images_dir: str | Path | None,
    exts: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp", ".gif"),
) -> dict[str, str]:
    """Resolve image path by matching id as filename (id + extension)."""
    if images_dir is None:
        return {}

    base = Path(images_dir)
    if not base.exists() or not base.is_dir():
        raise ValueError(f"Images directory does not exist or is not a directory: {base}")

    resolved: dict[str, str] = {}
    for item_id in ids:
        candidate = base / item_id
        if candidate.exists() and candidate.is_file():
            resolved[item_id] = str(candidate.resolve())
            continue

        for ext in exts:
            candidate = base / f"{item_id}{ext}"
            if candidate.exists() and candidate.is_file():
                resolved[item_id] = str(candidate.resolve())
                break

    return resolved
