"""Secondary index maintenance and lookup.

All index entries live in ``indexes_dir``.  Key layout::

    Tier 1 (structured fields):
    ("archived",   <bool>,      <run_hash>) -> b""
    ("active",     <bool>,      <run_hash>) -> b""
    ("experiment", <exp_name>,   <run_hash>) -> b""
    ("created_at", <timestamp>,  <run_hash>) -> b""
    ("tag",        <tag_name>,   <run_hash>) -> b""

    Tier 2 (hyperparameters — top-level scalars only):
    ("hparam",     <param_name>, <value>,    <run_hash>) -> b""

Values are always empty bytes — the run hash embedded in the key *is* the
payload.  Range scans on a prefix like ``("archived", False)`` yield all
non-archived run hashes.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from matyan_backend.fdb_types import transactional

from .fdb_client import get_directories

if TYPE_CHECKING:
    from matyan_backend.fdb_types import Database, DirectorySubspace, Transaction

_EMPTY = b""

# Index field name constants — Tier 1
_ARCHIVED = "archived"
_ACTIVE = "active"
_EXPERIMENT = "experiment"
_CREATED_AT = "created_at"
_TAG = "tag"

# Index field name constants — Tier 2
_HPARAM = "hparam"

_INDEXABLE_SCALAR_TYPES = (int, float, str, bool)


def _idx_dir() -> DirectorySubspace:
    return get_directories().indexes


# ---------------------------------------------------------------------------
# Write helpers
# ---------------------------------------------------------------------------


@transactional
def index_run(
    tr: Transaction,
    run_hash: str,
    *,
    is_archived: bool = False,
    active: bool = True,
    created_at: float = 0.0,
    experiment_name: str | None = None,
    tag_names: list[str] | None = None,
) -> None:
    """Write all Tier 1 index entries for a run (idempotent)."""
    idx = _idx_dir()
    tr[idx.pack((_ARCHIVED, is_archived, run_hash))] = _EMPTY
    tr[idx.pack((_ACTIVE, active, run_hash))] = _EMPTY
    tr[idx.pack((_CREATED_AT, created_at, run_hash))] = _EMPTY
    if experiment_name is not None:
        tr[idx.pack((_EXPERIMENT, experiment_name, run_hash))] = _EMPTY
    for tn in tag_names or []:
        tr[idx.pack((_TAG, tn, run_hash))] = _EMPTY


@transactional
def deindex_run(tr: Transaction, run_hash: str) -> None:
    """Remove *all* index entries (Tier 1 + Tier 2) for a run."""
    idx = _idx_dir()
    for prefix in (_ARCHIVED, _ACTIVE, _EXPERIMENT, _CREATED_AT, _TAG, _HPARAM):
        r = idx.range((prefix,))
        for kv in tr.get_range(r.start, r.stop):
            key_tuple = idx.unpack(kv.key)
            if key_tuple[-1] == run_hash:
                del tr[kv.key]


@transactional
def update_index_field(
    tr: Transaction,
    run_hash: str,
    field: str,
    old_value: object,
    new_value: object,
) -> None:
    """Atomically swap an index entry: delete old, write new.

    If *new_value* is ``None`` the entry is simply removed (no replacement).
    """
    idx = _idx_dir()
    if old_value is not None:
        del tr[idx.pack((field, old_value, run_hash))]
    if new_value is not None:
        tr[idx.pack((field, new_value, run_hash))] = _EMPTY


@transactional
def add_tag_index(tr: Transaction, run_hash: str, tag_name: str) -> None:
    idx = _idx_dir()
    tr[idx.pack((_TAG, tag_name, run_hash))] = _EMPTY


@transactional
def remove_tag_index(tr: Transaction, run_hash: str, tag_name: str) -> None:
    idx = _idx_dir()
    del tr[idx.pack((_TAG, tag_name, run_hash))]


@transactional
def remove_all_tag_indexes_for_tag(tr: Transaction, tag_name: str) -> None:
    """Remove all ``("tag", tag_name, *)`` index entries."""
    idx = _idx_dir()
    r = idx.range((_TAG, tag_name))
    del tr[r.start : r.stop]


@transactional
def rename_experiment_index(tr: Transaction, old_name: str, new_name: str) -> None:
    """Move all ``("experiment", old_name, *)`` entries to ``new_name``."""
    idx = _idx_dir()
    r = idx.range((_EXPERIMENT, old_name))
    run_hashes = [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]
    del tr[r.start : r.stop]
    for rh in run_hashes:
        tr[idx.pack((_EXPERIMENT, new_name, rh))] = _EMPTY


# ---------------------------------------------------------------------------
# Tier 2 — hparam index write helpers
# ---------------------------------------------------------------------------


@transactional
def index_hparams(tr: Transaction, run_hash: str, hparams: dict[str, Any]) -> None:
    """Write hparam index entries for all top-level scalar values (idempotent)."""
    idx = _idx_dir()
    for key, val in hparams.items():
        if isinstance(val, _INDEXABLE_SCALAR_TYPES):
            tr[idx.pack((_HPARAM, key, val, run_hash))] = _EMPTY


@transactional
def deindex_hparams(tr: Transaction, run_hash: str) -> None:
    """Remove all hparam index entries for a run."""
    idx = _idx_dir()
    r = idx.range((_HPARAM,))
    for kv in tr.get_range(r.start, r.stop):
        key_tuple = idx.unpack(kv.key)
        if key_tuple[-1] == run_hash:
            del tr[kv.key]


# ---------------------------------------------------------------------------
# Lookup helpers
# ---------------------------------------------------------------------------


@transactional
def lookup_by_archived(tr: Transaction, archived: bool) -> list[str]:
    idx = _idx_dir()
    r = idx.range((_ARCHIVED, archived))
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]


@transactional
def lookup_by_active(tr: Transaction, active: bool) -> list[str]:
    idx = _idx_dir()
    r = idx.range((_ACTIVE, active))
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]


@transactional
def lookup_by_experiment(tr: Transaction, experiment_name: str) -> list[str]:
    idx = _idx_dir()
    r = idx.range((_EXPERIMENT, experiment_name))
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]


@transactional
def lookup_by_created_at(
    tr: Transaction,
    start: float | None = None,
    end: float | None = None,
) -> list[str]:
    idx = _idx_dir()
    if start is None and end is None:
        r = idx.range((_CREATED_AT,))
        return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]

    begin_key = idx.pack((_CREATED_AT, start or 0.0))
    end_key = idx.pack((_CREATED_AT, end)) if end is not None else idx.range((_CREATED_AT,)).stop
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(begin_key, end_key)]


@transactional
def lookup_by_tag(tr: Transaction, tag_name: str) -> list[str]:
    idx = _idx_dir()
    r = idx.range((_TAG, tag_name))
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]


@transactional
def lookup_all_run_hashes(tr: Transaction) -> list[str]:
    """Return all indexed run hashes (ordered by creation time)."""
    idx = _idx_dir()
    r = idx.range((_CREATED_AT,))
    return [idx.unpack(kv.key)[2] for kv in tr.get_range(r.start, r.stop)]


# ---------------------------------------------------------------------------
# Tier 2 — hparam lookup helpers
# ---------------------------------------------------------------------------


@transactional
def lookup_by_hparam_eq(tr: Transaction, param_name: str, value: Any) -> list[str]:  # noqa: ANN401
    """Exact-match lookup: return run hashes where ``hparams[param_name] == value``."""
    idx = _idx_dir()
    r = idx.range((_HPARAM, param_name, value))
    return [idx.unpack(kv.key)[3] for kv in tr.get_range(r.start, r.stop)]


@transactional
def lookup_by_hparam_range(
    tr: Transaction,
    param_name: str,
    lo: Any = None,  # noqa: ANN401
    hi: Any = None,  # noqa: ANN401
) -> list[str]:
    """Range lookup on a hparam value.  Bounds are ``[lo, hi)`` (inclusive low, exclusive high).

    Pass ``None`` for an open bound.
    """
    idx = _idx_dir()
    full_range = idx.range((_HPARAM, param_name))
    begin_key = idx.pack((_HPARAM, param_name, lo)) if lo is not None else full_range.start
    end_key = idx.pack((_HPARAM, param_name, hi)) if hi is not None else full_range.stop
    return [idx.unpack(kv.key)[3] for kv in tr.get_range(begin_key, end_key)]


# ---------------------------------------------------------------------------
# Rebuild
# ---------------------------------------------------------------------------


@transactional
def _clear_all_indexes(tr: Transaction) -> None:
    idx = _idx_dir()
    for prefix in (_ARCHIVED, _ACTIVE, _EXPERIMENT, _CREATED_AT, _TAG, _HPARAM):
        r = idx.range((prefix,))
        del tr[r.start : r.stop]


def rebuild_indexes(db: Database) -> int:
    """Drop and rebuild all indexes (Tier 1 + Tier 2) from run data.

    Returns the number of runs indexed.
    """
    from . import entities, runs  # deferred to avoid circular import  # noqa: PLC0415

    _clear_all_indexes(db)

    all_hashes = runs.list_run_hashes(db)
    count = 0
    for rh in all_hashes:
        meta = runs.get_run_meta(db, rh)
        if not meta:
            continue

        exp_name: str | None = None
        exp_id = meta.get("experiment_id")
        if exp_id:
            exp = entities.get_experiment(db, exp_id)
            if exp:
                exp_name = exp.get("name")

        tag_names: list[str] = []
        tag_dicts = entities.get_tags_for_run(db, rh)
        for td in tag_dicts:
            n = td.get("name")
            if n:
                tag_names.append(n)

        index_run(
            db,
            rh,
            is_archived=meta.get("is_archived", False),
            active=meta.get("active", True),
            created_at=meta.get("created_at", 0.0),
            experiment_name=exp_name,
            tag_names=tag_names,
        )

        hparams = runs.get_run_attrs(db, rh, ("hparams",))
        if isinstance(hparams, dict):
            index_hparams(db, rh, hparams)

        count += 1

    return count
