from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Generic, Mapping, Sequence

from simple_inject import provide as inject_provide

from zyplib.experiment._types import (
    ParamsT,
    RecordConfig,
    TrialMode,
    matches_export_rule,
    serialize_value,
)
from zyplib.utils.fs import write_yaml


def _merge_records(
    store: dict[str, dict[str, Any]], namespace: str, key: str, value: Any
) -> None:
    store.setdefault(namespace, {})[key] = value


def _clone_records(records: Mapping[str, Mapping[str, Any]]) -> dict[str, dict[str, Any]]:
    return {namespace: dict(values) for namespace, values in records.items()}


@dataclass(slots=True)
class ExperimentContext(Generic[ParamsT]):
    """run 级上下文。

    `prepare_run()` 会在这个对象上工作，用来准备跨 trial 共享的数据，
    例如文件清单、固定 fold plan、预扫描索引等。

    Parameters
    ----------
    - `name`: `str`, experiment 名称
    - `params`: `ParamsT`, 当前运行参数
    - `run_dir`: `Path`, 当前 run 的输出目录
    - `record`: `RecordConfig | None`, record 行为配置
    - `shared`: `Any`, run 级共享对象, by default `None`

    Examples
    ----------
    ```python
    def prepare_run(ctx):
        ctx.set('subject_count', 12, namespace='result')
        return {'all_files': files}
    ```
    """

    name: str
    params: ParamsT
    run_dir: Path
    record: RecordConfig | None = None
    shared: Any = None
    _records: dict[str, dict[str, Any]] = field(
        default_factory=dict, init=False, repr=False
    )

    def set(self, key: str, value: Any, namespace: str = 'default') -> None:
        """写入 run 级 namespace 状态。"""

        _merge_records(self._records, namespace, key, value)

    def recorded_states(self) -> dict[str, dict[str, Any]]:
        """返回 run 级已记录状态。"""

        return _clone_records(self._records)

    def write_yaml(self, relative_path: str, data: Mapping[str, Any]) -> Path:
        """把 run 级数据写到当前 experiment 目录下。"""

        path = self.run_dir / relative_path
        write_yaml(path, serialize_value(data))
        return path


@dataclass(slots=True)
class TrialContext(Generic[ParamsT]):
    """trial 级上下文。

    trial 的数据准备、模型构建和训练逻辑通常都写在这里的消费链条上。
    它暴露了参数、trial 元信息、shared 数据与记录接口。

    Parameters
    ----------
    - `experiment`: `ExperimentContext[ParamsT]`, 所属 run 的上下文
    - `trial_name`: `str`, 当前 trial 名称
    - `trial_index`: `int`, 0-based trial 索引
    - `trial_mode`: `TrialMode`, 当前 trial 模式
    - `trial_seed`: `int | None`, 当前 trial 使用的随机种子
    - `trial_dir`: `Path`, 当前 trial 的输出目录

    Examples
    ----------
    ```python
    def build_data(ctx):
        print(ctx.trial_name, ctx.trial_seed)
        return data
    ```
    """

    experiment: ExperimentContext[ParamsT]
    trial_name: str
    trial_index: int
    trial_mode: TrialMode
    trial_seed: int | None
    trial_dir: Path
    _records: dict[str, dict[str, Any]] = field(
        default_factory=dict, init=False, repr=False
    )

    @property
    def params(self) -> ParamsT:
        """返回当前 trial 共享的参数对象。"""

        return self.experiment.params

    @property
    def run_dir(self) -> Path:
        """返回所属 experiment 的 run 目录。"""

        return self.experiment.run_dir

    @property
    def shared(self) -> Any:
        """返回 `prepare_run()` 生成的 run 级 shared 数据。"""

        return self.experiment.shared

    @property
    def record(self) -> RecordConfig | None:
        """返回当前 experiment 的 record 配置。"""

        return self.experiment.record

    @property
    def trial_no(self) -> int:
        """返回从 1 开始计数的 trial 编号。"""

        return self.trial_index + 1

    def set(self, key: str, value: Any, namespace: str = 'default') -> None:
        """写入 trial 级 namespace 状态，并按配置决定是否注入当前 scope。"""

        _merge_records(self._records, namespace, key, value)
        if (
            self.record
            and self.record.enabled
            and namespace in self.record.inject_namespaces
        ):
            inject_provide(key, value, namespace=namespace)

    def recorded_states(
        self, namespaces: Sequence[str] | None = None
    ) -> dict[str, dict[str, Any]]:
        """返回 trial 本地已记录状态。"""

        records = (
            self._records
            if namespaces is None
            else {namespace: self._records.get(namespace, {}) for namespace in namespaces}
        )
        return _clone_records(records)

    def snapshot_states(
        self, namespaces: Sequence[str] | None = None
    ) -> dict[str, dict[str, Any]]:
        """按 namespace 抓取当前 trial 的状态快照。"""

        if namespaces is None and self.record and self.record.enabled:
            filtered: dict[str, dict[str, Any]] = {}
            for namespace, values in self._records.items():
                exported = {
                    key: value
                    for key, value in values.items()
                    if matches_export_rule(self.record.export_rule, namespace, key, value)
                }
                if exported:
                    filtered[namespace] = exported
            return _clone_records(filtered)
        return self.recorded_states(namespaces)

    def write_yaml(self, relative_path: str, data: Mapping[str, Any]) -> Path:
        """把 trial 级数据写到当前 trial 目录下。"""

        path = self.trial_dir / relative_path
        write_yaml(path, serialize_value(data))
        return path


__all__ = ['ExperimentContext', 'TrialContext']
