utils/usage_logger.py

"""
Usage logger that pushes alignment runs to a HF Dataset repo.

Uses a ParquetScheduler (subclass of CommitScheduler) to buffer rows in memory
and periodically write+upload parquet files with embedded audio to the Hub.
Error logs use a separate CommitScheduler with JSONL files.
Falls back to local-only logging if schedulers can't initialize.

Scheduler creation is deferred to first use so that background threads don't
interfere with ZeroGPU's startup function scan.
"""

import hashlib
import io
import json
import threading
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
from uuid import uuid4

import numpy as np

# =========================================================================
# Directory setup
# =========================================================================

LOG_DIR = Path("usage_logs")
LOG_DIR.mkdir(parents=True, exist_ok=True)

ERROR_DIR = LOG_DIR / "errors"
ERROR_DIR.mkdir(parents=True, exist_ok=True)

ERROR_LOG_PATH = ERROR_DIR / f"error_log-{uuid4()}.jsonl"

# =========================================================================
# ParquetScheduler class definition (no instances created at import time)
# =========================================================================

_HAS_DEPS = False
try:
    import pyarrow as pa
    import pyarrow.parquet as pq
    from huggingface_hub import CommitScheduler
    from config import USAGE_LOG_DATASET_REPO, USAGE_LOG_PUSH_INTERVAL_MINUTES

    _HAS_DEPS = True
except Exception:
    pass

# HF features schema (column order matters — audio first for HF viewer widget)
_ALIGNER_SCHEMA: Dict[str, Dict[str, str]] = {
    # Identity
    "audio": {"_type": "Audio"},
    "audio_id": {"_type": "Value", "dtype": "string"},
    "timestamp": {"_type": "Value", "dtype": "string"},
    "user_id": {"_type": "Value", "dtype": "string"},
    # Input metadata
    "audio_duration_s": {"_type": "Value", "dtype": "float64"},
    "num_segments": {"_type": "Value", "dtype": "int32"},
    "surah": {"_type": "Value", "dtype": "int32"},
    # Segmentation settings
    "min_silence_ms": {"_type": "Value", "dtype": "int32"},
    "min_speech_ms": {"_type": "Value", "dtype": "int32"},
    "pad_ms": {"_type": "Value", "dtype": "int32"},
    "asr_model": {"_type": "Value", "dtype": "string"},
    "device": {"_type": "Value", "dtype": "string"},
    # Profiling
    "total_time": {"_type": "Value", "dtype": "float64"},
    "vad_queue_time": {"_type": "Value", "dtype": "float64"},
    "vad_gpu_time": {"_type": "Value", "dtype": "float64"},
    "asr_gpu_time": {"_type": "Value", "dtype": "float64"},
    "dp_total_time": {"_type": "Value", "dtype": "float64"},
    # Quality & retry
    "segments_passed": {"_type": "Value", "dtype": "int32"},
    "segments_failed": {"_type": "Value", "dtype": "int32"},
    "mean_confidence": {"_type": "Value", "dtype": "float64"},
    "tier1_retries": {"_type": "Value", "dtype": "int32"},
    "tier1_passed": {"_type": "Value", "dtype": "int32"},
    "tier2_retries": {"_type": "Value", "dtype": "int32"},
    "tier2_passed": {"_type": "Value", "dtype": "int32"},
    "reanchors": {"_type": "Value", "dtype": "int32"},
    "special_merges": {"_type": "Value", "dtype": "int32"},
    # Reciter stats
    "words_per_minute": {"_type": "Value", "dtype": "float64"},
    "phonemes_per_second": {"_type": "Value", "dtype": "float64"},
    "avg_segment_duration": {"_type": "Value", "dtype": "float64"},
    "std_segment_duration": {"_type": "Value", "dtype": "float64"},
    "avg_pause_duration": {"_type": "Value", "dtype": "float64"},
    "std_pause_duration": {"_type": "Value", "dtype": "float64"},
    # Session flags
    "resegmented": {"_type": "Value", "dtype": "bool"},
    "retranscribed": {"_type": "Value", "dtype": "bool"},
    # Segments, timestamps & error
    "segments": {"_type": "Value", "dtype": "string"},
    "word_timestamps": {"_type": "Value", "dtype": "string"},
    "char_timestamps": {"_type": "Value", "dtype": "string"},
    "error": {"_type": "Value", "dtype": "string"},
}

