"""Streaming search endpoints for runs and metrics.

All endpoints return ``StreamingResponse`` using the binary codec from
``api.streaming`` which the Aim UI understands.
"""

from __future__ import annotations

import asyncio
import struct as _struct
import time
from typing import TYPE_CHECKING

from fastapi import Header, HTTPException
from starlette.responses import StreamingResponse

from matyan_backend.api.streaming import (
    PROGRESS_REPORT_INTERVAL,
    collect_streamable_data,
    encode_tree,
    make_progress_key,
)
from matyan_backend.deps import FdbDb  # noqa: TC001
from matyan_backend.storage.runs import (
    get_all_contexts,
    get_run,
    get_run_attrs,
    get_run_traces_info,
)
from matyan_backend.storage.sequences import (
    read_and_sample_sequence,
    read_sequence,
    sample_sequence,
)

from ._collections import _NON_METRIC_DTYPES, iter_matching_runs, iter_matching_sequences
from ._pydantic_models import MetricAlignApiIn, RunTracesBatchApiIn  # noqa: TC001
from ._range_utils import parse_range
from ._run import rest_router_runs
from ._views import build_props_dict

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator

_SLEEP = 0.00001


def _numpy_to_encodable(values: list) -> dict:
    """Pack a list of numbers into the ``EncodedNumpyArray`` format the UI expects.

    Returns ``{"type": "numpy", "shape": N, "dtype": "float64", "blob": <bytes>}``
    where *blob* is little-endian float64 binary data.
    """
    floats = [float(v) for v in values]
    return {
        "type": "numpy",
        "shape": len(floats),
        "dtype": "float64",
        "blob": _struct.pack(f"<{len(floats)}d", *floats) if floats else b"",
    }


# ---------------------------------------------------------------------------
# Run search
# ---------------------------------------------------------------------------


@rest_router_runs.get("/search/run/")
async def run_search_api(  # noqa: C901
    db: FdbDb,
    q: str = "",
    limit: int = 0,
    offset: str = "",
    skip_system: bool = True,
    report_progress: bool = True,
    exclude_params: bool = False,
    exclude_traces: bool = False,
    x_timezone_offset: int = Header(default=0),
) -> StreamingResponse:
    async def _streamer() -> AsyncGenerator[bytes]:  # noqa: C901
        progress_idx = 0
        last_progress_time = time.time()
        count = 0
        past_offset = not offset

        for rv, (checked, total) in iter_matching_runs(db, q, tz_offset=x_timezone_offset):
            if not past_offset:
                if rv.hash == offset:
                    past_offset = True
                continue

            await asyncio.sleep(_SLEEP)

            if report_progress and time.time() - last_progress_time > PROGRESS_REPORT_INTERVAL:
                yield collect_streamable_data(encode_tree({make_progress_key(progress_idx): (checked, total)}))
                progress_idx += 1
                last_progress_time = time.time()

            meta = rv._meta  # noqa: SLF001
            meta["hash"] = rv.hash
            traces_info = get_run_traces_info(db, rv.hash)
            contexts = get_all_contexts(db, rv.hash)
            params = {} if exclude_params else (get_run_attrs(db, rv.hash) or {})
            if isinstance(params, dict):
                params.pop("__blobs__", None)
                if skip_system:
                    params.pop("__system_params", None)

            metric_traces: list[dict] = []
            if not exclude_traces:
                for t in traces_info:
                    dtype = t.get("dtype", "")
                    if dtype in _NON_METRIC_DTYPES:
                        continue
                    ctx_id = t.get("context_id", 0)
                    ctx = contexts.get(ctx_id, {})
                    last_val = t.get("last", 0.0)
                    metric_traces.append(
                        {
                            "name": t.get("name", ""),
                            "context": ctx,
                            "values": {
                                "last": last_val,
                                "last_step": t.get("last_step", 0),
                                "first": last_val,
                                "min": last_val,
                                "max": last_val,
                            },
                        },
                    )

            run_dict = {
                rv.hash: {
                    "params": params,
                    "traces": {"metric": metric_traces},
                    "props": build_props_dict(meta, db),
                },
            }
            yield collect_streamable_data(encode_tree(run_dict))

            if report_progress:
                yield collect_streamable_data(encode_tree({make_progress_key(progress_idx): (checked, total)}))
                progress_idx += 1
                last_progress_time = time.time()

            count += 1
            if limit and count >= limit:
                break

    return StreamingResponse(_streamer(), media_type="application/octet-stream")


