from __future__ import annotations

from dataclasses import asdict, is_dataclass
from numbers import Real
from pathlib import Path
from typing import TYPE_CHECKING, Any, Mapping, Sequence, cast

import numpy as np
from simple_inject import create_scope
from simple_inject import provide as inject_provide

from zyplib.experiment._types import (
    AutoRecordConfig,
    BuildDataFn,
    BuildModelFn,
    ExperimentData,
    ExperimentEffects,
    ExperimentReport,
    MetricsLike,
    ModelSetup,
    ParamsT,
    PrepareRunFn,
    RecordConfig,
    ResolvePlanFn,
    ResultAggregator,
    TrialMode,
    TrialPlan,
    TrialReport,
    TrialResult,
    TrainModelFn,
    matches_export_rule,
    normalize_metrics,
    serialize_value,
)
from zyplib.experiment.context import ExperimentContext, TrialContext
from zyplib.utils import TIME_FMT, now
from zyplib.utils.envs import env_status
from zyplib.utils.fs import write_yaml
from zyplib.utils.kv import aggregate_dicts
from zyplib.utils.print import PrintCapture

if TYPE_CHECKING:
    from torch.nn import Module
else:
    Module = Any


def _resolve_record_config(record: RecordConfig | bool) -> RecordConfig | None:
    if record is False:
        return None
    if record is True:
        return RecordConfig()
    return record


def _resolve_auto_record_config(
    config: AutoRecordConfig | bool,
) -> AutoRecordConfig | None:
    if config is False:
        return None
    if config is True:
        return AutoRecordConfig()
    return config


def _resolve_effects(effects: ExperimentEffects | None) -> ExperimentEffects:
    return effects or ExperimentEffects()


def _resolve_run_version(params: Any) -> str:
    for attr_name in ('version', 'VERSION', 'name', 'NAME'):
        value = getattr(params, attr_name, None)
        if value:
            return str(value)
    return now(TIME_FMT.DATE_TIME_UNDERSCORE)


def _resolve_run_dir(root_dir: str | Path, params: Any) -> Path:
    return Path(root_dir) / _resolve_run_version(params)


def _resolve_trial_plan(
    exp_ctx: ExperimentContext[ParamsT], plan: TrialPlan | ResolvePlanFn[ParamsT] | None
) -> TrialPlan:
    if callable(plan):
        return plan(exp_ctx)
    if plan is not None:
        return plan
    return TrialPlan.single()


def _trial_seeds(plan: TrialPlan) -> list[int | None]:
    if plan.trial_seeds is not None:
        return list(plan.trial_seeds)
    if plan.count == 1 and plan.mode == 'single':
        return [plan.base_seed]
    if plan.base_seed is None:
        return [None] * plan.count
    generator = np.random.default_rng(plan.base_seed)
    return [int(seed) for seed in generator.integers(0, 2**31 - 1, size=plan.count)]


def _trial_name(mode: TrialMode, index: int, count: int) -> str:
    if count == 1 and mode == 'single':
        return 'single'
    return f'{mode}-{index + 1:03d}'


def _default_aggregate(trials: Sequence[TrialReport]) -> dict[str, Any]:
    if not trials:
        return {}

    if len(trials) == 1:
        return {
            # key: {'mean': float(value), 'std': 0.0}
            key: value
            for key, value in trials[0].result.metrics.items()
        }

    metrics = [trial.result.metrics for trial in trials]

    def aggregate_metric(_: str, values: Sequence[Any]) -> dict[str, float]:
        if any(not isinstance(item, Real) for item in values):
            raise TypeError(
                'Default experiment aggregation only supports numeric metrics'
            )
        numeric_values = [float(item) for item in values]
        return {
            'mean': float(np.mean(numeric_values)),
            'std': float(np.std(numeric_values)),
        }

    return aggregate_dicts(aggregate_metric, metrics)  # pyright: ignore[reportArgumentType]


def _resolve_aggregate(aggregate: ResultAggregator | str) -> ResultAggregator:
    if callable(aggregate):
        return aggregate
    if aggregate != 'meanstd':
        raise ValueError(f'Unsupported aggregate mode: {aggregate}')
    return _default_aggregate


def _filter_scoped_state(
    scoped_state: Mapping[str, Mapping[str, Any]], record_config: RecordConfig | None
) -> dict[str, dict[str, Any]]:
    auto_keys = set()
    if record_config is not None:
        auto_keys = {
            record_config.inject_ctx_key,
            record_config.inject_params_key,
            record_config.inject_trial_key,
        }

    filtered: dict[str, dict[str, Any]] = {}
    for namespace, values in scoped_state.items():
        namespace_values = {
            key: value for key, value in values.items() if key not in auto_keys
        }
        if namespace_values:
            filtered[namespace] = namespace_values
    return filtered


def _should_export_record(
    record_config: RecordConfig, namespace: str, key: str, value: Any
) -> bool:
    return matches_export_rule(record_config.export_rule, namespace, key, value)


