src/duration.py

from __future__ import annotations

import math
from dataclasses import asdict, dataclass
from pathlib import Path

from .catalog import load_catalog
from .config import VIDEO_EXTENSIONS
from .media import normalize_uploads, probe_duration

DURATION_MODE_MANUAL = "Manual (fixed seconds)"
DURATION_MODE_SEMI_AUTO = "Semi-auto (input/model estimate)"
DURATION_MODES = [DURATION_MODE_SEMI_AUTO, DURATION_MODE_MANUAL]

ZERO_GPU_MIN_SECONDS = 30
ZERO_GPU_MAX_SECONDS = 300
DEFAULT_MANUAL_SECONDS = 180
DEFAULT_BASE_SECONDS_PER_MINUTE = 12.0
DEFAULT_SAFETY_MULTIPLIER = 1.25
DEFAULT_SAFE_MODE_ENABLED = False
SAFE_MODE_MULTIPLIER = 1.30
# Calibration runs showed that the first inference in a fresh ZeroGPU callback
# can pay a sizeable CUDA/ONNX kernel warm-up cost that is not proportional to
# source duration. Keep this as an explicit, inspectable term rather than hiding
# it in the per-minute multiplier.
ONNX_FIRST_INFERENCE_WARMUP_SECONDS = 30.0
NON_ONNX_FIRST_INFERENCE_WARMUP_SECONDS = 20.0
ADDITIONAL_MODEL_WARMUP_SECONDS = 8.0
ADDITIONAL_BATCH_ITEM_DISPATCH_SECONDS = 1.5
DURATION_CALIBRATION_REVISION = "cold-start-batch-20260718"


@dataclass(frozen=True)
class DurationEstimate:
    mode: str
    seconds: int
    unclamped_seconds: float
    base_unclamped_seconds: float
    total_input_seconds: float
    file_count: int
    video_count: int
    model_count: int
    model_factor_sum: float
    parameter_factor: float
    base_seconds_per_minute: float
    safety_multiplier: float
    safe_mode_requested: bool
    safe_mode_applied: bool
    safe_mode_multiplier: float
    fixed_overhead_seconds: float
    first_inference_warmup_seconds: float
    additional_batch_item_overhead_seconds: float
    onnx_model_count: int
    non_onnx_model_count: int
    calibration_revision: str
    unprobed_files: int

    def to_dict(self) -> dict:
        return asdict(self)


def _clamp_seconds(value: float) -> int:
    return int(max(ZERO_GPU_MIN_SECONDS, min(ZERO_GPU_MAX_SECONDS, math.ceil(value))))


def _input_stats(files) -> tuple[float, int, int, int]:
    total_seconds = 0.0
    file_count = 0
    video_count = 0
    unprobed = 0
    for path in normalize_uploads(files):
        file_count += 1
        if path.suffix.lower() in VIDEO_EXTENSIONS:
            video_count += 1
        try:
            seconds = float(probe_duration(path))
            if not math.isfinite(seconds) or seconds <= 0:
                raise ValueError("duration unavailable")
        except Exception:
            # Keep the duration callable non-fatal. One minute is a conservative
            # fallback for an upload that ffprobe cannot inspect before dispatch.
            seconds = 60.0
            unprobed += 1
        total_seconds += seconds
    if file_count == 0:
        total_seconds = 60.0
        file_count = 1
        unprobed = 1
    return total_seconds, file_count, video_count, unprobed


def _catalog_factor(model_id: str) -> float:
    item = load_catalog().get(str(model_id))
    if not item:
        return 1.15
    factor = 1.0
    category = str(item.get("category", "")).lower()
    text = " ".join(
        str(item.get(key, "")) for key in ("name", "model_filename", "architecture")
    ).lower()
    if "4-stem" in category or "4stem" in text or "6stem" in text:
        factor *= 1.45
    if "scnet" in text:
        factor *= 1.15
    if any(token in text for token in (" xl", "large", "super_big", "super big", "big_beta")):
        factor *= 1.15
    return factor


def _package_factor(filename: str) -> float:
    suffix = Path(str(filename)).suffix.lower()
    if suffix == ".onnx":
        return 0.75
    if suffix == ".pth":
        return 0.95
    return 1.0


def _model_stats(catalog_model_ids, package_model_filenames, custom_source: str) -> tuple[int, float]:
    catalog_ids = list(catalog_model_ids or [])
    package_files = list(package_model_filenames or [])
    factors = [_catalog_factor(model_id) for model_id in catalog_ids]
    factors.extend(_package_factor(filename) for filename in package_files)
    if str(custom_source or "None") in {"Hugging Face", "GitHub"}:
        factors.append(1.25)
    if not factors:
        factors = [1.0]
    return len(factors), float(sum(factors))


