logging_utils.py

"""Console/stdout logging (LOG-1) and optional per-inference debug logging to a private Hugging
Face Hub dataset repo (LOG-2 through LOG-9).

Modeled on the FireRed-Image-Edit-1.0-Fast reference project's logging_utils.py/inference.py
pipeline, adapted for video: logged media is referenced by path rather than embedded as bytes
(LOG-5), the logged output video is the same final result served to the user, and retention
prunes on both a total-storage-size cap and a per-directory file-count cap (LOG-7) — the latter exists
because the Hub commit endpoint rejects pushes once a directory (data/, images/, videos/)
holds more than 10000 files, which a size-only cap doesn't bound since data/ and images/
files are tiny compared to videos/.
"""

import json
import os
import tempfile
import threading
import time as _time
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

import torch
from diffusers.utils.export_utils import export_to_video
from huggingface_hub import CommitOperationAdd, CommitOperationDelete, hf_hub_download
from PIL.Image import Image as PILImage

if TYPE_CHECKING:
    from huggingface_hub import HfApi


def print_startup_env() -> None:
    print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"), flush=True)
    print("torch.__version__ =", torch.__version__, flush=True)
    print(f"CUDA device_count={torch.cuda.device_count()}, is_available={torch.cuda.is_available()}", flush=True)


def print_infer_start(prompt: str, negative_prompt: str, seed: int, steps: int, guidance_scale: float,
                       frame_multiplier: int, upscale_output: bool) -> None:
    print(f"[infer] ===== START ===== steps={steps}, guidance={guidance_scale}, seed={seed}, "
          f"frame_multiplier={frame_multiplier}, upscale={upscale_output}", flush=True)
    print(f"[infer] prompt={prompt[:120]!r}", flush=True)
    print(f"[infer] negative_prompt={negative_prompt[:120]!r}", flush=True)


def print_stage_start(stage: str) -> None:
    print(f"[{stage}] start", flush=True)


def print_stage_done(stage: str, elapsed: float) -> None:
    print(f"[{stage}] done — {elapsed:.1f}s", flush=True)


def print_stage_error(stage: str, error: Exception) -> None:
    print(f"[{stage}] ERROR: {type(error).__name__}: {error}", flush=True)


def print_infer_done(elapsed: float, output_fps: int, frame_count: int) -> None:
    print(f"[infer] ===== END t={elapsed:.1f}s ===== {frame_count} frames @ {output_fps}fps", flush=True)


def print_infer_error(error: Exception, elapsed: float) -> None:
    print(f"[infer] FAILED: {type(error).__name__}: {error} | t={elapsed:.1f}s", flush=True)


def print_frames_info(label: str, frames: Any, fps: Any) -> None:
    if frames is None:
        print(f"[export] {label}: frames=None, fps={fps}", flush=True)
        return
    n = len(frames)
    first = frames[0] if n else None
    kind = type(first).__name__ if first is not None else "n/a"
    shape = getattr(first, "shape", None)
    dtype = getattr(first, "dtype", None)
    # PIL Image exposes .size/.mode as plain attributes; torch.Tensor happens to have same-named
    # *methods* (Tensor.size(), Tensor.mode()) — skip those so callers passing tensors (e.g. the
    # RIFE->upscale GPU handoff) don't log a garbage bound-method repr in their place.
    size = getattr(first, "size", None)
    size = size if not callable(size) else None
    mode = getattr(first, "mode", None)
    mode = mode if not callable(mode) else None
    print(
        f"[export] {label}: n={n}, fps={fps}, frame_type={kind}, shape={shape}, dtype={dtype}, "
        f"pil_size={size}, pil_mode={mode}",
        flush=True,
    )


def print_export_start(video_path: str) -> None:
    print(f"[export] start -> {video_path}", flush=True)


def print_export_done(elapsed: float) -> None:
    print(f"[export] done — {elapsed:.1f}s", flush=True)


def print_export_error(error: Exception, tb: str) -> None:
    print(f"[export] FAILED: {type(error).__name__}: {error}\n{tb}", flush=True)


def print_log_skipped(has_token: bool, has_repo: bool) -> None:
    print(f"[log] skipped — token={'set' if has_token else 'missing'}, repo={'set' if has_repo else 'missing'}")


def print_log_queued(stem: str, n_files: int, pending: int) -> None:
    print(f"[log] queued {stem} ({n_files} file(s), pending={pending})")


def print_log_inference_warning(e: Exception, tb: str) -> None:
    print(f"[log] WARNING: {e}\n{tb}")


