"""Stage 3 — publish the artifact set to S3.

`latest.json` is written last and separately: a partially-uploaded run must
never be advertised as current.
"""

from __future__ import annotations

import json
from typing import Dict, List

from .artifact_store import delete_prefix, list_objects, list_subprefixes, upload_dir, write_text
from .config import Config
from .enumerate_assets import load_manifest
from .export import ARTIFACTS


def run_prefix(cfg: Config, manifest: Dict) -> str:
    base = cfg.require("ARTIFACT_S3_URI").rstrip("/")
    return f"{base}/{manifest['corpus_name']}/{manifest['run_id']}"


def verify_uploaded(cfg: Config, prefix: str) -> List[str]:
    """List what actually landed, so a silent partial upload is caught."""
    return [rel for rel, _ in list_objects(cfg, prefix)]


def prune(cfg: Config, manifest: Dict) -> None:
    keep = cfg.get_int("ARTIFACT_RETENTION_RUNS", 1000)
    base = cfg.require("ARTIFACT_S3_URI").rstrip("/") + f"/{manifest['corpus_name']}"
    runs = sorted(list_subprefixes(cfg, base))
    if len(runs) <= keep:
        print(f"  retention: {len(runs)} runs stored, keeping {keep} — nothing to prune")
        return
    for stale in runs[: len(runs) - keep]:
        if stale == manifest["run_id"]:
            continue  # never remove the run we just wrote
        print(f"  pruning {stale}")
        delete_prefix(cfg, f"{base}/{stale}")


def run(cfg: Config, update_latest: bool = True) -> int:
    if not cfg.get("ARTIFACT_S3_URI").strip():
        print(
            "ARTIFACT_S3_URI is not set — this is the pipeline's one unresolved input.\n"
            "Set it in config/pipeline.env once the target bucket is known. Every other\n"
            "stage runs without it; only this one is blocked."
        )
        return 2

    manifest = load_manifest()
    if not ARTIFACTS.exists():
        raise SystemExit("no artifacts found — run `export` first.")

    prefix = run_prefix(cfg, manifest)
    print(f"uploading {ARTIFACTS} -> {prefix}")
    local = upload_dir(cfg, ARTIFACTS, prefix)

    remote = set(verify_uploaded(cfg, prefix))
    missing = [name for name in local if name not in remote]
    if missing:
        print(f"\nFAIL: {len(missing)} file(s) missing remotely: {missing[:5]}")
        return 1
    print(f"  verified {len(local)} objects")

    if not update_latest:
        # Every Space without RESTORE_RUN_ID follows latest.json, so leaving
        # it alone publishes the run without switching any live Space to it.
        # Pruning is skipped too: it must never race the run that is current.
        print(f"\npublished {manifest['run_id']} — latest.json NOT updated (--no-latest); "
              f"use RESTORE_RUN_ID={manifest['run_id']} to load it")
        return 0

    # Written last, and only now: this is what marks the run as current.
    pointer = {
        "run_id": manifest["run_id"],
        "corpus_sha256": manifest["corpus_sha256"],
        "asset_count": manifest.get("asset_count"),
        "prefix": prefix,
    }
    base = cfg.require("ARTIFACT_S3_URI").rstrip("/")
    write_text(cfg, f"{base}/{manifest['corpus_name']}/latest.json", json.dumps(pointer, indent=2) + "\n")

    prune(cfg, manifest)
    print(f"\npublished {manifest['run_id']}")
    return 0