voice-transcriptor/src/voice_transcriptor/services/transcription.py

from __future__ import annotations

import random
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import openai

from voice_transcriptor.models import AppSettings
from voice_transcriptor.services.job_manifest import ChunkStatus, JobManifestRepository
from voice_transcriptor.services.preprocessing import CancellationToken, PreprocessingCancelled


BRAZILIAN_PROMPT = (
    "Conversa em português brasileiro. Preserve a língua falada; não traduza. "
    "Preserve ortografia e pontuação brasileiras, números, nomes próprios, "
    "termos técnicos e siglas com máxima fidelidade."
)


class TranscriptionError(RuntimeError): pass
class PermanentTranscriptionError(TranscriptionError): pass
class RetryExhaustedError(TranscriptionError): pass


@dataclass(frozen=True, slots=True)
class RetryPolicy:
    max_attempts: int = 5
    initial_delay_seconds: float = 1.0
    max_delay_seconds: float = 30.0
    jitter_ratio: float = 0.2


@dataclass(frozen=True, slots=True)
class TranscriptionProgress:
    completed: int
    total: int
    current_chunk: int | None
    elapsed_seconds: float
    phase: str
    message: str
    api_error: str | None = None


def normalize_language(value: str) -> str:
    normalized = value.strip().replace("_", "-")
    return normalized.split("-", 1)[0].lower()


def build_prompt(context: str) -> str:
    extra = context.strip()
    return f"{BRAZILIAN_PROMPT}\n\n{extra}" if extra else BRAZILIAN_PROMPT


class OpenAITranscriptionClient:
    def __init__(self, client) -> None: self.client = client

    def transcribe(self, path: Path, model: str, language: str, prompt: str) -> str:
        with path.open("rb") as audio_file:
            response = self.client.audio.transcriptions.create(
                file=audio_file,
                model=model,
                language=normalize_language(language),
                prompt=prompt,
                response_format="json",
            )
        text = getattr(response, "text", None)
        if not isinstance(text, str):
            raise PermanentTranscriptionError("The transcription API returned no text.")
        return text


class TranscriptionService:
    def __init__(
        self,
        client: OpenAITranscriptionClient,
        manifests: JobManifestRepository,
        retry_policy: RetryPolicy | None = None,
        sleep: Callable[[float], None] = time.sleep,
        monotonic: Callable[[], float] = time.monotonic,
        random_value: Callable[[], float] = random.random,
    ) -> None:
        self.client = client; self.manifests = manifests
        self.retry_policy = retry_policy or RetryPolicy()
        self.sleep = sleep; self.monotonic = monotonic; self.random_value = random_value

    def run(
        self,
        manifest_path: Path,
        settings: AppSettings,
        token: CancellationToken | None = None,
        progress: Callable[[TranscriptionProgress], None] | None = None,
    ) -> Path:
        token = token or CancellationToken(); started = self.monotonic()
        try:
            token.raise_if_cancelled()
            manifest = self.manifests.recover_for_resume(manifest_path)
            total = manifest["total_chunks"]
            for item in manifest["chunks"]:
                if item["status"] == ChunkStatus.COMPLETED: continue
                token.raise_if_cancelled(); index = item["index"]
                text = self._transcribe_with_retry(manifest_path, item, settings, token, progress, started, total)
                self.manifests.mark_completed(manifest_path, index, text)
                self.manifests.assemble_transcript(manifest_path)
                completed = self.manifests.load(manifest_path)["completed_chunks"]
                self._emit(progress, completed, total, index + 1, started, "transcribing", f"Completed chunk {index + 1} of {total}")
            self.manifests.mark_job_state(manifest_path, "completed")
            return self.manifests.assemble_transcript(manifest_path)
        except PreprocessingCancelled:
            self.manifests.mark_job_state(manifest_path, "cancelled")
            raise

    def _transcribe_with_retry(self, manifest_path: Path, item: dict, settings: AppSettings, token: CancellationToken, progress, started: float, total: int) -> str:
        policy = self.retry_policy; index = item["index"]
        for attempt in range(1, policy.max_attempts + 1):
            token.raise_if_cancelled(); self.manifests.mark_processing(manifest_path, index)
            self._emit(progress, self.manifests.load(manifest_path)["completed_chunks"], total, index + 1, started, "transcribing", f"Transcribing chunk {index + 1} of {total}")
            try:
                return self.client.transcribe(manifest_path.parent / item["path"], settings.model, settings.language, build_prompt(settings.context_vocabulary))
            except PreprocessingCancelled: raise
            except Exception as exc:
                status = getattr(exc, "status_code", None)
                retryable = self._retryable(exc, status)
                safe = self._safe_error(status, retryable)
                if not retryable:
                    self.manifests.mark_failed(manifest_path, index, safe)
                    raise PermanentTranscriptionError(safe) from None
                if attempt >= policy.max_attempts:
                    self.manifests.mark_failed(manifest_path, index, safe)
                    raise RetryExhaustedError(f"{safe} Retry limit reached.") from None
                delay = self._delay(exc, attempt)
                self._emit(progress, self.manifests.load(manifest_path)["completed_chunks"], total, index + 1, started, "retrying", f"API temporarily unavailable; retrying in {delay:g} seconds.", safe)
                token.raise_if_cancelled(); self.sleep(delay); token.raise_if_cancelled()
        raise AssertionError("unreachable")

    def _delay(self, exc: Exception, attempt: int) -> float:
        headers = getattr(getattr(exc, "response", None), "headers", {}) or {}
        retry_after = headers.get("retry-after") or headers.get("Retry-After")
        try: server_delay = float(retry_after)
        except (TypeError, ValueError): server_delay = 0
        base = max(server_delay, self.retry_policy.initial_delay_seconds * (2 ** (attempt - 1)))
        base = min(base, self.retry_policy.max_delay_seconds)
        jitter = base * self.retry_policy.jitter_ratio * ((self.random_value() * 2) - 1)
        return max(0, min(self.retry_policy.max_delay_seconds, base + jitter))

    @staticmethod
    def _retryable(exc: Exception, status: int | None) -> bool:
        transient_types = (openai.RateLimitError, openai.APIConnectionError, openai.APITimeoutError)
        return isinstance(exc, transient_types) or status in (408, 409, 429) or (isinstance(status, int) and status >= 500)

    @staticmethod
    def _safe_error(status: int | None, retryable: bool) -> str:
        kind = "transient" if retryable else "permanent"
        suffix = f" (HTTP {status})" if isinstance(status, int) else ""
        return f"OpenAI API {kind} error{suffix}."

    def _emit(self, callback, completed: int, total: int, current: int | None, started: float, phase: str, message: str, api_error: str | None = None) -> None:
        if callback: callback(TranscriptionProgress(completed, total, current, max(0, self.monotonic() - started), phase, message, api_error))