def _capture_trial_records(
    trial_ctx: TrialContext[ParamsT], scoped_state: Mapping[str, Mapping[str, Any]] | None
) -> dict[str, dict[str, Any]]:
    record_config = trial_ctx.record
    if record_config is None or not record_config.enabled:
        return {}

    local_records = trial_ctx.snapshot_states()
    inject_records = _filter_scoped_state(scoped_state or {}, record_config)
    namespaces = set(local_records) | set(inject_records)

    merged_records: dict[str, dict[str, Any]] = {}
    for namespace in namespaces:
        namespace_records: dict[str, Any] = {}
        namespace_records.update(inject_records.get(namespace, {}))
        namespace_records.update(local_records.get(namespace, {}))
        exported_records = {
            key: value
            for key, value in namespace_records.items()
            if _should_export_record(record_config, namespace, key, value)
        }
        if exported_records:
            merged_records[namespace] = serialize_value(exported_records)
    return merged_records


def _normalize_model_setup(value: ModelSetup | Module) -> ModelSetup:
    if isinstance(value, ModelSetup):
        return value
    return ModelSetup(model=value)


def _normalize_trial_result(value: TrialResult | Mapping[str, Any]) -> TrialResult:
    if isinstance(value, TrialResult):
        return TrialResult(
            metrics=normalize_metrics(value.metrics),
            summary=value.summary,
            artifacts={str(key): str(item) for key, item in value.artifacts.items()},
        )
    return TrialResult(metrics=normalize_metrics(value))


def _auto_record_setups(
    ctx: TrialContext[ParamsT], data: ExperimentData, model_setup: ModelSetup
) -> None:
    if ctx.record is None or not ctx.record.enabled:
        return

    auto_record = _resolve_auto_record_config(ctx.record.auto_record)
    if auto_record is None:
        return

    if auto_record.data_info and data.info:
        ctx.set('data_info', data.info, namespace='result')
    if auto_record.model_hparams and model_setup.hparams:
        ctx.set('model_hparam', model_setup.hparams)
    if auto_record.optimizer and model_setup.optimizer is not None:
        ctx.set('optimizer', model_setup.optimizer, namespace='train')
    if auto_record.lr_scheduler and model_setup.lr_scheduler is not None:
        ctx.set('lr_scheduler', model_setup.lr_scheduler, namespace='train')
    if auto_record.trial_meta:
        ctx.set('trial_name', ctx.trial_name)
        ctx.set('trial_mode', ctx.trial_mode)
        if ctx.trial_seed is not None:
            ctx.set('trial_seed', ctx.trial_seed)


def build_default_trial_fabric(ctx: TrialContext[Any]):
    """为单个 trial 构造默认 Fabric。"""

    from zyplib.train import default_fabric

    return default_fabric(
        root_dir=str(ctx.run_dir), dir_name='train', version=ctx.trial_name, keep_best_k=3
    )


def _write_run_metadata(
    exp_ctx: ExperimentContext[ParamsT], env: dict[str, Any], record_env: bool
) -> None:
    params_payload = exp_ctx.params
    if is_dataclass(params_payload) and not isinstance(params_payload, type):
        params_payload = asdict(params_payload)
    write_yaml(
        exp_ctx.run_dir / 'params.yaml',
        {'name': exp_ctx.name, 'params': serialize_value(params_payload)},
    )
    if record_env:
        write_yaml(exp_ctx.run_dir / 'env.yaml', serialize_value(env))