def _model_runtime_profile(catalog_model_ids, package_model_filenames, custom_source: str) -> tuple[int, int]:
    """Return ONNX/non-ONNX model counts for cold-start calibration."""
    filenames: list[str] = []
    catalog = load_catalog()
    for model_id in list(catalog_model_ids or []):
        item = catalog.get(str(model_id)) or {}
        filename = str(item.get("model_filename") or item.get("filename") or "")
        filenames.append(filename)
    filenames.extend(str(value or "") for value in list(package_model_filenames or []))
    custom_count = 1 if str(custom_source or "None") in {"Hugging Face", "GitHub"} else 0
    onnx_count = sum(1 for filename in filenames if Path(filename).suffix.lower() == ".onnx")
    non_onnx_count = max(0, len(filenames) - onnx_count) + custom_count
    if onnx_count + non_onnx_count == 0:
        non_onnx_count = 1
    return onnx_count, non_onnx_count


def calculate_duration_estimate(
    files,
    catalog_model_ids,
    package_model_filenames,
    custom_source,
    ensemble_algorithm,
    output_format,
    pitch_shift,
    chunk_duration,
    duration_mode,
    manual_gpu_seconds,
    semi_auto_base_seconds,
    semi_auto_safety_multiplier,
    semi_auto_safe_mode=DEFAULT_SAFE_MODE_ENABLED,
    total_input_seconds_override=None,
) -> DurationEstimate:
    mode = str(duration_mode or DURATION_MODE_SEMI_AUTO)
    total_seconds, file_count, video_count, unprobed = _input_stats(files)
    if total_input_seconds_override is not None:
        try:
            override = float(total_input_seconds_override)
        except (TypeError, ValueError):
            override = 0.0
        if math.isfinite(override) and override > 0:
            total_seconds = override
            unprobed = 0

    model_count, model_factor_sum = _model_stats(
        catalog_model_ids, package_model_filenames, str(custom_source or "None")
    )
    onnx_model_count, non_onnx_model_count = _model_runtime_profile(
        catalog_model_ids, package_model_filenames, str(custom_source or "None")
    )

    manual = float(manual_gpu_seconds or DEFAULT_MANUAL_SECONDS)
    base = max(1.0, float(semi_auto_base_seconds or DEFAULT_BASE_SECONDS_PER_MINUTE))
    safety = max(0.5, float(semi_auto_safety_multiplier or DEFAULT_SAFETY_MULTIPLIER))
    safe_requested = bool(semi_auto_safe_mode)

    if mode == DURATION_MODE_MANUAL:
        seconds = _clamp_seconds(manual)
        return DurationEstimate(
            mode=mode,
            seconds=seconds,
            unclamped_seconds=manual,
            base_unclamped_seconds=manual,
            total_input_seconds=total_seconds,
            file_count=file_count,
            video_count=video_count,
            model_count=model_count,
            model_factor_sum=model_factor_sum,
            parameter_factor=1.0,
            base_seconds_per_minute=base,
            safety_multiplier=safety,
            safe_mode_requested=safe_requested,
            safe_mode_applied=False,
            safe_mode_multiplier=1.0,
            fixed_overhead_seconds=0.0,
            first_inference_warmup_seconds=0.0,
            additional_batch_item_overhead_seconds=0.0,
            onnx_model_count=onnx_model_count,
            non_onnx_model_count=non_onnx_model_count,
            calibration_revision=DURATION_CALIBRATION_REVISION,
            unprobed_files=unprobed,
        )

    parameter_factor = 1.0
    if int(pitch_shift or 0) != 0:
        parameter_factor *= 1.12
    if int(chunk_duration or 0) > 0:
        parameter_factor *= 1.08
    if str(output_format or "FLAC").upper() in {"MP3", "OGG", "OPUS", "M4A"}:
        parameter_factor *= 1.06
    algorithm = str(ensemble_algorithm or "avg_wave").lower()
    if "fft" in algorithm:
        parameter_factor *= 1.08
    elif algorithm.startswith("uvr_") or algorithm == "ensemble_wav":
        parameter_factor *= 1.12
    if model_count > 1:
        parameter_factor *= 1.0 + min(0.25, 0.05 * (model_count - 1))

    fixed_overhead = 18.0 + 4.0 * model_count + 2.0 * file_count + 4.0 * video_count
    first_inference_warmup = (
        ONNX_FIRST_INFERENCE_WARMUP_SECONDS
        if onnx_model_count > 0
        else NON_ONNX_FIRST_INFERENCE_WARMUP_SECONDS
    )
    if model_count > 1:
        first_inference_warmup += ADDITIONAL_MODEL_WARMUP_SECONDS * (model_count - 1)
    additional_batch_item_overhead = ADDITIONAL_BATCH_ITEM_DISPATCH_SECONDS * max(0, file_count - 1)
    compute = (total_seconds / 60.0) * base * model_factor_sum * parameter_factor
    base_raw = (
        fixed_overhead
        + first_inference_warmup
        + additional_batch_item_overhead
        + compute
    ) * safety
    safe_applied = safe_requested
    effective_safe_multiplier = SAFE_MODE_MULTIPLIER if safe_applied else 1.0
    raw = base_raw * effective_safe_multiplier
    seconds = _clamp_seconds(raw)
    return DurationEstimate(
        mode=mode,
        seconds=seconds,
        unclamped_seconds=raw,
        base_unclamped_seconds=base_raw,
        total_input_seconds=total_seconds,
        file_count=file_count,
        video_count=video_count,
        model_count=model_count,
        model_factor_sum=model_factor_sum,
        parameter_factor=parameter_factor,
        base_seconds_per_minute=base,
        safety_multiplier=safety,
        safe_mode_requested=safe_requested,
        safe_mode_applied=safe_applied,
        safe_mode_multiplier=effective_safe_multiplier,
        fixed_overhead_seconds=fixed_overhead,
        first_inference_warmup_seconds=first_inference_warmup,
        additional_batch_item_overhead_seconds=additional_batch_item_overhead,
        onnx_model_count=onnx_model_count,
        non_onnx_model_count=non_onnx_model_count,
        calibration_revision=DURATION_CALIBRATION_REVISION,
        unprobed_files=unprobed,
    )


