from __future__ import annotations

from pathlib import Path

from persistence.backend import write_json
from schemas import build_artifact_layout_payload

ARTIFACT_STORAGE_LAYOUT = {
    "root_directory": "artifacts",
    "path_template": (
        "artifacts/{type_directory}/{artifact_key}/"
        "step-{step_label}__epoch-{epoch_label}__seq-{sequence_id:08d}.{extension}"
    ),
    "metadata_sidecar_suffix": ".meta.json",
    "type_directories": {
        "histogram": "histograms",
        "tensor_summary": "summaries",
        "embedding_snapshot": "embeddings",
        "batch_diagnostics_snapshot": "snapshots",
        "checkpoint_metadata": "checkpoints",
        "similarity_analysis_output": "analysis",
        "alert_evidence_bundle": "evidence",
    },
}


def ensure_artifact_layout(run_path: str | Path) -> Path:
    target = Path(run_path) / "artifacts" / "layout.json"
    if target.exists():
        return target
    payload = build_artifact_layout_payload(layout=ARTIFACT_STORAGE_LAYOUT)
    write_json(target, payload)
    return target