if _HAS_DEPS:
    class ParquetScheduler(CommitScheduler):
        """Buffers rows in memory and uploads a parquet file each interval.

        Audio values are stored as file paths in the row dict; on push they are
        read as bytes and embedded in the parquet using the HF Audio struct.
        """

        def __init__(
            self,
            *,
            repo_id: str,
            schema: Optional[Dict[str, Dict[str, str]]] = None,
            every: Union[int, float] = 5,
            path_in_repo: Optional[str] = "data",
            repo_type: Optional[str] = "dataset",
            private: bool = False,
        ) -> None:
            super().__init__(
                repo_id=repo_id,
                folder_path="dummy",  # not used — we upload directly
                every=every,
                path_in_repo=path_in_repo,
                repo_type=repo_type,
                private=private,
            )
            self._rows: List[Dict[str, Any]] = []
            self._schema = schema

        def append(self, row: Dict[str, Any]) -> None:
            with self.lock:
                self._rows.append(row)

        def push_to_hub(self) -> None:
            with self.lock:
                rows = self._rows
                self._rows = []
            if not rows:
                return

            print(f"[USAGE_LOG] Pushing {len(rows)} alignment row(s) to Hub.")

            schema: Dict[str, Dict] = dict(self._schema) if self._schema else {}
            paths_to_cleanup: List[Path] = []

            for row in rows:
                for key, value in row.items():
                    if key not in schema:
                        schema[key] = _infer_schema(key, value)

                    if value is not None and schema[key].get("_type") in ("Image", "Audio"):
                        file_path = Path(value)
                        if file_path.is_file():
                            row[key] = {
                                "path": file_path.name,
                                "bytes": file_path.read_bytes(),
                            }
                            paths_to_cleanup.append(file_path)
                        else:
                            row[key] = None

            for row in rows:
                for feature in schema:
                    if feature not in row:
                        row[feature] = None

            table = pa.Table.from_pylist(rows)
            table = table.replace_schema_metadata(
                {"huggingface": json.dumps({"info": {"features": schema}})}
            )

            archive = None
            try:
                import tempfile
                archive = tempfile.NamedTemporaryFile(suffix=".parquet", delete=False)
                pq.write_table(table, archive.name)
                self.api.upload_file(
                    repo_id=self.repo_id,
                    repo_type=self.repo_type,
                    revision=self.revision,
                    path_in_repo=f"{self.path_in_repo}/{uuid4()}.parquet",
                    path_or_fileobj=archive.name,
                )
                print("[USAGE_LOG] Parquet commit completed.")
            except Exception as e:
                print(f"[USAGE_LOG] Failed to upload parquet: {e}")
            finally:
                if archive:
                    archive.close()
                    Path(archive.name).unlink(missing_ok=True)

            for path in paths_to_cleanup:
                path.unlink(missing_ok=True)

    def _infer_schema(key: str, value: Any) -> Dict[str, str]:
        if "image" in key:
            return {"_type": "Image"}
        if "audio" in key:
            return {"_type": "Audio"}
        if isinstance(value, bool):
            return {"_type": "Value", "dtype": "bool"}
        if isinstance(value, int):
            return {"_type": "Value", "dtype": "int64"}
        if isinstance(value, float):
            return {"_type": "Value", "dtype": "float64"}
        if isinstance(value, bytes):
            return {"_type": "Value", "dtype": "binary"}
        return {"_type": "Value", "dtype": "string"}


# =========================================================================
# Lazy scheduler initialization (deferred to first use)
# =========================================================================

_aligner_scheduler = None
_error_scheduler = None
_schedulers_initialized = False
_init_lock = threading.Lock()
_fallback_lock = threading.Lock()