# ---------------------------------------------------------------------------
# Metric search
# ---------------------------------------------------------------------------


@rest_router_runs.get("/search/metric/")
async def metric_search_api(
    db: FdbDb,
    q: str = "",
    p: int = 50,
    x_axis: str | None = None,
    skip_system: bool = True,
    report_progress: bool = True,
    x_timezone_offset: int = Header(default=0),
) -> StreamingResponse:
    async def _streamer() -> AsyncGenerator[bytes]:
        progress_idx = 0
        last_progress_time = time.time()

        seen_runs: dict[str, dict] = {}

        for rv, sv, (checked, total) in iter_matching_sequences(db, q, tz_offset=x_timezone_offset, seq_type="metric"):
            await asyncio.sleep(_SLEEP)

            if report_progress and time.time() - last_progress_time > PROGRESS_REPORT_INTERVAL:
                yield collect_streamable_data(encode_tree({make_progress_key(progress_idx): (checked, total)}))
                progress_idx += 1
                last_progress_time = time.time()

            run_hash = rv.hash

            if run_hash not in seen_runs:
                meta = rv._meta.copy()  # noqa: SLF001
                meta["hash"] = run_hash
                run_params = get_run_attrs(db, run_hash) or {}
                if isinstance(run_params, dict):
                    run_params.pop("__blobs__", None)
                    if skip_system:
                        run_params.pop("__system_params", None)
                seen_runs[run_hash] = {
                    "params": run_params,
                    "traces": [],
                    "props": build_props_dict(meta, db),
                }

            ctx_id = sv._trace_info.get("context_id", 0)  # noqa: SLF001
            sampled = sample_sequence(db, run_hash, ctx_id, sv.name, p)

            steps = sampled.get("steps", [])
            trace_view: dict = {
                "name": sv.name,
                "context": sv._context,  # noqa: SLF001
                "slice": [0, steps[-1] if steps else 0, 1],
                "values": _numpy_to_encodable(sampled.get("val", [])),
                "iters": _numpy_to_encodable(steps),
                "epochs": _numpy_to_encodable(sampled.get("epoch", [])),
                "timestamps": _numpy_to_encodable(sampled.get("time", [])),
            }

            if x_axis:
                x_seq = read_and_sample_sequence(
                    db,
                    run_hash,
                    ctx_id,
                    x_axis,
                    density=p,
                )
                trace_view["x_axis_values"] = _numpy_to_encodable(x_seq.get("val", []))
                trace_view["x_axis_iters"] = _numpy_to_encodable(x_seq.get("steps", []))

            seen_runs[run_hash]["traces"].append(trace_view)

        for run_hash, run_data in seen_runs.items():
            yield collect_streamable_data(encode_tree({run_hash: run_data}))

        if report_progress:
            yield collect_streamable_data(encode_tree({make_progress_key(progress_idx): (0, 0)}))

    return StreamingResponse(_streamer(), media_type="application/octet-stream")


# ---------------------------------------------------------------------------
# Active runs
# ---------------------------------------------------------------------------


