from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any, Iterable


DEFAULT_REQUIRED = ("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json")


def verify_artifacts(artifact_dir: str | Path, required: Iterable[str] = DEFAULT_REQUIRED) -> dict[str, Any]:
    root = Path(artifact_dir)
    if not root.exists():
        raise FileNotFoundError(f"Artifact directory not found: {root}")
    if not root.is_dir():
        raise ValueError(f"Artifact path is not a directory: {root}")

    missing = [name for name in required if not (root / name).is_file()]
    if missing:
        raise ValueError(f"Artifact directory missing required files: {', '.join(missing)}")

    _validate_json(root / "final_metrics.json")
    _validate_json(root / "run_manifest.json")
    _validate_jsonl(root / "metrics.jsonl")

    files = sorted(path for path in root.rglob("*") if path.is_file())
    manifest = {
        "artifact_dir": str(root),
        "file_count": len(files),
        "files": [
            {
                "path": str(path.relative_to(root)),
                "bytes": path.stat().st_size,
                "sha256": sha256_file(path),
            }
            for path in files
        ],
    }
    (root / "artifact_manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    (root / "checksums.txt").write_text(
        "".join(f"{item['sha256']}  {item['path']}\n" for item in manifest["files"])
    )
    return manifest


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _validate_json(path: Path) -> None:
    try:
        json.loads(path.read_text())
    except json.JSONDecodeError as exc:
        raise ValueError(f"Invalid JSON artifact {path}: {exc}") from exc


def _validate_jsonl(path: Path) -> None:
    with path.open() as file:
        for line_number, line in enumerate(file, start=1):
            stripped = line.strip()
            if not stripped:
                continue
            try:
                json.loads(stripped)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSONL artifact {path}:{line_number}: {exc}") from exc