def _ensure_schedulers() -> None:
    global _aligner_scheduler, _error_scheduler, _schedulers_initialized
    if _schedulers_initialized:
        return
    with _init_lock:
        if _schedulers_initialized:
            return
        _schedulers_initialized = True
        if not _HAS_DEPS:
            print("[USAGE_LOG] Dependencies missing (local-only mode).")
            return
        try:
            _aligner_scheduler = ParquetScheduler(
                repo_id=USAGE_LOG_DATASET_REPO,
                schema=_ALIGNER_SCHEMA,
                every=USAGE_LOG_PUSH_INTERVAL_MINUTES,
                path_in_repo="data",
                repo_type="dataset",
                private=True,
            )
            _error_scheduler = CommitScheduler(
                repo_id=USAGE_LOG_DATASET_REPO,
                repo_type="dataset",
                folder_path=ERROR_DIR,
                path_in_repo="data/errors",
                private=True,
                every=USAGE_LOG_PUSH_INTERVAL_MINUTES,
            )
        except Exception as e:
            print(f"[USAGE_LOG] Scheduler init failed (local-only mode): {e}")


# =========================================================================
# Helpers
# =========================================================================


def _get_error_lock():
    _ensure_schedulers()
    if _error_scheduler is not None:
        return _error_scheduler.lock
    return _fallback_lock


def get_user_id(request) -> str:
    """SHA-256 hash (12-char) of IP+UA from a gr.Request, or 'unknown'."""
    try:
        headers = request.headers
        ip = (
            headers.get("x-forwarded-for", "").split(",")[0].strip()
            or headers.get("x-real-ip", "")
            or ""
        )
        ua = headers.get("user-agent", "")
        return hashlib.sha256(f"{ip}|{ua}".encode()).hexdigest()[:12]
    except Exception:
        return "unknown"


def _compute_audio_id(audio: np.ndarray, ts: datetime) -> str:
    """Content hash (16-char) + compact timestamp."""
    audio_hash = hashlib.sha256(audio.tobytes()).hexdigest()[:16]
    return f"{audio_hash}:{ts.strftime('%Y%m%dT%H%M%S')}"


def _encode_audio_flac(audio: np.ndarray, sample_rate: int, audio_id: str) -> str:
    """Encode audio to a temp FLAC file; returns the file path."""
    import soundfile as sf

    tmp_dir = LOG_DIR / "tmp_audio"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    safe_id = audio_id.replace(":", "-")
    filepath = tmp_dir / f"{safe_id}.flac"
    sf.write(str(filepath), audio, sample_rate, format="FLAC")
    return str(filepath)


def _sync_row_to_scheduler(row: Dict[str, Any]) -> None:
    """Ensure *row* is represented in the scheduler buffer.

    gr.State may deserialize the dict (creating a copy), and push_to_hub
    detaches rows from the buffer.  This helper finds the original row by
    audio_id and updates it, or re-appends if it was already pushed.
    """
    if _aligner_scheduler is None:
        return
    audio_id = row.get("audio_id")
    if not audio_id:
        return
    with _aligner_scheduler.lock:
        for buffered in _aligner_scheduler._rows:
            if buffered.get("audio_id") == audio_id:
                # Update the buffered row in-place (handles gr.State copies)
                buffered.update(row)
                return
        # Row was already pushed — re-append (audio file may be gone, that's ok)
        _aligner_scheduler._rows.append(row)


# =========================================================================
# Public logging API
# =========================================================================


