"""Run/sequence collection iterators backed by FDB.

Provides query-filtered iteration over runs, yielding per-run sequence
collections together with progress tuples for streaming progress reporting.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from matyan_backend.storage.runs import (
    get_all_contexts,
    get_run_meta,
    get_run_traces_info,
    list_run_hashes,
)

from ._planner import plan_query
from ._query import RestrictedPythonQuery
from ._views import RunView, SequenceView

if TYPE_CHECKING:
    from collections.abc import Iterator

    from matyan_backend.fdb_types import Database


def iter_matching_runs(
    db: Database,
    query: str = "",
    *,
    tz_offset: int = 0,
) -> Iterator[tuple[RunView, tuple[int, int]]]:
    """Iterate all runs, yielding only those matching the AimQL *query*.

    Yields ``(run_view, (counter, total))`` tuples so callers can report
    progress.
    """
    candidates = plan_query(db, query)
    hashes = candidates if candidates is not None else list_run_hashes(db)

    q = RestrictedPythonQuery(query)
    total = len(hashes)

    for idx, h in enumerate(hashes, start=1):
        meta = get_run_meta(db, h)
        if not meta:
            continue
        meta["hash"] = h

        rv = RunView(db, h, meta, tz_offset=tz_offset)
        match = q.check(run=rv)
        if match:
            yield rv, (idx, total)


_SEQ_TYPE_TO_DTYPES: dict[str, frozenset[str]] = {
    "images": frozenset({"image"}),
    "texts": frozenset({"text"}),
    "distributions": frozenset({"distribution"}),
    "audios": frozenset({"audio"}),
    "figures": frozenset({"figure"}),
}

_NON_METRIC_DTYPES = frozenset().union(
    *_SEQ_TYPE_TO_DTYPES.values(),
    {"logs", "log_records"},
)


def iter_matching_sequences(
    db: Database,
    query: str = "",
    *,
    tz_offset: int = 0,
    seq_type: str = "metric",
) -> Iterator[tuple[RunView, SequenceView, tuple[int, int]]]:
    """Iterate all runs and their traces, yielding matching run+sequence pairs.

    Yields ``(run_view, seq_view, (counter, total))`` for every matching
    sequence across all runs.
    """
    allowed_dtypes = _SEQ_TYPE_TO_DTYPES.get(seq_type)

    candidates = plan_query(db, query)
    hashes = candidates if candidates is not None else list_run_hashes(db)

    q = RestrictedPythonQuery(query)
    total = len(hashes)

    for idx, h in enumerate(hashes, start=1):
        meta = get_run_meta(db, h)
        if not meta:
            continue
        meta["hash"] = h

        rv = RunView(db, h, meta, tz_offset=tz_offset)
        traces = get_run_traces_info(db, h)
        contexts = get_all_contexts(db, h)

        for t in traces:
            dtype = t.get("dtype", "")
            if allowed_dtypes is not None:
                if dtype not in allowed_dtypes:
                    continue
            elif dtype in _NON_METRIC_DTYPES:
                continue

            ctx_id = t.get("context_id", 0)
            ctx = contexts.get(ctx_id, {})
            sv = SequenceView(t.get("name", ""), ctx, rv, trace_info=t)

            match = q.check(run=rv, **{seq_type: sv})
            if match:
                yield rv, sv, (idx, total)
