src/preparation.py

from __future__ import annotations

import hashlib
import json
import math
import os
import platform
import re
import secrets
import shutil
import time
from dataclasses import asdict
from datetime import datetime, timezone
from importlib import metadata
from pathlib import Path

import torch

from .catalog import load_catalog
from .config import JOB_ROOT, JOB_TTL_SECONDS, MAX_MODELS, MODEL_ROOT, OUTPUT_FORMATS, SAFE_PACKAGE_MODELS
from .dataset_samples import DEFAULT_DATASET_SAMPLE_ID, resolve_dataset_sample
from .downloads import (
    CustomModelRecord,
    ensure_runtime_dirs,
    file_sha256,
    package_model_cache_files,
    prefetch_package_models,
    resolve_catalog_models,
    resolve_github_model,
    resolve_hf_model,
)
from .duration import calculate_duration_estimate
from .errors import (
    INPUT_VALIDATION,
    MEDIA_PREPARATION,
    MODEL_DOWNLOAD,
    PREPARED_JOB_VALIDATION,
    StageError,
    stage_error,
)
from .jobs import JobPaths, cleanup_old_jobs, create_job, get_job
from .longform import (
    CHUNK_MODE_AUTO,
    CHUNK_MODE_DISABLED,
    CHUNK_MODE_FIXED,
    DEFAULT_FIXED_CHUNK_SECONDS,
    DEFAULT_PREVIEW_SECONDS,
    RANGE_MODE_FULL,
    build_long_form_plan,
)
from .media import normalize_uploads, prepare_inputs, probe_duration, validate_uploads
from .observability import DEFAULT_LOG_LEVEL, JobLogCapture, ProgressBridge, ProgressReporter, normalize_log_level

PREPARATION_SCHEMA = "sesa-preparation-v21"
REPRODUCIBILITY_SCHEMA = "sesa-reproducibility-v21"
PREPARED_STATE_SCHEMA = "sesa-prepared-state-v1"

BITRATE_AUTO = "Auto"
BUNDLE_LAYOUT_FLAT = "Flat outputs"
BUNDLE_LAYOUT_BY_INPUT = "Group by input"


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _package_version(name: str) -> str | None:
    try:
        return metadata.version(name)
    except metadata.PackageNotFoundError:
        return None


def _json_write(path: Path, payload: dict) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8")
    os.replace(temporary, path)
    return path


def _canonical_digest(payload: dict) -> str:
    serialized = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(serialized.encode("utf-8")).hexdigest()


def _digest_for_private_config(config: dict) -> str:
    payload = dict(config)
    payload.pop("config_sha256", None)
    payload.pop("access_token", None)
    return _canonical_digest(payload)


def _public_config(config: dict) -> dict:
    """Return a deep-copied public payload without secrets or host paths."""
    payload = json.loads(json.dumps(config))
    payload.pop("access_token", None)
    source = payload.get("source")
    if isinstance(source, dict):
        source.pop("local_path", None)
    for item in payload.get("inputs", []) or []:
        if isinstance(item, dict):
            item.pop("source_path", None)
    models = payload.get("models")
    if isinstance(models, dict):
        for item in models.get("custom_records", []) or []:
            if isinstance(item, dict):
                item.pop("model_path", None)
                item.pop("config_path", None)
        for item in models.get("records", []) or []:
            if not isinstance(item, dict):
                continue
            item.pop("model_path", None)
            item.pop("config_path", None)
            for file_record in item.get("files", []) or []:
                if isinstance(file_record, dict):
                    file_record.pop("path", None)
    return payload


def public_preparation_config(config: dict) -> dict:
    """Public frozen Preflight payload; its top-level schema remains Preparation."""
    return _public_config(config)


def public_reproducibility_config(config: dict) -> dict:
    """Public final payload with an unambiguous reproducibility top-level schema."""
    payload = _public_config(config)
    preparation_schema = payload.get("schema")
    reproducibility_schema = payload.pop("reproducibility_schema", REPRODUCIBILITY_SCHEMA)
    payload["schema"] = reproducibility_schema
    payload["preparation_schema"] = preparation_schema
    return payload


def _safe_relative(job: JobPaths, path: Path) -> str:
    resolved = path.resolve()
    if job.root.resolve() not in resolved.parents:
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared path escaped the job workspace.", retryable=False)
    return resolved.relative_to(job.root.resolve()).as_posix()