def _execute_experiment(
    *,
    name: str,
    params: ParamsT,
    run_dir: Path,
    prepare_run: PrepareRunFn[ParamsT] | None,
    plan: TrialPlan | ResolvePlanFn[ParamsT] | None,
    build_data: BuildDataFn[ParamsT],
    build_model: BuildModelFn[ParamsT],
    train_model: TrainModelFn[ParamsT],
    aggregate: ResultAggregator | str,
    effects: ExperimentEffects | None,
) -> ExperimentReport[ParamsT]:
    resolved_effects = _resolve_effects(effects)
    record_config = _resolve_record_config(resolved_effects.record)
    run_dir.mkdir(parents=True, exist_ok=True)

    env = env_status() if resolved_effects.record_env else {}
    exp_ctx = ExperimentContext(
        name=name, params=params, run_dir=run_dir, record=record_config
    )
    _write_run_metadata(exp_ctx, env, resolved_effects.record_env)

    if prepare_run is not None:
        exp_ctx.shared = prepare_run(exp_ctx)

    resolved_plan = _resolve_trial_plan(exp_ctx, plan)
    seeds = _trial_seeds(resolved_plan)

    trials: list[TrialReport] = []
    for trial_index, trial_seed in enumerate(seeds):
        trial_name = _trial_name(resolved_plan.mode, trial_index, resolved_plan.count)
        trial_dir = run_dir / 'trials' / trial_name
        trial_dir.mkdir(parents=True, exist_ok=True)
        trial_ctx = TrialContext(
            experiment=exp_ctx,
            trial_name=trial_name,
            trial_index=trial_index,
            trial_mode=resolved_plan.mode,
            trial_seed=trial_seed,
            trial_dir=trial_dir,
        )

        if record_config and record_config.enabled:
            with create_scope() as scope:
                inject_provide(record_config.inject_ctx_key, trial_ctx)
                inject_provide(record_config.inject_params_key, trial_ctx.params)
                inject_provide(
                    record_config.inject_trial_key,
                    {
                        'trial_name': trial_ctx.trial_name,
                        'trial_index': trial_ctx.trial_index,
                        'trial_no': trial_ctx.trial_no,
                        'trial_mode': trial_ctx.trial_mode,
                        'trial_seed': trial_ctx.trial_seed,
                    },
                )
                data = build_data(trial_ctx)
                model_setup = _normalize_model_setup(build_model(trial_ctx))
                _auto_record_setups(trial_ctx, data, model_setup)
                raw_result = train_model(trial_ctx, model_setup, data)
            scoped_state = scope.scoped_state()
        else:
            data = build_data(trial_ctx)
            model_setup = _normalize_model_setup(build_model(trial_ctx))
            _auto_record_setups(trial_ctx, data, model_setup)
            raw_result = train_model(trial_ctx, model_setup, data)
            scoped_state = {}

        trial_result = _normalize_trial_result(raw_result)
        trial_report = TrialReport(
            trial_name=trial_ctx.trial_name,
            trial_index=trial_ctx.trial_index,
            trial_mode=trial_ctx.trial_mode,
            trial_seed=trial_ctx.trial_seed,
            result=trial_result,
            records=_capture_trial_records(trial_ctx, scoped_state),
        )
        write_yaml(trial_dir / 'report.yaml', trial_report.to_dict())
        trials.append(trial_report)

    summary = _resolve_aggregate(aggregate)(trials)
    report = ExperimentReport(
        name=name, params=params, run_dir=run_dir, summary=summary, trials=trials, env=env
    )
    write_yaml(
        run_dir / 'report.yaml',
        report.to_dict(
            include_trials=resolved_plan.count == 1
        ),  # 默认不包含 trails 级别 (因为已经在各自的 report.yaml 里了)
    )

    run_records = exp_ctx.recorded_states()
    if run_records:
        write_yaml(run_dir / 'run-state.yaml', run_records)

    return report


def run_experiment(
    *,
    name: str,
    params: ParamsT,
    root_dir: str,
    build_data: BuildDataFn[ParamsT],
    build_model: BuildModelFn[ParamsT],
    train_model: TrainModelFn[ParamsT],
    prepare_run: PrepareRunFn[ParamsT] | None = None,
    plan: TrialPlan | ResolvePlanFn[ParamsT] | None = None,
    aggregate: ResultAggregator | str = 'meanstd',
    effects: ExperimentEffects | None = None,
) -> ExperimentReport[ParamsT]:
    """以函数式风格执行一个 experiment。"""

    resolved_effects = _resolve_effects(effects)
    run_dir = _resolve_run_dir(root_dir, params)
    run_dir.mkdir(parents=True, exist_ok=True)

    if not resolved_effects.capture_stdout:
        return _execute_experiment(
            name=name,
            params=params,
            run_dir=run_dir,
            prepare_run=prepare_run,
            plan=plan,
            build_data=build_data,
            build_model=build_model,
            train_model=train_model,
            aggregate=aggregate,
            effects=resolved_effects,
        )

    stdout_path = run_dir / 'stdout-log.txt'
    with PrintCapture(str(stdout_path), flush_interval=50):
        return _execute_experiment(
            name=name,
            params=params,
            run_dir=run_dir,
            prepare_run=prepare_run,
            plan=plan,
            build_data=build_data,
            build_model=build_model,
            train_model=train_model,
            aggregate=aggregate,
            effects=resolved_effects,
        )


def run_experiment_cli(
    *,
    name: str,
    params_type: type[ParamsT],
    root_dir: str,
    build_data: BuildDataFn[ParamsT],
    build_model: BuildModelFn[ParamsT],
    train_model: TrainModelFn[ParamsT],
    prepare_run: PrepareRunFn[ParamsT] | None = None,
    plan: TrialPlan | ResolvePlanFn[ParamsT] | None = None,
    aggregate: ResultAggregator | str = 'meanstd',
    effects: ExperimentEffects | None = None,
    args: Sequence[str] | None = None,
) -> ExperimentReport[ParamsT]:
    """用 `tyro` 解析参数后执行函数式 experiment。"""

    import tyro  # pyright: ignore[reportMissingImports]

    params = cast(
        ParamsT, tyro.cli(params_type, args=list(args) if args is not None else None)
    )
    return run_experiment(
        name=name,
        params=params,
        root_dir=root_dir,
        build_data=build_data,
        build_model=build_model,
        train_model=train_model,
        prepare_run=prepare_run,
        plan=plan,
        aggregate=aggregate,
        effects=effects,
    )


__all__ = [
    'BuildModelFn',
    'BuildDataFn',
    'PrepareRunFn',
    'ResolvePlanFn',
    'TrainModelFn',
    'build_default_trial_fabric',
    'run_experiment',
    'run_experiment_cli',
]
