from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from one_run.run_id import make_run_id
from one_run.constants import (
    ARTIFACTS_DIR_NAME,
    CONFIG_SNAPSHOT_NAME,
    LOGS_DIR_NAME,
    METADATA_NAME,
    METRICS_DIR_NAME,
    LOCAL_EXIT_CODE_NAME,
    LOCAL_PID_NAME,
    RUNTIME_NAME,
    SUMMARY_NAME,
    SUMMARY_MIN_NAME,
    EVENTS_NAME,
    STDERR_LOG_NAME,
    STATUS_NAME,
    STDOUT_LOG_NAME,
)


@dataclass(frozen=True)
class RunPaths:
    run_dir: Path
    config_snapshot: Path
    metadata: Path
    status: Path
    runtime: Path
    summary: Path
    summary_min: Path
    events: Path
    local_pid: Path
    local_exit_code: Path
    logs_dir: Path
    stdout_log: Path
    stderr_log: Path
    artifacts_dir: Path
    metrics_dir: Path


def resolve_run_paths(run_dir: Path) -> RunPaths:
    logs = run_dir / LOGS_DIR_NAME
    artifacts = run_dir / ARTIFACTS_DIR_NAME
    metrics = run_dir / METRICS_DIR_NAME

    return RunPaths(
        run_dir=run_dir,
        config_snapshot=run_dir / CONFIG_SNAPSHOT_NAME,
        metadata=run_dir / METADATA_NAME,
        status=run_dir / STATUS_NAME,
        runtime=run_dir / RUNTIME_NAME,
        summary=run_dir / SUMMARY_NAME,
        summary_min=run_dir / SUMMARY_MIN_NAME,
        events=run_dir / EVENTS_NAME,
        local_pid=run_dir / LOCAL_PID_NAME,
        local_exit_code=run_dir / LOCAL_EXIT_CODE_NAME,
        logs_dir=logs,
        stdout_log=logs / STDOUT_LOG_NAME,
        stderr_log=logs / STDERR_LOG_NAME,
        artifacts_dir=artifacts,
        metrics_dir=metrics,
    )


def create_run_paths(run_dir: Path) -> RunPaths:
    paths = resolve_run_paths(run_dir)

    run_dir.mkdir(parents=True, exist_ok=False)
    paths.logs_dir.mkdir(parents=True, exist_ok=True)
    paths.artifacts_dir.mkdir(parents=True, exist_ok=True)
    paths.metrics_dir.mkdir(parents=True, exist_ok=True)
    return paths


def allocate_run_paths(runs_root: Path, experiment_name: str, *, max_attempts: int = 100) -> tuple[str, RunPaths]:
    base_run_id = make_run_id(experiment_name)
    for attempt in range(max_attempts):
        run_id = base_run_id if attempt == 0 else f"{base_run_id}-{attempt:02d}"
        try:
            return run_id, create_run_paths(runs_root / run_id)
        except FileExistsError:
            continue
    raise FileExistsError(f"failed to allocate unique run_id after {max_attempts} attempts")