def _resolve_custom_source(
    source_kind: str,
    hf_repo: str,
    hf_weight: str,
    hf_config: str,
    hf_revision: str,
    github_weight_url: str,
    github_config_url: str,
    github_sha256: str,
) -> list[CustomModelRecord]:
    if source_kind == "Hugging Face":
        if not (str(hf_repo).strip() and str(hf_weight).strip() and str(hf_config).strip()):
            raise ValueError("Complete the Hugging Face repository, weight, and YAML fields.")
        return [resolve_hf_model(hf_repo, hf_weight, hf_config, hf_revision or "main")]
    if source_kind == "GitHub":
        if not (str(github_weight_url).strip() and str(github_config_url).strip()):
            raise ValueError("Complete both GitHub release-weight and YAML URLs.")
        return [resolve_github_model(github_weight_url, github_config_url, github_sha256)]
    return []


def _input_source(upload_files, sample_id: str) -> tuple[list[Path], dict]:
    uploads = normalize_uploads(upload_files)
    if uploads:
        return validate_uploads(uploads), {
            "source_kind": "upload",
            "sample_id": None,
            "license_id": None,
            "attribution": None,
        }
    selected = str(sample_id or "").strip()
    if not selected:
        raise ValueError("Upload at least one file or select a public sample.")
    resolved = resolve_dataset_sample(selected or DEFAULT_DATASET_SAMPLE_ID)
    return [Path(resolved.local_path)], {
        "source_kind": "dataset-sample",
        **asdict(resolved),
    }


def _provider_snapshot() -> dict:
    try:
        import onnxruntime as ort

        providers = list(ort.get_available_providers())
        ort_version = getattr(ort, "__version__", None)
    except Exception as exc:
        providers = []
        ort_version = None
        ort_error = f"{type(exc).__name__}: {exc}"
    else:
        ort_error = None
    cuda_visible = bool(torch.cuda.is_available())
    return {
        "preparation_torch_cuda_visible": cuda_visible,
        "preparation_torch_cuda_note": (
            "ZeroGPU may expose CUDA emulation outside the decorated callback; actual CUDA is checked again inside the callback."
        ),
        "onnxruntime_version": ort_version,
        "onnx_available_providers": providers,
        "onnx_cuda_provider_available": "CUDAExecutionProvider" in providers,
        "onnx_provider_inspection_error": ort_error,
    }


def _estimated_output_bytes(total_seconds: float, output_format: str, bitrate: str, sample_rate: int, stem_count: int) -> int:
    duration = max(0.0, float(total_seconds))
    stems = max(1, int(stem_count))
    fmt = str(output_format or "FLAC").upper()
    rate = max(8000, int(sample_rate or 44100))
    if bitrate and bitrate != BITRATE_AUTO and str(bitrate).lower().endswith("k"):
        try:
            bits_per_second = float(str(bitrate)[:-1]) * 1000.0
            return int(duration * bits_per_second / 8.0 * stems * 1.08)
        except ValueError:
            pass
    raw = duration * rate * 2 * 2 * stems
    factor = 0.62 if fmt == "FLAC" else 1.0
    if fmt in {"MP3", "OGG", "OPUS", "M4A"}:
        factor = 0.18
    return int(raw * factor)


def _model_cache_before(catalog_ids: list[str], package_filenames: list[str], custom_names: list[str]) -> dict:
    catalog = load_catalog()
    catalog_before = {}
    for model_id in catalog_ids:
        item = catalog.get(model_id) or {}
        filename = Path(str(item.get("weight_filename") or "")).name
        catalog_before[model_id] = bool(filename and list(MODEL_ROOT.rglob(filename)))
    package_before = {name: bool(package_model_cache_files(name)) for name in package_filenames}
    custom_before = {name: bool(name and list(MODEL_ROOT.rglob(Path(name).name))) for name in custom_names}
    return {"catalog": catalog_before, "package": package_before, "custom": custom_before}


