"""Project-level aggregation queries.

These functions scan FDB to produce summary statistics consumed by the
``/projects/*`` API endpoints.  Project settings (name, pinned sequences, etc.)
are stored under the ``system`` directory.
"""

from __future__ import annotations

import datetime
from collections import defaultdict
from typing import TYPE_CHECKING, Any

from matyan_backend.fdb_types import transactional

from .encoding import decode_value, encode_value
from .entities import list_experiments
from .fdb_client import get_directories
from .runs import get_context, get_run_attrs, get_run_traces_info, list_run_hashes, list_runs_meta
from .tree import _LEAF_SENTINEL

if TYPE_CHECKING:
    from matyan_backend.fdb_types import DirectorySubspace, Transaction

_PROJECT_PREFIX = "project"


def _sys_dir() -> DirectorySubspace:
    return get_directories().system


@transactional
def get_project_info(tr: Transaction) -> dict:
    sd = _sys_dir()
    r = sd.range((_PROJECT_PREFIX,))
    fields: dict[str, Any] = {}
    for kv in tr.get_range(r.start, r.stop):
        key_tuple = sd.unpack(kv.key)
        if len(key_tuple) >= 2:
            fields[key_tuple[1]] = decode_value(kv.value)
    return {
        "name": fields.get("name", "My project"),
        "path": "",
        "description": fields.get("description", ""),
        "telemetry_enabled": 0,
    }


@transactional
def set_project_info(tr: Transaction, **fields: Any) -> None:  # noqa: ANN401
    sd = _sys_dir()
    for k, v in fields.items():
        tr[sd.pack((_PROJECT_PREFIX, k, _LEAF_SENTINEL))] = encode_value(v)


@transactional
def get_project_activity(tr: Transaction, tz_offset: int = 0) -> dict:
    """Compute aggregate activity stats across all runs."""
    runs = list_runs_meta(tr)
    experiments = list_experiments(tr)

    num_runs = len(runs)
    num_archived = sum(1 for r in runs if r.get("is_archived"))
    num_active = sum(1 for r in runs if r.get("active"))

    activity_map: dict[str, int] = defaultdict(int)
    for r in runs:
        created = r.get("created_at")
        if created:
            ts = created - tz_offset * 60
            day = datetime.datetime.fromtimestamp(ts, tz=datetime.UTC).strftime("%Y-%m-%dT00:00:00")
            activity_map[day] += 1

    return {
        "num_experiments": len(experiments),
        "num_runs": num_runs,
        "num_archived_runs": num_archived,
        "num_active_runs": num_active,
        "activity_map": dict(activity_map),
    }


_INTERNAL_ATTR_KEYS = frozenset({"__system_params", "__blobs__"})
_NON_METRIC_DTYPES = frozenset({
    "image", "text", "distribution", "audio", "figure",
    "logs", "log_records",
})


@transactional
def get_project_params(tr: Transaction, sequence_types: tuple[str, ...] = ()) -> dict:
    """Aggregate distinct hparams and trace names across all runs.

    Returns a dict keyed by sequence type (``metric``, ``images``, etc.) with
    nested structure representing the parameter/trace name tree.
    """
    hashes = list_run_hashes(tr)

    dtype_to_bucket = {
        "image": "images",
        "text": "texts",
        "distribution": "distributions",
        "audio": "audios",
        "figure": "figures",
    }

    all_params: dict[str, Any] = {}
    buckets: dict[str, dict[str, list[dict]]] = {
        "metric": {},
        "images": {},
        "texts": {},
        "figures": {},
        "distributions": {},
        "audios": {},
    }

    wanted = set(sequence_types) if sequence_types else set(buckets)

    for h in hashes:
        attrs = get_run_attrs(tr, h)
        if isinstance(attrs, dict):
            _merge_dict(all_params, attrs)

        traces = get_run_traces_info(tr, h)
        for t in traces:
            name = t.get("name", "")
            ctx_id = t.get("context_id", 0)
            dtype = t.get("dtype", "")
            if dtype in _NON_METRIC_DTYPES and dtype not in dtype_to_bucket:
                continue
            bucket = dtype_to_bucket.get(dtype, "metric")
            if bucket not in wanted:
                continue
            ctx = get_context(tr, h, ctx_id) or {}
            if name not in buckets[bucket]:
                buckets[bucket][name] = []
            if ctx not in buckets[bucket][name]:
                buckets[bucket][name].append(ctx)

    for key in _INTERNAL_ATTR_KEYS:
        all_params.pop(key, None)

    return {"params": all_params, **buckets}


def _merge_dict(target: dict, source: dict) -> None:
    for k, v in source.items():
        if isinstance(v, dict) and isinstance(target.get(k), dict):
            _merge_dict(target[k], v)
        else:
            target[k] = v


# ---------------------------------------------------------------------------
# Pinned sequences
# ---------------------------------------------------------------------------

_PINNED = "pinned_sequences"


@transactional
def get_pinned_sequences(tr: Transaction) -> list:
    sd = _sys_dir()
    raw = tr[sd.pack((_PINNED, _LEAF_SENTINEL))]
    if not raw.present():
        return []
    return decode_value(raw)


@transactional
def set_pinned_sequences(tr: Transaction, sequences: list) -> list:
    sd = _sys_dir()
    tr[sd.pack((_PINNED, _LEAF_SENTINEL))] = encode_value(sequences)
    return sequences