def _extract_duration_args(args, kwargs) -> tuple:
    def value(index: int, name: str, default=None):
        if len(args) > index:
            return args[index]
        return kwargs.get(name, default)

    return (
        value(0, "files"),
        value(1, "catalog_model_ids", []),
        value(2, "package_model_filenames", []),
        value(3, "custom_source", "None"),
        value(11, "ensemble_algorithm", "avg_wave"),
        value(12, "output_format", "FLAC"),
        value(14, "pitch_shift", 0),
        value(15, "chunk_duration", 0),
        value(16, "duration_mode", DURATION_MODE_SEMI_AUTO),
        value(17, "manual_gpu_seconds", DEFAULT_MANUAL_SECONDS),
        value(18, "semi_auto_base_seconds", DEFAULT_BASE_SECONDS_PER_MINUTE),
        value(19, "semi_auto_safety_multiplier", DEFAULT_SAFETY_MULTIPLIER),
        value(20, "semi_auto_safe_mode", DEFAULT_SAFE_MODE_ENABLED),
    )


def estimate_gpu_seconds(*args, **kwargs) -> int:
    try:
        return calculate_duration_estimate(*_extract_duration_args(args, kwargs)).seconds
    except Exception:
        return DEFAULT_MANUAL_SECONDS


def estimate_gpu_details(*args, **kwargs) -> DurationEstimate:
    return calculate_duration_estimate(*_extract_duration_args(args, kwargs))


def format_duration_estimate(*args, **kwargs) -> str:
    estimate = estimate_gpu_details(*args, **kwargs)
    if estimate.mode == DURATION_MODE_MANUAL:
        return f"**ZeroGPU duration: {estimate.seconds} seconds** — manual fixed value; Safe mode is not applied."
    fallback = (
        f" · ffprobe fallback used for {estimate.unprobed_files} file(s)"
        if estimate.unprobed_files
        else ""
    )
    safe_text = (
        f"Safe mode ON: base {math.ceil(estimate.base_unclamped_seconds)}s × "
        f"{estimate.safe_mode_multiplier:.2f} → {estimate.seconds}s"
        if estimate.safe_mode_applied
        else f"Safe mode OFF: {estimate.seconds}s"
    )
    return (
        f"**ZeroGPU duration estimate: {estimate.seconds} seconds** — {safe_text}  \n"
        f"Input {estimate.total_input_seconds:.1f}s / {estimate.file_count} file(s), "
        f"models {estimate.model_count} (factor sum {estimate.model_factor_sum:.2f}), "
        f"cold-start {estimate.first_inference_warmup_seconds:.0f}s, "
        f"Batch dispatch {estimate.additional_batch_item_overhead_seconds:.1f}s, "
        f"parameter factor {estimate.parameter_factor:.2f}, calibration {estimate.safety_multiplier:.2f}"
        f"{fallback}."
    )