def _serialize_custom_record(record: CustomModelRecord, cache_before: bool | None = None) -> dict:
    return {
        "kind": "custom",
        "key": record.key,
        "display_name": record.display_name,
        "filename": record.filename,
        "architecture": record.architecture,
        "model_path": str(record.model_path),
        "model_size_bytes": record.model_path.stat().st_size,
        "model_sha256": file_sha256(record.model_path),
        "config_path": str(record.config_path),
        "config_size_bytes": record.config_path.stat().st_size,
        "config_sha256": file_sha256(record.config_path),
        "cache_before": cache_before,
        "cache_after": True,
        "downloaded_during_preparation": cache_before is False,
    }


def _preflight_markdown(config: dict) -> str:
    duration = config["duration"]
    providers = config["providers"]
    models = config["models"]
    output = config["output"]
    source = config["source"]
    long_form = config.get("long_form") or {}
    batch = config.get("batch") or {}
    downloads = sum(1 for item in models["records"] if item.get("downloaded_during_preparation"))
    provider_text = ", ".join(providers.get("onnx_available_providers") or []) or "not detected during Preparation"
    safe = "ON" if duration.get("safe_mode_applied") else "OFF"
    sample_text = source.get("label") or source.get("sample_id") or "uploaded file(s)"
    raw_duration = float(duration.get("unclamped_seconds") or duration.get("seconds") or 0)
    duration_warning = (
        f"\n- Runtime warning: estimate {raw_duration:.0f}s saturated the {duration['seconds']}s declaration ceiling; "
        "the platform may warn and might abort the task; completion is not guaranteed."
        if raw_duration > float(duration.get("seconds") or 0) + 0.5
        else ""
    )
    chunk_text = (
        f"{long_form.get('resolved_chunk_seconds')}s × about {long_form.get('chunk_count_total')} chunk(s)"
        if int(long_form.get("resolved_chunk_seconds") or 0) > 0
        else "disabled"
    )
    return (
        "### Preflight ready\n"
        f"- Job: `{config['job_id']}`\n"
        f"- Source: {source.get('source_kind')} · {sample_text}\n"
        f"- Inputs: {len(config['inputs'])} · selected {config['input_total_seconds']:.1f} seconds\n"
        f"- Batch: {'enabled' if batch.get('enabled') else 'single item'} · shared model load · sequential item execution · "
        f"continue after item failure {'ON' if batch.get('continue_on_item_error') else 'OFF'}\n"
        f"- Range: {long_form.get('range_mode', 'Full input')} · chunk policy {long_form.get('chunk_mode', 'Disabled')} → {chunk_text}\n"
        f"- Models: {len(models['selected_filenames'])} · downloaded during Preparation: {downloads}\n"
        f"- ZeroGPU declaration: **{duration['seconds']} seconds** · Safe mode {safe}"
        f"{duration_warning}\n"
        f"- Duration calibration: cold-start {float(duration.get('first_inference_warmup_seconds') or 0):.0f}s · "
        f"additional Batch dispatch {float(duration.get('additional_batch_item_overhead_seconds') or 0):.1f}s · "
        f"profile `{duration.get('calibration_revision') or 'legacy'}`\n"
        f"- Output: {output['format']} · bitrate {output['bitrate']} · {output['sample_rate']} Hz · {output['single_stem']}\n"
        f"- Estimated output size: {output['estimated_output_bytes'] / (1024 * 1024):.1f} MiB\n"
        f"- ONNX providers seen before GPU allocation: `{provider_text}`\n"
        "- Settings are frozen in the preparation JSON. Change settings, then prepare again to create a new isolated job."
    )


def prepared_state(job_id: str, access_token: str) -> str:
    return json.dumps({"schema": PREPARED_STATE_SCHEMA, "job_id": job_id, "access_token": access_token}, sort_keys=True)