def print_log_inference_total(elapsed: float) -> None:
    print(f"[log] log_inference total: {elapsed:.3f}s")


def print_log_batch_upload_warning(e: Exception) -> None:
    print(f"[log] batch upload warning: {e}")


def print_log_committed(n_files: int, n_pruned: int) -> None:
    print(f"[log] committed {n_files} file(s), pruned {n_pruned}")


def print_log_list_existing_files_failed(e: Exception) -> None:
    print(f"[log] could not list existing files (empty repo?): {e}")


def print_log_squash_marker_not_found(e: Exception) -> None:
    print(f"[log] squash marker not found ({e}), proceeding with squash")


def print_log_squashed_history(repo_id: str) -> None:
    print(f"[log] squashed history for {repo_id}")


def print_log_squash_warning(e: Exception) -> None:
    print(f"[log] squash warning: {e}")


def _path_struct() -> Any:
    import pyarrow as pa
    return pa.struct([("bytes", pa.binary()), ("path", pa.string())])


def _path_value(path_in_repo: str | None) -> dict[str, Any]:
    return {"bytes": None, "path": path_in_repo}


def _build_table(image_path_in_repo: str | None, video_path_in_repo: str | None, prompt: str,
                  negative_prompt: str, seed: int, steps: int, guidance_scale: float,
                  interpolation_enabled: bool, interpolation_multiplier: int, upscale_enabled: bool,
                  output_width: int | None, output_height: int | None, output_fps: int | None,
                  output_duration_seconds: float | None, generation_duration_seconds: float,
                  success: bool, error_message: str, now: datetime) -> Any:
    import pyarrow as pa

    media_struct = _path_struct()
    hf_meta = json.dumps({"info": {"features": {
        "timestamp":                    {"dtype": "float64", "_type": "Value"},
        "prompt":                       {"dtype": "string",  "_type": "Value"},
        "negative_prompt":              {"dtype": "string",  "_type": "Value"},
        "seed":                         {"dtype": "int32",   "_type": "Value"},
        "steps":                        {"dtype": "int32",   "_type": "Value"},
        "guidance_scale":               {"dtype": "float32", "_type": "Value"},
        "input_image":                  {"_type": "Image"},
        "interpolation_enabled":        {"dtype": "bool",    "_type": "Value"},
        "interpolation_multiplier":     {"dtype": "int32",   "_type": "Value"},
        "upscale_enabled":              {"dtype": "bool",    "_type": "Value"},
        "output_video":                 {"_type": "Video"},
        "output_width":                 {"dtype": "int32",   "_type": "Value"},
        "output_height":                {"dtype": "int32",   "_type": "Value"},
        "output_fps":                   {"dtype": "int32",   "_type": "Value"},
        "output_duration_seconds":      {"dtype": "float32", "_type": "Value"},
        "generation_duration_seconds":  {"dtype": "float32", "_type": "Value"},
        "success":                      {"dtype": "bool",    "_type": "Value"},
        "error_message":                {"dtype": "string",  "_type": "Value"},
    }}}).encode()

    schema = pa.schema([
        ("timestamp",                   pa.float64()),
        ("prompt",                      pa.string()),
        ("negative_prompt",             pa.string()),
        ("seed",                        pa.int32()),
        ("steps",                       pa.int32()),
        ("guidance_scale",              pa.float32()),
        ("input_image",                 media_struct),
        ("interpolation_enabled",       pa.bool_()),
        ("interpolation_multiplier",    pa.int32()),
        ("upscale_enabled",             pa.bool_()),
        ("output_video",                media_struct),
        ("output_width",                pa.int32()),
        ("output_height",               pa.int32()),
        ("output_fps",                  pa.int32()),
        ("output_duration_seconds",     pa.float32()),
        ("generation_duration_seconds", pa.float32()),
        ("success",                     pa.bool_()),
        ("error_message",               pa.string()),
    ], metadata={b"huggingface": hf_meta})

    def _opt_i32(v: int | None) -> Any:
        return pa.array([v], type=pa.int32())

    def _opt_f32(v: float | None) -> Any:
        return pa.array([v], type=pa.float32())

    return pa.table({
        "timestamp":                   pa.array([now.timestamp()], type=pa.float64()),
        "prompt":                      pa.array([prompt], type=pa.string()),
        "negative_prompt":             pa.array([negative_prompt], type=pa.string()),
        "seed":                        pa.array([int(seed)], type=pa.int32()),
        "steps":                       pa.array([int(steps)], type=pa.int32()),
        "guidance_scale":              pa.array([float(guidance_scale)], type=pa.float32()),
        "input_image":                 pa.array([_path_value(image_path_in_repo)], type=media_struct),
        "interpolation_enabled":       pa.array([bool(interpolation_enabled)], type=pa.bool_()),
        "interpolation_multiplier":    pa.array([int(interpolation_multiplier)], type=pa.int32()),
        "upscale_enabled":             pa.array([bool(upscale_enabled)], type=pa.bool_()),
        "output_video":                pa.array([_path_value(video_path_in_repo)], type=media_struct),
        "output_width":                _opt_i32(output_width),
        "output_height":               _opt_i32(output_height),
        "output_fps":                  _opt_i32(output_fps),
        "output_duration_seconds":     _opt_f32(output_duration_seconds),
        "generation_duration_seconds": pa.array([float(generation_duration_seconds)], type=pa.float32()),
        "success":                     pa.array([bool(success)], type=pa.bool_()),
        "error_message":               pa.array([str(error_message)], type=pa.string()),
    }, schema=schema)