@rest_router_runs.get("/active/")
async def active_runs_api(
    db: FdbDb,
    report_progress: bool = True,
) -> StreamingResponse:
    async def _streamer() -> AsyncGenerator[bytes]:
        progress_idx = 0

        for rv, (checked, total) in iter_matching_runs(db, ""):
            if not rv.active:
                continue

            await asyncio.sleep(_SLEEP)
            meta = rv._meta.copy()  # noqa: SLF001
            meta["hash"] = rv.hash
            traces_info = get_run_traces_info(db, rv.hash)
            contexts = get_all_contexts(db, rv.hash)

            metric_traces: list[dict] = []
            for t in traces_info:
                ctx_id = t.get("context_id", 0)
                ctx = contexts.get(ctx_id, {})
                metric_traces.append(
                    {
                        "name": t.get("name", ""),
                        "context": ctx,
                        "values": {
                            "last": t.get("last", 0.0),
                            "last_step": t.get("last_step", 0),
                        },
                    },
                )

            run_dict = {
                rv.hash: {
                    "traces": {"metric": metric_traces},
                    "props": build_props_dict(meta, db),
                },
            }
            yield collect_streamable_data(encode_tree(run_dict))

            if report_progress:
                yield collect_streamable_data(encode_tree({make_progress_key(progress_idx): (checked, total)}))
                progress_idx += 1

    return StreamingResponse(_streamer(), media_type="application/octet-stream")


# ---------------------------------------------------------------------------
# Metric alignment
# ---------------------------------------------------------------------------


@rest_router_runs.post("/search/metric/align/")
async def metric_custom_align_api(request_data: MetricAlignApiIn, db: FdbDb) -> StreamingResponse:  # noqa: C901
    async def _streamer() -> AsyncGenerator[bytes]:
        for run_data in request_data.runs:
            await asyncio.sleep(_SLEEP)
            run_hash = run_data.run_id
            run = get_run(db, run_hash)
            if not run:
                continue

            traces_info = get_run_traces_info(db, run_hash)
            contexts = get_all_contexts(db, run_hash)

            traces_list = []
            for trace_req in run_data.traces:
                ctx_id: int | None = None
                for t in traces_info:
                    c = contexts.get(t.get("context_id", 0), {})
                    if t.get("name") == trace_req.name and c == trace_req.context:
                        ctx_id = t.get("context_id", 0)
                        break

                if ctx_id is None:
                    continue

                x_ctx_id: int | None = None
                for t in traces_info:
                    c = contexts.get(t.get("context_id", 0), {})
                    if t.get("name") == request_data.align_by and c == trace_req.context:
                        x_ctx_id = t.get("context_id", 0)
                        break

                if x_ctx_id is None:
                    continue

                max_step = trace_req.slice[1] if trace_req.slice[1] > 0 else None
                x_seq = read_sequence(db, run_hash, x_ctx_id, request_data.align_by, end_step=max_step)

                traces_list.append(
                    {
                        "name": trace_req.name,
                        "context": trace_req.context,
                        "x_axis_values": _numpy_to_encodable(x_seq.get("val", [])),
                        "x_axis_iters": _numpy_to_encodable(x_seq.get("steps", [])),
                    },
                )

            yield collect_streamable_data(encode_tree({run_hash: traces_list}))

    return StreamingResponse(_streamer(), media_type="application/octet-stream")


# ---------------------------------------------------------------------------
# Metric batch
# ---------------------------------------------------------------------------


@rest_router_runs.post("/{run_id}/metric/get-batch/")
async def run_metric_batch_api(
    run_id: str,
    body: RunTracesBatchApiIn,
    db: FdbDb,
    record_range: str = "",
    record_density: int = 50,
) -> list[dict]:
    run = get_run(db, run_id)
    if not run:
        raise HTTPException(status_code=404)

    rr = parse_range(record_range)
    effective_end = rr.stop - 1 if rr.stop is not None else None
    traces_info = get_run_traces_info(db, run_id)
    contexts = get_all_contexts(db, run_id)

    results = []
    for req in body:
        ctx_id: int | None = None
        for t in traces_info:
            c = contexts.get(t.get("context_id", 0), {})
            if t.get("name") == req.name and c == req.context:
                ctx_id = t.get("context_id", 0)
                break
        if ctx_id is None:
            continue

        seq = read_and_sample_sequence(
            db,
            run_id,
            ctx_id,
            req.name,
            start_step=rr.start,
            end_step=effective_end,
            density=record_density,
        )
        results.append(
            {
                "name": req.name,
                "context": req.context,
                "iters": seq.get("steps", []),
                "values": seq.get("val", []),
            },
        )

    return results