def parse_prepared_state(value) -> tuple[JobPaths, dict]:
    try:
        state = json.loads(str(value or ""))
    except Exception as exc:
        raise StageError(PREPARED_JOB_VALIDATION, "Prepare the job before starting separation.", hint="Click Prepare / refresh preflight.") from exc
    if not isinstance(state, dict) or state.get("schema") != PREPARED_STATE_SCHEMA:
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job state is invalid.", retryable=False)
    job = get_job(str(state.get("job_id") or ""))
    config_path = job.config_dir / "preparation.json"
    marker = job.config_dir / ".prepared"
    if not config_path.is_file() or not marker.is_file():
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job is incomplete or expired.", hint="Prepare the job again.")
    config = json.loads(config_path.read_text(encoding="utf-8"))
    if not secrets.compare_digest(str(state.get("access_token") or ""), str(config.get("access_token") or "")):
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job access token does not match.", retryable=False)
    if config.get("schema") != PREPARATION_SCHEMA or config.get("status") != "PREPARED":
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job is not ready.", hint="Prepare the job again.")
    expected_digest = str(config.get("config_sha256") or "")
    marker_digest = marker.read_text(encoding="utf-8").strip()
    actual_digest = _digest_for_private_config(config)
    if not expected_digest or not secrets.compare_digest(expected_digest, marker_digest):
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job integrity marker does not match.", retryable=False)
    if not secrets.compare_digest(expected_digest, actual_digest):
        raise StageError(PREPARED_JOB_VALIDATION, "Prepared job configuration changed after Preparation.", retryable=False)
    return job, config


def cleanup_prepared_state(value) -> None:
    """Best-effort session cleanup. TTL cleanup remains the fallback."""
    try:
        state = json.loads(str(value or ""))
        job = get_job(str(state.get("job_id") or ""))
    except Exception:
        return
    from .jobs import safe_remove_job

    safe_remove_job(job)


def estimate_prepared_gpu_seconds(prepared_state_value, *args, **kwargs) -> int:
    del args, kwargs
    try:
        _, config = parse_prepared_state(prepared_state_value)
        return int(config.get("duration", {}).get("seconds") or 60)
    except Exception:
        return 60