def _make_stem(now: datetime, uid: str) -> str:
    return f"{now.strftime('%Y-%m-%d-%H%M%S')}-{uid}"


def _write_temp_jpeg(image: PILImage, quality: int = 85) -> str:
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        path = tmp.name
    image.convert("RGB").save(path, format="JPEG", quality=quality)
    return path


def _export_temp_video(frames: list[Any], fps: int) -> str:
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        path = tmp.name
    export_to_video(frames, path, fps=fps, quality=6)
    return path


def _write_parquet(table: Any) -> str:
    import pyarrow.parquet as pq
    with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
        path = tmp.name
    pq.write_table(table, path)
    return path


_LOGGED_PREFIXES = ("data/", "images/", "videos/")


def _list_existing_files_with_sizes(api: "HfApi", repo_id: str) -> list[tuple[str, int]]:
    try:
        entries = list(api.list_repo_tree(repo_id, repo_type="dataset", recursive=True))
    except Exception as e:
        print_log_list_existing_files_failed(e)
        return []
    return [(f.path, getattr(f, "size", 0) or 0) for f in entries if f.path.startswith(_LOGGED_PREFIXES)]


def _stem_of(path: str) -> str:
    # "data/2026-08-29-120000-abcd1234.parquet" -> "2026-08-29-120000-abcd1234"
    name = path.split("/", 1)[1] if "/" in path else path
    return name.rsplit(".", 1)[0]


def _group_by_stem(existing: list[tuple[str, int]]) -> dict[str, dict[str, Any]]:
    groups: dict[str, dict[str, Any]] = {}
    for path, size in existing:
        stem = _stem_of(path)
        group = groups.setdefault(stem, {"paths": [], "size": 0})
        group["paths"].append(path)
        group["size"] += size
    return groups


def _build_delete_ops(existing: list[tuple[str, int]], new_batch: list[tuple[str, str]],
                       max_bytes: int, max_files: int) -> list[CommitOperationDelete]:
    groups = _group_by_stem(existing)
    new_stems = {_stem_of(p) for p, _ in new_batch}
    total_bytes = sum(g["size"] for g in groups.values()) + sum(os.path.getsize(local) for _, local in new_batch)
    total_files = len(groups) + len(new_stems)

    def _over_cap() -> bool:
        return (max_bytes > 0 and total_bytes > max_bytes) or (max_files > 0 and total_files > max_files)

    if not _over_cap():
        return []
    ops: list[CommitOperationDelete] = []
    for stem in sorted(groups):
        if not _over_cap():
            break
        group = groups[stem]
        ops.extend(CommitOperationDelete(path_in_repo=p) for p in group["paths"])
        total_bytes -= group["size"]
        total_files -= 1
    return ops


def _squash_if_needed(api: "HfApi", repo_id: str) -> None:
    marker = "metadata/last_squash.txt"
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    try:
        try:
            local = hf_hub_download(repo_id=repo_id, filename=marker, repo_type="dataset", token=api.token)
            if open(local).read().strip() == today:
                return
        except Exception as e:
            print_log_squash_marker_not_found(e)
        api.super_squash_history(repo_id=repo_id, repo_type="dataset")
        api.upload_file(path_or_fileobj=today.encode(), path_in_repo=marker, repo_id=repo_id, repo_type="dataset")
        print_log_squashed_history(repo_id)
    except Exception as e:
        print_log_squash_warning(e)


def _delete_temp_files(paths: list[str]) -> None:
    for path in paths:
        try:
            os.unlink(path)
        except Exception:
            pass


