"""
Staging layer for mlvault-mlflow.

Artifacts are copied here instantly on log_artifact (like `git add`).
`mlvault commit` reads from here, bundles, encrypts, and uploads to Jackal.

Layout:
  ~/.mlvault/mlflow-staging/{mlflow_run_id}/
    artifacts/          -- mirrored artifact tree
    meta.json           -- run metadata written on first artifact
"""

import json
import shutil
from datetime import datetime, timezone
from pathlib import Path

_STAGING_ROOT = Path.home() / ".mlvault" / "mlflow-staging"


def _staging_dir(mlflow_run_id: str) -> Path:
    return _STAGING_ROOT / mlflow_run_id


def stage_file(mlflow_run_id: str, local_file: str, artifact_path: str | None) -> None:
    """Copy a local file into the staging area, preserving artifact_path sub-directory."""
    src = Path(local_file)
    if not src.is_file():
        raise FileNotFoundError(f"Artifact not found: {local_file}")

    dest_base = _staging_dir(mlflow_run_id) / "artifacts"
    if artifact_path:
        dest_dir = dest_base / artifact_path
    else:
        dest_dir = dest_base
    dest_dir.mkdir(parents=True, exist_ok=True)

    shutil.copy2(src, dest_dir / src.name)


def write_meta(mlflow_run_id: str, project: str) -> None:
    """Write (or update) meta.json for a staged run."""
    meta_path = _staging_dir(mlflow_run_id) / "meta.json"
    if meta_path.exists():
        return  # already written on first artifact
    meta_path.parent.mkdir(parents=True, exist_ok=True)
    meta = {
        "mlflow_run_id": mlflow_run_id,
        "project": project,
        "staged_at": datetime.now(timezone.utc).isoformat(),
    }
    meta_path.write_text(json.dumps(meta, indent=2))


def list_staged_files(mlflow_run_id: str) -> list[Path]:
    """Return all staged artifact files for a run."""
    artifacts_dir = _staging_dir(mlflow_run_id) / "artifacts"
    if not artifacts_dir.exists():
        return []
    return [p for p in artifacts_dir.rglob("*") if p.is_file()]


def get_meta(mlflow_run_id: str) -> dict:
    meta_path = _staging_dir(mlflow_run_id) / "meta.json"
    if not meta_path.exists():
        raise FileNotFoundError(f"No staged run found: {mlflow_run_id}")
    return json.loads(meta_path.read_text())


def list_all_staged() -> list[dict]:
    """Return meta for every staged mlflow run."""
    result = []
    if not _STAGING_ROOT.exists():
        return result
    for run_dir in sorted(_STAGING_ROOT.iterdir()):
        meta_path = run_dir / "meta.json"
        if meta_path.exists():
            result.append(json.loads(meta_path.read_text()))
    return result