def prepare_job(
    files,
    sample_id,
    catalog_model_ids,
    package_model_filenames,
    custom_source,
    hf_repo,
    hf_weight,
    hf_config,
    hf_revision,
    github_weight_url,
    github_config_url,
    github_sha256,
    ensemble_algorithm,
    output_format,
    output_bitrate,
    output_sample_rate,
    normalization_threshold,
    amplification_threshold,
    single_stem,
    bundle_layout,
    allow_cpu_fallback,
    pitch_shift,
    chunk_duration,
    duration_mode,
    manual_gpu_seconds,
    semi_auto_base_seconds,
    semi_auto_safety_multiplier,
    semi_auto_safe_mode,
    log_level=DEFAULT_LOG_LEVEL,
    progress=None,
    *,
    range_mode=RANGE_MODE_FULL,
    range_start_seconds=0,
    range_end_seconds=0,
    preview_seconds=DEFAULT_PREVIEW_SECONDS,
    chunk_mode=None,
    fixed_chunk_seconds=DEFAULT_FIXED_CHUNK_SECONDS,
    batch_continue_on_item_error=True,
) -> tuple[str, str, str, str]:
    cleanup_old_jobs()
    ensure_runtime_dirs()
    job = create_job()
    level_name = normalize_log_level(log_level)
    preparation_log = job.logs_dir / f"sesa_prepare_{job.job_id}.log"
    bridge = ProgressBridge(progress)
    access_token = secrets.token_urlsafe(32)
    started = time.perf_counter()
    config: dict = {
        "schema": PREPARATION_SCHEMA,
        "reproducibility_schema": REPRODUCIBILITY_SCHEMA,
        "status": "PREPARING",
        "job_id": job.job_id,
        "access_token": access_token,
        "created_at_utc": _utc_now(),
        "expires_after_seconds": JOB_TTL_SECONDS,
    }

    with JobLogCapture(preparation_log, job.job_id, level=level_name) as log:
        reporter = ProgressReporter(bridge, log.event, log.debug_event, level=level_name)
        try:
            reporter.update(0.02, "Creating isolated Preparation workspace", stage="preparation", force_info=True)
            try:
                sources, source_metadata = _input_source(files, sample_id)
            except Exception as exc:
                raise stage_error(INPUT_VALIDATION, exc, hint="Upload a supported file or select a public sample.")
            log.event("preparation-source-resolved", source_kind=source_metadata.get("source_kind"), count=len(sources))

            reporter.update(0.10, "Validating, probing, and planning media range", stage="input-validation", force_info=True)
            try:
                validated = validate_uploads(sources)
                source_records = []
                source_durations = []
                for path in validated:
                    source_duration = round(float(probe_duration(path)), 6)
                    source_durations.append(source_duration)
                    source_records.append({
                        "display_name": path.name,
                        "source_path": str(path),
                        "source_size_bytes": path.stat().st_size,
                        "source_sha256": file_sha256(path),
                        "source_duration_seconds": source_duration,
                    })
            except Exception as exc:
                raise stage_error(INPUT_VALIDATION, exc, hint="Check file format, size, and readability.")

            catalog_ids = list(catalog_model_ids or [])
            package_files = list(package_model_filenames or [])
            raw_model_count = len(catalog_ids) + len(package_files) + (
                1 if str(custom_source or "None") != "None" else 0
            )
            try:
                selected_chunk_mode = chunk_mode
                selected_fixed_chunk = fixed_chunk_seconds
                if selected_chunk_mode is None:
                    if int(chunk_duration or 0) > 0:
                        selected_chunk_mode = CHUNK_MODE_FIXED
                        selected_fixed_chunk = int(chunk_duration)
                    else:
                        selected_chunk_mode = CHUNK_MODE_DISABLED
                long_form_plan = build_long_form_plan(
                    source_durations,
                    range_mode=range_mode,
                    range_start_seconds=range_start_seconds,
                    range_end_seconds=range_end_seconds,
                    preview_seconds=preview_seconds,
                    chunk_mode=selected_chunk_mode,
                    fixed_chunk_seconds=selected_fixed_chunk,
                    model_count=max(1, raw_model_count),
                )
                planned_ranges = (
                    long_form_plan.input_ranges if long_form_plan.selection_applied else None
                )
            except Exception as exc:
                raise stage_error(
                    INPUT_VALIDATION,
                    exc,
                    hint="Check the selected start/end range and chunk policy.",
                )


            for record, range_record in zip(source_records, long_form_plan.input_ranges):
                record["selection"] = dict(range_record)

            reporter.update(0.22, "Copying and normalizing selected media before GPU allocation", stage="media-preparation", force_info=True)
            try:
                if planned_ranges is None:
                    prepared, total_duration = prepare_inputs(job, validated)
                else:
                    prepared, total_duration = prepare_inputs(
                        job,
                        validated,
                        input_ranges=planned_ranges,
                    )
            except Exception as exc:
                raise stage_error(MEDIA_PREPARATION, exc, hint="Check FFmpeg/ffprobe and the media stream.")
            for record, path in zip(source_records, prepared):
                record.update({
                    "prepared_path": _safe_relative(job, path),
                    "prepared_size_bytes": path.stat().st_size,
                    "prepared_sha256": file_sha256(path),
                    "prepared_duration_seconds": round(float(probe_duration(path)), 6),
                })

            safe_package_filenames = {filename for _, filename in SAFE_PACKAGE_MODELS}
            unknown_package = [name for name in package_files if name not in safe_package_filenames]
            if unknown_package:
                raise StageError(INPUT_VALIDATION, f"Package model is not allowlisted: {unknown_package[0]}", retryable=False)
            if len(catalog_ids) + len(package_files) + (1 if str(custom_source or "None") != "None" else 0) > MAX_MODELS:
                raise StageError(INPUT_VALIDATION, f"A maximum of {MAX_MODELS} models can be used in one job.")
            if not catalog_ids and not package_files and str(custom_source or "None") == "None":
                raise StageError(INPUT_VALIDATION, "Select at least one model.")
            custom_names = [hf_weight] if custom_source == "Hugging Face" else [Path(str(github_weight_url)).name] if custom_source == "GitHub" else []
            cache_before = _model_cache_before(catalog_ids, package_files, custom_names)

            reporter.update(0.36, "Prefetching and locking model files outside GPU allocation", stage="model-download", force_info=True)
            try:
                custom_records = resolve_catalog_models(catalog_ids)
                custom_records.extend(_resolve_custom_source(
                    custom_source or "None", hf_repo or "", hf_weight or "", hf_config or "",
                    hf_revision or "main", github_weight_url or "", github_config_url or "", github_sha256 or "",
                ))
                package_records = prefetch_package_models(package_files)
            except Exception as exc:
                raise stage_error(MODEL_DOWNLOAD, exc, hint="Retry Preparation; shared cache downloads are file-locked.")

            records: list[dict] = []
            for record in custom_records:
                before = cache_before["catalog"].get(record.key)
                if before is None:
                    before = cache_before["custom"].get(record.filename)
                records.append(_serialize_custom_record(record, before))
            records.extend(package_records)
            selected_filenames = []
            seen = set()
            for record in custom_records:
                if record.filename not in seen:
                    selected_filenames.append(record.filename)
                    seen.add(record.filename)
            for filename in package_files:
                if filename not in seen:
                    selected_filenames.append(filename)
                    seen.add(filename)

            reporter.update(0.64, "Inspecting providers and freezing output plan", stage="preflight", force_info=True)
            provider_info = _provider_snapshot()
            selected_onnx = any(str(name).lower().endswith(".onnx") for name in selected_filenames)
            provider_info["selected_onnx_model"] = selected_onnx
            if selected_onnx:
                provider_info["preflight_provider_status"] = (
                    "pending callback recheck; CUDAExecutionProvider visible during Preparation"
                    if provider_info.get("onnx_cuda_provider_available")
                    else "pending callback recheck; CUDAExecutionProvider not visible during Preparation"
                )
                provider_info["onnx_provider_recheck"] = {
                    "applicability": "required",
                    "status": "pending",
                    "reason": "At least one prepared model is ONNX; available providers and loaded-session registered providers are checked inside the GPU callback.",
                }
            else:
                provider_info["preflight_provider_status"] = "not applicable; no ONNX model selected"
                provider_info["onnx_provider_recheck"] = {
                    "applicability": "not-applicable",
                    "status": "not-applicable",
                    "reason": "No prepared model filename ends with .onnx.",
                }
            output_format_value = str(output_format or "FLAC").upper()
            if output_format_value not in OUTPUT_FORMATS:
                raise StageError(INPUT_VALIDATION, f"Unsupported output format: {output_format_value}", retryable=False)
            bitrate = str(output_bitrate or BITRATE_AUTO)
            if bitrate != BITRATE_AUTO and not re.fullmatch(r"[1-9][0-9]{1,3}k", bitrate):
                raise StageError(INPUT_VALIDATION, "Output bitrate must look like 96k or 320k.", retryable=False)
            sample_rate = int(output_sample_rate or 44100)
            if sample_rate not in {24000, 32000, 44100, 48000}:
                raise StageError(INPUT_VALIDATION, "Unsupported output sample rate.", retryable=False)
            if str(bundle_layout or BUNDLE_LAYOUT_FLAT) not in {BUNDLE_LAYOUT_FLAT, BUNDLE_LAYOUT_BY_INPUT}:
                raise StageError(INPUT_VALIDATION, "Unsupported bundle layout.", retryable=False)
            normalization = max(0.0, min(1.0, float(normalization_threshold if normalization_threshold is not None else 0.9)))
            amplification = max(0.0, min(1.0, float(amplification_threshold if amplification_threshold is not None else 0.0)))
            stem_count = 1 if str(single_stem or "All stems") != "All stems" else 2
            duration = calculate_duration_estimate(
                files=prepared,
                catalog_model_ids=catalog_ids,
                package_model_filenames=package_files,
                custom_source=custom_source,
                ensemble_algorithm=ensemble_algorithm,
                output_format=output_format_value,
                pitch_shift=pitch_shift,
                chunk_duration=long_form_plan.resolved_chunk_seconds,
                duration_mode=duration_mode,
                manual_gpu_seconds=manual_gpu_seconds,
                semi_auto_base_seconds=semi_auto_base_seconds,
                semi_auto_safety_multiplier=semi_auto_safety_multiplier,
                semi_auto_safe_mode=semi_auto_safe_mode,
                total_input_seconds_override=total_duration,
            ).to_dict()
            duration["request_contract"] = {
                "meaning": "maximum requested ZeroGPU callback runtime",
                "overrun_behavior": (
                    "The platform may warn that the task might be aborted. An overrun that completes is a functional "
                    "success only and does not prove the request is safe."
                ),
                "shorter_request_queue_priority": True,
            }

            output_plan = {
                "format": output_format_value,
                "bitrate": bitrate,
                "sample_rate": sample_rate,
                "normalization_threshold": normalization,
                "amplification_threshold": amplification,
                "single_stem": str(single_stem or "All stems"),
                "bundle_layout": str(bundle_layout or BUNDLE_LAYOUT_FLAT),
                "physical_naming": "sesa_<job_uuid>_input_<index>_<sequence>_<stem>.<ext>",
                "estimated_output_bytes": _estimated_output_bytes(total_duration, output_format_value, bitrate, sample_rate, stem_count),
            }
            config.update({
                "status": "PREPARED",
                "prepared_at_utc": _utc_now(),
                "preparation_wall_seconds": round(time.perf_counter() - started, 6),
                "source": source_metadata,
                "inputs": source_records,
                "input_total_seconds": round(float(total_duration), 6),
                "batch": {
                    "schema": "sesa-batch-plan-v2",
                    "enabled": len(source_records) > 1,
                    "input_count": len(source_records),
                    "execution_order": "sequential-inputs-shared-loaded-model",
                    "model_load_scope": "once-per-attempt",
                    "continue_on_item_error": bool(batch_continue_on_item_error),
                    "output_mapping": "explicit-per-separate-call",
                    "retry_policy": "failed-and-unprocessed-items-only",
                    "cooperative_stop_boundary": "between-batch-items",
                    "status": "PREPARED",
                },
                "models": {
                    "catalog_ids": catalog_ids,
                    "package_filenames": package_files,
                    "custom_source": str(custom_source or "None"),
                    "custom_records": [
                        {
                            "key": item.key,
                            "display_name": item.display_name,
                            "filename": item.filename,
                            "model_path": str(item.model_path),
                            "config_path": str(item.config_path),
                            "architecture": item.architecture,
                        }
                        for item in custom_records
                    ],
                    "selected_filenames": selected_filenames,
                    "records": records,
                    "ensemble_algorithm": str(ensemble_algorithm or "avg_wave"),
                },
                "output": output_plan,
                "parameters": {
                    "pitch_shift": int(pitch_shift or 0),
                    "chunk_duration": int(long_form_plan.resolved_chunk_seconds),
                    "chunk_mode": long_form_plan.chunk_mode,
                    "allow_cpu_fallback": bool(allow_cpu_fallback),
                    "log_level": level_name,
                },
                "long_form": long_form_plan.to_dict(),
                "duration": duration,
                "providers": provider_info,
                "runtime": {
                    "python": platform.python_version(),
                    "platform": platform.platform(),
                    "packages": {
                        "gradio": _package_version("gradio"),
                        "spaces": _package_version("spaces"),
                        "torch": _package_version("torch"),
                        "audio-separator": _package_version("audio-separator"),
                        "onnxruntime-gpu": _package_version("onnxruntime-gpu"),
                        "huggingface-hub": _package_version("huggingface-hub"),
                    },
                },
                "error": None,
                "execution": None,
                "outputs": [],
            })
            config["config_sha256"] = _digest_for_private_config(config)
            config_path = _json_write(job.config_dir / "preparation.json", config)
            preflight_config = public_preparation_config(config)
            reproducibility_config = public_reproducibility_config(config)
            preflight_path = _json_write(job.config_dir / "preflight.json", preflight_config)
            _json_write(job.config_dir / "reproducibility.json", reproducibility_config)
            (job.config_dir / ".prepared").write_text(config["config_sha256"] + "\n", encoding="utf-8")
            log.event(
                "preparation-complete",
                model_count=len(selected_filenames),
                input_count=len(source_records),
                duration_seconds=duration["seconds"],
                config_sha256=config["config_sha256"],
            )
            reporter.finish("Preflight ready", stage="preparation")
            return prepared_state(job.job_id, access_token), _preflight_markdown(config), str(preflight_path), str(preparation_log)
        except Exception as exc:
            error = stage_error(INPUT_VALIDATION, exc)
            config.update({
                "status": "FAILED_PREPARATION",
                "failed_at_utc": _utc_now(),
                "preparation_wall_seconds": round(time.perf_counter() - started, 6),
                "error": error.to_dict(),
            })
            error_path = _json_write(job.config_dir / "preparation_error.json", config)
            log.event("preparation-failed", error_stage=error.stage, error_type=type(exc).__name__, error=str(exc))
            reporter.finish("Preparation failed before GPU allocation", stage="preparation")
            failure = StageError(error.stage, error.message, retryable=error.retryable, hint=error.hint)
            failure.diagnostic_path = str(error_path)
            failure.log_path = str(preparation_log)
            failure.job_id = job.job_id
            raise failure from exc