def log_alignment(
    *,
    audio: np.ndarray,
    sample_rate: int,
    request=None,
    # Input metadata
    audio_duration_s: float,
    num_segments: int,
    surah: int,
    # Settings
    min_silence_ms: int,
    min_speech_ms: int,
    pad_ms: int,
    asr_model: str,
    device: str,
    # Profiling
    total_time: float,
    vad_queue_time: float,
    vad_gpu_time: float,
    asr_gpu_time: float,
    dp_total_time: float,
    # Quality & retry
    segments_passed: int,
    segments_failed: int,
    mean_confidence: float,
    tier1_retries: int,
    tier1_passed: int,
    tier2_retries: int,
    tier2_passed: int,
    reanchors: int,
    special_merges: int,
    # Reciter stats
    words_per_minute: float,
    phonemes_per_second: float,
    avg_segment_duration: float,
    std_segment_duration: float,
    avg_pause_duration: float,
    std_pause_duration: float,
    # Segments
    log_segments: List[dict],
) -> Optional[Dict[str, Any]]:
    """Log an alignment run. Returns the row dict reference for in-place mutation.

    The returned dict can be stored in gr.State and mutated on
    resegment/retranscribe/timestamps before the scheduler pushes.
    """
    _ensure_schedulers()
    try:
        ts = datetime.now()
        audio_id = _compute_audio_id(audio, ts)
        user_id = get_user_id(request) if request else "unknown"

        # Build the segments JSON: array of run objects
        segments_runs = [{
            "min_silence_ms": int(min_silence_ms),
            "min_speech_ms": int(min_speech_ms),
            "pad_ms": int(pad_ms),
            "asr_model": asr_model,
            "segments": log_segments,
        }]

        # Encode audio to FLAC temp file (scheduler embeds bytes on push)
        audio_path = _encode_audio_flac(audio, sample_rate, audio_id)

        row: Dict[str, Any] = {
            "audio": audio_path,
            "audio_id": audio_id,
            "timestamp": ts.isoformat(timespec="seconds"),
            "user_id": user_id,
            # Input metadata
            "audio_duration_s": audio_duration_s,
            "num_segments": num_segments,
            "surah": surah,
            # Settings (latest)
            "min_silence_ms": int(min_silence_ms),
            "min_speech_ms": int(min_speech_ms),
            "pad_ms": int(pad_ms),
            "asr_model": asr_model,
            "device": device,
            # Profiling
            "total_time": total_time,
            "vad_queue_time": vad_queue_time,
            "vad_gpu_time": vad_gpu_time,
            "asr_gpu_time": asr_gpu_time,
            "dp_total_time": dp_total_time,
            # Quality & retry
            "segments_passed": segments_passed,
            "segments_failed": segments_failed,
            "mean_confidence": mean_confidence,
            "tier1_retries": tier1_retries,
            "tier1_passed": tier1_passed,
            "tier2_retries": tier2_retries,
            "tier2_passed": tier2_passed,
            "reanchors": reanchors,
            "special_merges": special_merges,
            # Reciter stats
            "words_per_minute": words_per_minute,
            "phonemes_per_second": phonemes_per_second,
            "avg_segment_duration": avg_segment_duration,
            "std_segment_duration": std_segment_duration,
            "avg_pause_duration": avg_pause_duration,
            "std_pause_duration": std_pause_duration,
            # Session flags
            "resegmented": False,
            "retranscribed": False,
            # Segments & error
            "segments": json.dumps(segments_runs),
            "word_timestamps": None,
            "char_timestamps": None,
            "error": None,
        }

        if _aligner_scheduler is not None:
            _aligner_scheduler.append(row)
        else:
            _write_fallback(row)

        return row

    except Exception as e:
        print(f"[USAGE_LOG] Failed to log alignment: {e}")
        return None


