from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .spec import StudySpec, TrialSpec


@dataclass(frozen=True)
class InfraJob:
    name: str
    gpus: int
    gpu_model: str | None
    hosts: list[str]
    command: str
    artifacts: list[str]
    env: dict[str, str]


def _toml_scalar(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, int):
        return str(value)
    text = str(value).replace("\\", "\\\\").replace('"', '\\"')
    return f'"{text}"'


def _toml_list(values: list[Any]) -> str:
    return "[" + ", ".join(_toml_scalar(item) for item in values) + "]"


def _toml_inline_table(mapping: dict[str, str]) -> str:
    parts = [f"{key} = {_toml_scalar(value)}" for key, value in sorted(mapping.items())]
    return "{ " + ", ".join(parts) + " }"


def build_trial_job(*, study: StudySpec, trial: TrialSpec, repo_root: Path) -> InfraJob:
    trial_path = Path(trial.artifact_dir) / "trial_spec.json"
    rel_trial_path = trial_path.resolve().relative_to(repo_root.resolve())
    rel_trial_dir = Path(trial.artifact_dir).resolve().relative_to(repo_root.resolve())
    command = (
        f"{study.engine.python_executable} -m aituner.cli worker run-trial "
        f"--trial-spec {rel_trial_path}"
    )
    env = {"PYTHONPATH": "src"}
    return InfraJob(
        name=f"{study.study_id}-{trial.trial_id}",
        gpus=study.hardware.gpu_count,
        gpu_model=study.hardware.gpu_model,
        hosts=list(study.hardware.host_candidates),
        command=command,
        artifacts=[str(rel_trial_dir)],
        env=env,
    )


def append_job(jobs_path: Path, job: InfraJob) -> None:
    jobs_path.parent.mkdir(parents=True, exist_ok=True)
    with jobs_path.open("a", encoding="utf-8") as handle:
        if jobs_path.stat().st_size == 0:
            handle.write("version = 1\n")
        handle.write("\n[[jobs]]\n")
        handle.write(f"name = {_toml_scalar(job.name)}\n")
        handle.write(f"gpus = {job.gpus}\n")
        if job.gpu_model:
            handle.write(f"gpu_model = {_toml_scalar(job.gpu_model)}\n")
        if job.hosts:
            handle.write(f"hosts = {_toml_list(job.hosts)}\n")
        handle.write(f"command = {_toml_scalar(job.command)}\n")
        if job.artifacts:
            handle.write(f"artifacts = {_toml_list(job.artifacts)}\n")
        if job.env:
            handle.write(f"env = {_toml_inline_table(job.env)}\n")