class LogUploader:
    def __init__(self, token: str | None, repo_id: str | None, max_bytes: int, max_files: int,
                 batch_interval: int = 60) -> None:
        self._token = token
        self._repo_id = repo_id
        self._max_bytes = max_bytes
        self._max_files = max_files
        self._batch_interval = batch_interval
        self._pending: list[tuple[str, str]] = []
        self._lock = threading.Lock()
        if self.enabled:
            threading.Thread(target=self._loop, daemon=True, name="log-uploader").start()

    @property
    def enabled(self) -> bool:
        return bool(self._token and self._repo_id)

    def log_inference(self, input_image: PILImage | None, output_frames: list[Any] | None,
                       output_fps: int | None, prompt: str, negative_prompt: str, seed: int, steps: int,
                       guidance_scale: float, interpolation_enabled: bool, interpolation_multiplier: int,
                       upscale_enabled: bool, output_width: int | None, output_height: int | None,
                       generation_duration_seconds: float, success: bool, error_message: str = "") -> None:
        if not self.enabled:
            print_log_skipped(bool(self._token), bool(self._repo_id))
            return
        t0 = _time.perf_counter()
        local_files: list[str] = []
        try:
            now = datetime.now(timezone.utc)
            stem = _make_stem(now, uuid.uuid4().hex[:8])
            batch: list[tuple[str, str]] = []

            image_path_in_repo = None
            if input_image is not None:
                local_jpeg = _write_temp_jpeg(input_image)
                local_files.append(local_jpeg)
                image_path_in_repo = f"images/{stem}.jpg"
                batch.append((image_path_in_repo, local_jpeg))

            video_path_in_repo = None
            output_duration_seconds = None
            if output_frames is not None and output_fps:
                local_mp4 = _export_temp_video(output_frames, output_fps)
                local_files.append(local_mp4)
                video_path_in_repo = f"videos/{stem}.mp4"
                output_duration_seconds = len(output_frames) / output_fps
                batch.append((video_path_in_repo, local_mp4))

            table = _build_table(
                image_path_in_repo, video_path_in_repo, prompt, negative_prompt, seed, steps,
                guidance_scale, interpolation_enabled, interpolation_multiplier, upscale_enabled,
                output_width, output_height, output_fps, output_duration_seconds,
                generation_duration_seconds, success, error_message, now,
            )
            local_parquet = _write_parquet(table)
            local_files.append(local_parquet)
            batch.insert(0, (f"data/{stem}.parquet", local_parquet))

            self._enqueue_many(batch)
            print_log_queued(stem, len(batch), len(self._pending))
        except Exception as e:
            import traceback as _tb
            print_log_inference_warning(e, _tb.format_exc())
            _delete_temp_files(local_files)
        print_log_inference_total(_time.perf_counter() - t0)

    def _enqueue_many(self, files: list[tuple[str, str]]) -> None:
        with self._lock:
            self._pending.extend(files)

    def _drain(self) -> list[tuple[str, str]]:
        with self._lock:
            batch = self._pending[:]
            self._pending.clear()
        return batch

    def _requeue(self, batch: list[tuple[str, str]]) -> None:
        with self._lock:
            self._pending[:0] = batch

    def _loop(self) -> None:
        while True:
            _time.sleep(self._batch_interval)
            self._flush()

    def _flush(self) -> None:
        batch = self._drain()
        if not batch:
            return
        try:
            self._commit_batch(batch)
            _delete_temp_files([local for _, local in batch])
        except Exception as e:
            print_log_batch_upload_warning(e)
            self._requeue(batch)

    def _commit_batch(self, batch: list[tuple[str, str]]) -> None:
        from huggingface_hub import HfApi
        assert self._repo_id is not None
        api = HfApi(token=self._token)
        api.create_repo(repo_id=self._repo_id, repo_type="dataset", private=True, exist_ok=True)
        existing = _list_existing_files_with_sizes(api, self._repo_id)
        add_ops = [CommitOperationAdd(path_in_repo=p, path_or_fileobj=local) for p, local in batch]
        del_ops = _build_delete_ops(existing, batch, self._max_bytes, self._max_files)
        api.create_commit(
            repo_id=self._repo_id, repo_type="dataset",
            operations=[*add_ops, *del_ops],
            commit_message=f"[log] batch {len(batch)} file(s)" + (f", prune {len(del_ops)}" if del_ops else ""),
        )
        print_log_committed(len(batch), len(del_ops))
        _squash_if_needed(api, self._repo_id)