def update_alignment_row(
    row: Dict[str, Any],
    *,
    action: str,
    # Input metadata (overwritten)
    audio_duration_s: float,
    num_segments: int,
    surah: int,
    # Settings for this run
    min_silence_ms: int,
    min_speech_ms: int,
    pad_ms: int,
    asr_model: str,
    device: str,
    # Profiling
    total_time: float,
    vad_queue_time: float,
    vad_gpu_time: float,
    asr_gpu_time: float,
    dp_total_time: float,
    # Quality & retry
    segments_passed: int,
    segments_failed: int,
    mean_confidence: float,
    tier1_retries: int,
    tier1_passed: int,
    tier2_retries: int,
    tier2_passed: int,
    reanchors: int,
    special_merges: int,
    # Reciter stats
    words_per_minute: float,
    phonemes_per_second: float,
    avg_segment_duration: float,
    std_segment_duration: float,
    avg_pause_duration: float,
    std_pause_duration: float,
    # Segments
    log_segments: List[dict],
) -> None:
    """Mutate an existing row dict in-place and ensure it's in the scheduler buffer.

    After mutation, syncs the row into the scheduler's buffer so the update
    is captured even if gr.State returned a deserialized copy or if the
    original row was already pushed to Hub.

    Args:
        row: The dict returned by log_alignment(), stored in gr.State.
        action: "resegment" or "retranscribe".
    """
    try:
        # Overwrite run-level fields
        row["audio_duration_s"] = audio_duration_s
        row["num_segments"] = num_segments
        row["surah"] = surah
        row["min_silence_ms"] = int(min_silence_ms)
        row["min_speech_ms"] = int(min_speech_ms)
        row["pad_ms"] = int(pad_ms)
        row["asr_model"] = asr_model
        row["device"] = device
        row["total_time"] = total_time
        row["vad_queue_time"] = vad_queue_time
        row["vad_gpu_time"] = vad_gpu_time
        row["asr_gpu_time"] = asr_gpu_time
        row["dp_total_time"] = dp_total_time
        row["segments_passed"] = segments_passed
        row["segments_failed"] = segments_failed
        row["mean_confidence"] = mean_confidence
        row["tier1_retries"] = tier1_retries
        row["tier1_passed"] = tier1_passed
        row["tier2_retries"] = tier2_retries
        row["tier2_passed"] = tier2_passed
        row["reanchors"] = reanchors
        row["special_merges"] = special_merges
        row["words_per_minute"] = words_per_minute
        row["phonemes_per_second"] = phonemes_per_second
        row["avg_segment_duration"] = avg_segment_duration
        row["std_segment_duration"] = std_segment_duration
        row["avg_pause_duration"] = avg_pause_duration
        row["std_pause_duration"] = std_pause_duration

        # Set session flag
        if action == "resegment":
            row["resegmented"] = True
        elif action == "retranscribe":
            row["retranscribed"] = True

        # Append new run to segments array
        segments_runs = json.loads(row.get("segments") or "[]")
        segments_runs.append({
            "min_silence_ms": int(min_silence_ms),
            "min_speech_ms": int(min_speech_ms),
            "pad_ms": int(pad_ms),
            "asr_model": asr_model,
            "segments": log_segments,
        })
        row["segments"] = json.dumps(segments_runs)

        # Sync with scheduler buffer — the row from gr.State may be a
        # deserialized copy, or the original may have already been pushed.
        _sync_row_to_scheduler(row)

    except Exception as e:
        print(f"[USAGE_LOG] Failed to update alignment row: {e}")


def update_word_timestamps(
    row: Dict[str, Any],
    word_timestamps_json: str,
    char_timestamps_json: Optional[str] = None,
) -> None:
    """Set word and char timestamps fields on an existing row and sync to scheduler."""
    try:
        row["word_timestamps"] = word_timestamps_json
        if char_timestamps_json is not None:
            row["char_timestamps"] = char_timestamps_json
        _sync_row_to_scheduler(row)
    except Exception as e:
        print(f"[USAGE_LOG] Failed to update word timestamps: {e}")


def log_error(user_id: str, error_message: str) -> None:
    """Log a pipeline error to JSONL."""
    try:
        with _get_error_lock():
            with ERROR_LOG_PATH.open("a") as f:
                json.dump({
                    "timestamp": datetime.now().isoformat(timespec="seconds"),
                    "user_id": user_id,
                    "error_message": error_message or "",
                }, f)
                f.write("\n")
    except Exception:
        pass


def _write_fallback(row: Dict[str, Any]) -> None:
    """Local-only fallback: write JSONL (without audio)."""
    fallback_path = LOG_DIR / "alignments_fallback.jsonl"
    with _fallback_lock:
        with fallback_path.open("a") as f:
            fallback_row = {k: v for k, v in row.items() if k != "audio"}
            json.dump(fallback_row, f)
            f.write("\n")