from __future__ import annotations

from dataclasses import dataclass
import gc
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
from tempfile import TemporaryDirectory
from typing import Callable, Iterable
import wave


DEFAULT_MODEL_DIR = Path(r"H:\LLM_Models\safetensor\gemma4\E2B-it")
DEFAULT_MODELS_ROOT = DEFAULT_MODEL_DIR.parent
DEFAULT_SAMPLE_RATE = 16_000
DEFAULT_CHUNK_SECONDS = 30
DEFAULT_VIDEO_CHUNK_SECONDS = 60
DEFAULT_TEMPERATURE = 1.0
DEFAULT_TOP_P = 0.95
DEFAULT_TOP_K = 64
AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".flac", ".aac", ".ogg", ".opus"}
VIDEO_EXTENSIONS = {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v"}


os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


@dataclass(frozen=True)
class TranscriberSettings:
    model_dir: Path = DEFAULT_MODEL_DIR
    language: str = "ja"
    mode: str = "speech"
    chunk_seconds: int = DEFAULT_CHUNK_SECONDS
    max_new_tokens: int = 256
    dtype: str = "auto"
    device_map: str = "auto"
    max_memory: str = ""
    ffmpeg_path: str = ""
    sample_rate: int = DEFAULT_SAMPLE_RATE


def lazy_imports():
    try:
        import librosa
        import numpy as np
        import torch
        from transformers import (
            AutoModelForCausalLM,
            AutoModelForMultimodalLM,
            AutoProcessor,
        )
    except ModuleNotFoundError as exc:
        missing = exc.name or "required package"
        raise RuntimeError(
            "Missing dependency: "
            f"{missing}\n"
            "Run `powershell -ExecutionPolicy Bypass -File "
            "D:\\Workspace\\Python\\LLM_GEMMA4_TEST\\setup.ps1` first."
        ) from exc

    return (
        librosa,
        np,
        torch,
        AutoModelForCausalLM,
        AutoModelForMultimodalLM,
        AutoProcessor,
    )


def resolve_torch_dtype(torch_module, dtype_name: str):
    if dtype_name == "auto":
        return "auto"
    return getattr(torch_module, dtype_name)


def parse_max_memory(max_memory_text: str) -> dict[str, str] | None:
    if not max_memory_text.strip():
        return None
    try:
        parsed = json.loads(max_memory_text)
    except json.JSONDecodeError as exc:
        raise RuntimeError(
            "Invalid max_memory JSON. Example: "
            '{"cuda:0":"16GiB","cuda:1":"10GiB","cpu":"32GiB"}'
        ) from exc

    if not isinstance(parsed, dict) or not parsed:
        raise RuntimeError(
            "max_memory must be a JSON object like "
            '{"cuda:0":"16GiB","cuda:1":"10GiB","cpu":"32GiB"}'
        )

    normalized: dict[int | str, str] = {}
    for key, value in parsed.items():
        if not isinstance(key, str) or not isinstance(value, str):
            raise RuntimeError("max_memory keys and values must be strings.")
        lowered = key.lower()
        if lowered.startswith("cuda:"):
            device_index = lowered.split(":", 1)[1]
            if not device_index.isdigit():
                raise RuntimeError(
                    "CUDA device keys must look like cuda:0 or cuda:1."
                )
            normalized[int(device_index)] = value
        elif lowered.isdigit():
            normalized[int(lowered)] = value
        elif lowered in {"cpu", "mps", "disk"}:
            normalized[lowered] = value
        else:
            raise RuntimeError(
                "Unsupported max_memory device key. Use 0, 1, cuda:0, cuda:1, cpu, mps, or disk."
            )
    return normalized


def load_audio(librosa_module, audio_path: Path, sample_rate: int):
    try:
        waveform, actual_sr = librosa_module.load(
            audio_path.as_posix(), sr=sample_rate, mono=True
        )
    except Exception as exc:
        raise RuntimeError(
            "Failed to load audio. Install ffmpeg if MP3 decoding is unavailable.\n"
            "You can also convert the file to WAV and retry.\n"
            f"Details: {exc}"
        ) from exc

    if waveform.size == 0:
        raise RuntimeError("The audio file is empty.")

    return waveform, actual_sr


def is_video_file(path: Path) -> bool:
    return path.suffix.lower() in VIDEO_EXTENSIONS


def extract_audio_from_video(
    media_path: Path,
    output_path: Path,
    sample_rate: int,
    ffmpeg_path: str = "",
) -> None:
    resolved_ffmpeg = resolve_ffmpeg_executable(ffmpeg_path)

    command = [
        resolved_ffmpeg,
        "-y",
        "-i",
        str(media_path),
        "-vn",
        "-ac",
        "1",
        "-ar",
        str(sample_rate),
        str(output_path),
    ]
    try:
        completed = subprocess.run(
            command,
            check=True,
            capture_output=True,
        )
    except FileNotFoundError as exc:
        raise RuntimeError(
            "ffmpeg was not found. Set an ffmpeg.exe path in the UI/CLI, or install it in PATH."
        ) from exc
    except subprocess.CalledProcessError as exc:
        details_bytes = exc.stderr or exc.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "Failed to extract audio from the video file with ffmpeg.\n"
            f"Details: {details}"
        ) from exc

    if not output_path.exists() or output_path.stat().st_size == 0:
        details_bytes = completed.stderr or completed.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "ffmpeg finished but no audio file was created.\n"
            f"Details: {details}"
        )


def resolve_ffmpeg_executable(ffmpeg_path: str = "") -> str:
    ffmpeg_candidates: list[str] = []
    if ffmpeg_path.strip():
        ffmpeg_candidates.append(ffmpeg_path.strip())
    ffmpeg_candidates.extend(
        [
            "ffmpeg",
            r"C:\ffmpeg\bin\ffmpeg.exe",
            r"C:\Program Files\ffmpeg\bin\ffmpeg.exe",
            r"C:\ProgramData\chocolatey\bin\ffmpeg.exe",
            str(Path.home() / "scoop" / "apps" / "ffmpeg" / "current" / "bin" / "ffmpeg.exe"),
        ]
    )

    for candidate in ffmpeg_candidates:
        if candidate == "ffmpeg":
            located = shutil.which("ffmpeg")
            if located:
                return located
            continue

        expanded = Path(candidate).expanduser()
        if expanded.exists():
            return str(expanded)

    raise RuntimeError(
        "ffmpeg was not found. Set an ffmpeg.exe path in the UI/CLI, or install it in PATH."
    )


def resolve_ffprobe_executable(ffmpeg_path: str = "") -> str:
    candidates: list[str] = []
    if ffmpeg_path.strip():
        ffmpeg_file = Path(ffmpeg_path.strip()).expanduser()
        if ffmpeg_file.name.lower() == "ffmpeg.exe":
            candidates.append(str(ffmpeg_file.with_name("ffprobe.exe")))
        candidates.append(str(ffmpeg_file))
    candidates.extend(
        [
            "ffprobe",
            r"C:\ffmpeg\bin\ffprobe.exe",
            r"C:\Program Files\ffmpeg\bin\ffprobe.exe",
            r"C:\ProgramData\chocolatey\bin\ffprobe.exe",
            str(Path.home() / "scoop" / "apps" / "ffmpeg" / "current" / "bin" / "ffprobe.exe"),
        ]
    )

    for candidate in candidates:
        if candidate == "ffprobe":
            located = shutil.which("ffprobe")
            if located:
                return located
            continue
        expanded = Path(candidate).expanduser()
        if expanded.exists():
            return str(expanded)

    raise RuntimeError(
        "ffprobe was not found. Set an ffmpeg.exe path in the UI/CLI, or install ffmpeg/ffprobe in PATH."
    )


def get_video_duration_seconds(media_path: Path, ffmpeg_path: str = "") -> float | None:
    ffprobe_path = resolve_ffprobe_executable(ffmpeg_path)
    command = [
        ffprobe_path,
        "-v",
        "error",
        "-show_entries",
        "format=duration",
        "-of",
        "json",
        str(media_path),
    ]
    try:
        completed = subprocess.run(
            command,
            check=True,
            capture_output=True,
        )
        payload = json.loads(completed.stdout.decode("utf-8", errors="replace"))
        duration = payload.get("format", {}).get("duration")
        if duration is None:
            return None
        return float(duration)
    except Exception:
        return None


def extract_frames_from_video(
    media_path: Path,
    output_dir: Path,
    ffmpeg_path: str = "",
    max_frames: int = 8,
    start_time: float | None = None,
    duration_seconds: float | None = None,
) -> list[Path]:
    resolved_ffmpeg = resolve_ffmpeg_executable(ffmpeg_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    frame_pattern = output_dir / "frame_%03d.jpg"
    effective_duration = duration_seconds
    if effective_duration is None:
        effective_duration = get_video_duration_seconds(media_path, ffmpeg_path)
    if effective_duration is not None and effective_duration > 0:
        fps_value = max_frames / effective_duration
    else:
        fps_value = 1 / 8
    command = [resolved_ffmpeg, "-y"]
    if start_time is not None and start_time > 0:
        command.extend(["-ss", f"{start_time:.3f}"])
    command.extend(["-i", str(media_path)])
    if duration_seconds is not None and duration_seconds > 0:
        command.extend(["-t", f"{duration_seconds:.3f}"])
    command.extend(
        [
            "-vf",
            f"fps={fps_value:.6f},scale=1024:-1",
            "-frames:v",
            str(max_frames),
            str(frame_pattern),
        ]
    )
    try:
        completed = subprocess.run(
            command,
            check=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as exc:
        details_bytes = exc.stderr or exc.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "Failed to extract frames from the video file with ffmpeg.\n"
            f"Details: {details}"
        ) from exc

    frames = sorted(output_dir.glob("frame_*.jpg"))
    if not frames:
        details_bytes = completed.stderr or completed.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "ffmpeg finished but no video frames were created.\n"
            f"Details: {details}"
        )
    return frames


def split_video_segment(
    media_path: Path,
    output_path: Path,
    start_time: float,
    duration_seconds: float,
    ffmpeg_path: str = "",
) -> None:
    resolved_ffmpeg = resolve_ffmpeg_executable(ffmpeg_path)
    command = [
        resolved_ffmpeg,
        "-y",
        "-ss",
        f"{start_time:.3f}",
        "-i",
        str(media_path),
        "-t",
        f"{duration_seconds:.3f}",
        "-c:v",
        "libx264",
        "-preset",
        "veryfast",
        "-crf",
        "23",
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        str(output_path),
    ]
    try:
        completed = subprocess.run(
            command,
            check=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as exc:
        details_bytes = exc.stderr or exc.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "Failed to split the video file with ffmpeg.\n"
            f"Details: {details}"
        ) from exc

    if not output_path.exists() or output_path.stat().st_size == 0:
        details_bytes = completed.stderr or completed.stdout or b""
        details = details_bytes.decode("utf-8", errors="replace").strip()
        raise RuntimeError(
            "ffmpeg finished but no video segment was created.\n"
            f"Details: {details}"
        )


def iter_chunks(waveform, chunk_size: int) -> Iterable[tuple[int, object]]:
    chunk_index = 0
    for start in range(0, len(waveform), chunk_size):
        yield chunk_index, waveform[start : start + chunk_size]
        chunk_index += 1


def build_transcription_prompt(language: str, mode: str) -> str:
    if mode == "song":
        return (
            "Transcribe the following singing in its original language. "
            "Follow these specific instructions for formatting the answer:\n"
            "* Only output the lyrics.\n"
            "* Do not translate.\n"
            "* Do not add any other text.\n"
            "* Do not add timestamps, explanations, or speaker labels."
        )
    return (
        "Transcribe the following speech segment in its original language. "
        "Follow these specific instructions for formatting the answer:\n"
        "* Only output the transcription, with no newlines.\n"
        "* Do not add any other text.\n"
        "* When transcribing numbers, write the digits, i.e. write 1.7 and not one point seven, and write 3 instead of three."
    )


def describe_missing_chat_template_files(model_dir: Path) -> str:
    expected_files = [
        "chat_template.jinja",
        "preprocessor_config.json",
        "special_tokens_map.json",
        "tokenizer.model",
    ]
    missing = [name for name in expected_files if not (model_dir / name).exists()]
    if not missing:
        return "No obvious missing template files were detected."
    return "Missing files: " + ", ".join(missing)


def build_transcription_prompt_variants(language: str, mode: str) -> list[str]:
    language_key = language.lower()
    if mode == "song":
        if language_key == "ja":
            return [
                build_transcription_prompt(language, mode),
                "歌詞だけ。",
                "歌詞のみ。",
            ]
        return [
            build_transcription_prompt(language, mode),
            "Output only the lyrics.",
        ]
    if language_key == "ja":
        return [
            build_transcription_prompt(language, mode),
            "文字起こしだけ。",
            "話している内容だけを書いてください。",
        ]
    return [
        build_transcription_prompt(language, mode),
        "Output only the transcript.",
    ]


def build_audio_chat_messages(prompt_text: str) -> list[dict[str, object]]:
    return [
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": "__audio__"},
                {"type": "text", "text": prompt_text.strip()},
            ],
        }
    ]


def build_audio_transcription_messages(
    audio_source: str,
    prompt_text: str,
) -> list[dict[str, object]]:
    return [
        {
            "role": "user",
            "content": [
                {"type": "audio", "audio": audio_source},
                {"type": "text", "text": prompt_text.strip()},
            ],
        }
    ]


def build_chat_prompt(system_prompt: str, messages: list[dict[str, str]]) -> str:
    lines: list[str] = []
    if system_prompt.strip():
        lines.append(f"System:<|think|> {system_prompt.strip()}")
    for message in messages:
        role = message["role"]
        content = message["content"].strip()
        if not content:
            continue
        if role == "user":
            lines.append(f"User: {content}")
        else:
            lines.append(f"Assistant: {content}")
    lines.append("Assistant:")
    return "\n".join(lines)


def build_chat_messages(
    system_prompt: str,
    messages: list[dict[str, str]],
) -> list[dict[str, object]]:
    chat_messages: list[dict[str, object]] = []
    if system_prompt.strip():
        chat_messages.append(
            {
                "role": "system",
                "content": [{"type": "text", "text": system_prompt.strip()}],
            }
        )
    for message in messages:
        role = message.get("role", "").strip()
        content = message.get("content", "").strip()
        if role not in {"user", "assistant"} or not content:
            continue
        chat_messages.append(
            {
                "role": role,
                "content": [{"type": "text", "text": content}],
            }
        )
    return chat_messages


def build_text_chat_template_messages(
    system_prompt: str,
    messages: list[dict[str, str]],
) -> list[dict[str, str]]:
    chat_messages: list[dict[str, str]] = []
    if system_prompt.strip():
        chat_messages.append({"role": "system", "content": system_prompt.strip()})
    for message in messages:
        role = message.get("role", "").strip()
        content = message.get("content", "").strip()
        if role not in {"user", "assistant"} or not content:
            continue
        chat_messages.append({"role": role, "content": content})
    return chat_messages


def build_image_prompt(language: str, user_prompt: str) -> str:
    prompt = user_prompt.strip()
    if prompt:
        return prompt
    if language.lower() == "ja":
        return "この画像を詳しく説明してください。文字があれば読み取ってください。"
    return "Describe this image in detail and read any visible text."


def build_video_prompt(language: str, user_prompt: str) -> str:
    prompt = user_prompt.strip()
    if prompt:
        return prompt
    return "Describe this video."


def build_video_chunk_prompt(
    language: str,
    user_prompt: str,
    start_seconds: float,
    end_seconds: float,
) -> str:
    base_prompt = build_video_prompt(language, user_prompt)
    return (
        f"This is the video segment from {start_seconds:.0f}s to {end_seconds:.0f}s. "
        f"{base_prompt}"
    )


def build_video_summary_prompt(language: str, segment_texts: list[str]) -> str:
    joined_segments = "\n\n".join(
        f"Segment {index + 1}:\n{text}"
        for index, text in enumerate(segment_texts)
        if text.strip()
    )
    if language.lower() == "ja":
        return (
            "以下は動画の各区間の説明です。全体として自然につながるように、"
            "冒頭から終盤まで統合した最終説明を日本語で作成してください。本文だけを書いてください。\n\n"
            f"{joined_segments}"
        )
    return (
        "The following are descriptions of each video segment. "
        "Write one coherent final description of the full video from beginning to end. Output only the description.\n\n"
        f"{joined_segments}"
    )


def clean_chat_text(text: str) -> str:
    separators = (
        "\nUser:",
        "\nAssistant:",
        "User:",
        "Assistant:",
        "<eos>",
    )
    cleaned = text.strip()
    for separator in separators:
        if separator in cleaned:
            cleaned = cleaned.split(separator, 1)[0].strip()
    return cleaned


def normalize_parsed_response(parsed) -> str:
    if parsed is None:
        return ""
    if isinstance(parsed, str):
        return parsed.strip()
    if isinstance(parsed, list):
        parts: list[str] = []
        for part in parsed:
            normalized = normalize_parsed_response(part)
            if normalized:
                parts.append(normalized)
        return "\n".join(parts).strip()
    if isinstance(parsed, dict):
        role = parsed.get("role")
        content = parsed.get("content")
        if isinstance(content, str):
            stripped = content.strip()
            if stripped:
                return stripped
        if isinstance(content, list):
            normalized_content = normalize_parsed_response(content)
            if normalized_content:
                return normalized_content
        for key in ("response", "text", "content", "output", "answer", "final"):
            value = parsed.get(key)
            normalized_value = normalize_parsed_response(value)
            if normalized_value:
                return normalized_value
        if role == "assistant":
            return ""
        return str(parsed).strip()
    return str(parsed).strip()


def parse_chat_result(parsed, raw_text: str) -> ChatResult:
    raw_stripped = raw_text.strip()
    if isinstance(parsed, dict):
        content = normalize_parsed_response(parsed.get("content"))
        thinking = normalize_parsed_response(parsed.get("thinking"))
        if thinking and raw_stripped:
            content_without_thinking = raw_stripped
            if "<|channel|>thought" in content_without_thinking and "<channel|>" in content_without_thinking:
                content_without_thinking = content_without_thinking.split("<channel|>", 1)[1]
            content_without_thinking = re.sub(r"<turn\|>\s*$", "", content_without_thinking).strip()
            content_without_thinking = strip_special_only_response(content_without_thinking)
            if content_without_thinking:
                content = content or content_without_thinking
        if content or thinking:
            return ChatResult(content=content, thinking=thinking, raw_text=raw_stripped)
    normalized = normalize_parsed_response(parsed)
    return ChatResult(content=normalized, thinking="", raw_text=raw_stripped)


def strip_special_only_response(text: str) -> str:
    stripped = text.strip()
    if not stripped:
        return ""
    normalized = re.sub(r"<[^>\n]+>", "", stripped).strip()
    if not normalized:
        return ""
    return stripped


def is_empty_or_special_response(text: str) -> bool:
    return not strip_special_only_response(text)


def clean_transcript_text(text: str, prompt: str) -> str:
    timestamp_line_pattern = re.compile(
        r"^\d{2}:\d{2}:\d{2},\d{3}(?:\s*-->\s*\d{2}:\d{2}:\d{2},\d{3})?$"
    )
    dot_only_pattern = re.compile(r"^[.\s]{5,}$")
    banned_prefixes = (
        "* ",
        "- ",
        "Please respond only",
        "No comments.",
        "here are the lyrics",
        "Here are the lyrics",
        "If you want to sing along",
        "This video and its lyrics are not for commercial use.",
        "Only output",
        "Do not add",
        "Output only",
        "No explanation",
        "Transcribe this",
        "Transcribe the following",
        "Follow these specific instructions",
        "この歌声を",
        "この音声を",
        "この歌の歌詞を書き起こしてください",
        "元の言語のまま、歌詞本文だけを出力してください",
        "文字起こし結果のみを書いてください",
        "話している内容だけを書いてください",
        "歌詞だけ",
        "歌詞のみ",
        "歌詞を日本語にしてください",
        "字幕形式にしないでください",
        "タイムスタンプ",
        "歌詞の本文だけを出力してください",
        "本文だけを出力してください",
    )
    prompt_lines = {
        line.strip()
        for line in prompt.splitlines()
        if line.strip() and line.strip() != "<|audio|>"
    }
    cleaned_lines: list[str] = []
    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if set(line) == {"-"}:
            continue
        if timestamp_line_pattern.match(line):
            continue
        if dot_only_pattern.match(line):
            continue
        if line in prompt_lines:
            continue
        if any(line.startswith(prefix) for prefix in banned_prefixes):
            continue
        cleaned_lines.append(line)

    return "\n".join(cleaned_lines).strip()


def salvage_transcript_text(text: str, prompt: str) -> str:
    cleaned = text.replace(prompt, "").replace("<|audio|>", "").strip()
    lines: list[str] = []
    for raw_line in cleaned.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if line.startswith("Please respond only"):
            continue
        if line.startswith("No comments."):
            continue
        if line.lower().startswith("here are the lyrics"):
            continue
        if line.startswith("この歌の歌詞を書き起こしてください"):
            continue
        if line.startswith("元の言語のまま、歌詞本文だけを出力してください"):
            continue
        if line.startswith("文字起こし結果のみを書いてください"):
            continue
        if line.startswith("If you want to sing along"):
            continue
        if line.startswith("This video and its lyrics are not for commercial use."):
            continue
        if "Do not add any other text" in line:
            continue
        if set(line) == {"-"}:
            continue
        lines.append(line)
    return "\n".join(lines).strip()


def normalize_transcript_output(text: str, language: str, mode: str) -> str:
    normalized = text.strip()
    if not normalized:
        return normalized

    if language.lower() == "ja":
        if mode == "song":
            lines = [line.replace(" ", "") for line in normalized.splitlines()]
            return "\n".join(line for line in lines if line).strip()
        return normalized.replace(" ", "")

    return re.sub(r"[ \t]+", " ", normalized)


def is_low_quality_transcript(text: str) -> bool:
    stripped = text.strip()
    if not stripped:
        return True

    if re.fullmatch(r"[\d年月日時分秒、。,\-/: ]+", stripped):
        return True

    if len(stripped) <= 20:
        repeated = re.fullmatch(r"(.{1,10})\1{1,}", stripped)
        if repeated:
            return True

    if "\n" not in stripped:
        parts = [part for part in re.split(r"(。|、|\s+)", stripped) if part and not part.isspace()]
        text_only_parts = [part for part in parts if part not in {"。", "、"}]
        if len(text_only_parts) >= 2 and len(set(text_only_parts)) == 1:
            return True

    return False


def dedupe_repeated_phrases(text: str) -> str:
    stripped = text.strip()
    if not stripped:
        return stripped

    for unit_size in range(1, min(20, len(stripped) // 2 + 1)):
        unit = stripped[:unit_size]
        if unit * (len(stripped) // unit_size) == stripped[: unit_size * (len(stripped) // unit_size)]:
            remainder = stripped[unit_size * (len(stripped) // unit_size) :]
            if not remainder:
                return unit

    lines: list[str] = []
    prev = None
    for line in stripped.splitlines():
        if line == prev:
            continue
        lines.append(line)
        prev = line
    return "\n".join(lines).strip()


@dataclass(frozen=True)
class ChunkTranscriptionResult:
    chunk_index: int
    prompt_text: str
    raw_text: str
    cleaned_text: str


@dataclass(frozen=True)
class VideoAnalysisChunkResult:
    chunk_index: int
    start_seconds: float
    end_seconds: float
    prompt_text: str
    raw_text: str
    cleaned_text: str


@dataclass(frozen=True)
class ChatResult:
    content: str
    thinking: str
    raw_text: str


class GemmaTranscriber:
    def __init__(self, settings: TranscriberSettings):
        self.settings = settings
        (
            librosa,
            np,
            torch,
            AutoModelForCausalLM,
            AutoModelForMultimodalLM,
            AutoProcessor,
        ) = lazy_imports()
        self.librosa = librosa
        self.np = np
        self.torch = torch
        self.AutoModelForCausalLM = AutoModelForCausalLM
        self.AutoModelForMultimodalLM = AutoModelForMultimodalLM

        if not settings.model_dir.is_dir():
            raise RuntimeError(f"Model directory not found: {settings.model_dir}")
        if settings.mode not in {"speech", "song"}:
            raise RuntimeError(f"Unsupported mode: {settings.mode}")
        if settings.chunk_seconds > 30:
            raise RuntimeError(
                "Gemma4 audio input supports up to 30 seconds per chunk. "
                f"Current chunk_seconds: {settings.chunk_seconds}"
            )

        dtype = resolve_torch_dtype(torch, settings.dtype)
        max_memory = parse_max_memory(settings.max_memory)
        self.model_dtype = dtype
        self.model_max_memory = max_memory
        self.processor = AutoProcessor.from_pretrained(settings.model_dir)
        self.model = None
        self.model_kind = ""
        self.model_device = None
        self._load_model("multimodal")

    def _build_model_kwargs(self):
        model_kwargs = {
            "pretrained_model_name_or_path": self.settings.model_dir,
            "dtype": self.model_dtype,
            "device_map": self.settings.device_map,
        }
        if self.model_max_memory is not None:
            model_kwargs["max_memory"] = self.model_max_memory
        return model_kwargs

    def _load_model(self, model_kind: str) -> None:
        if self.model is not None and self.model_kind == model_kind:
            return
        if self.model is not None:
            try:
                self.model.cpu()
            except Exception:
                pass
            del self.model
            self.model = None
            gc.collect()
            if self.torch.cuda.is_available():
                self.torch.cuda.empty_cache()
                try:
                    self.torch.cuda.ipc_collect()
                except Exception:
                    pass

        model_kwargs = self._build_model_kwargs()
        if model_kind == "causal":
            self.model = self.AutoModelForCausalLM.from_pretrained(**model_kwargs)
        else:
            self.model = self.AutoModelForMultimodalLM.from_pretrained(**model_kwargs)
        self.model_kind = model_kind
        self.model_device = self._resolve_model_device()

    def _ensure_multimodal_model(self) -> None:
        self._load_model("multimodal")

    def _ensure_causal_model(self) -> None:
        self._load_model("causal")

    def _resolve_model_device(self):
        hf_device_map = getattr(self.model, "hf_device_map", None)
        if isinstance(hf_device_map, dict):
            cuda_indexes: list[int] = []
            for device in hf_device_map.values():
                if isinstance(device, int):
                    cuda_indexes.append(device)
                    continue
                if isinstance(device, str) and device.startswith("cuda:"):
                    suffix = device.split(":", 1)[1]
                    if suffix.isdigit():
                        cuda_indexes.append(int(suffix))
            if cuda_indexes:
                return self.torch.device(f"cuda:{min(cuda_indexes)}")

        model_device = getattr(self.model, "device", None)
        if model_device is not None:
            return model_device
        return next(self.model.parameters()).device

    def _move_inputs(self, inputs):
        moved = {}
        for key, value in inputs.items():
            if hasattr(value, "to"):
                target_dtype = None
                if hasattr(value, "dtype") and getattr(value.dtype, "is_floating_point", False):
                    target_dtype = self.model.dtype
                if target_dtype is not None:
                    moved[key] = value.to(self.model_device, dtype=target_dtype)
                else:
                    moved[key] = value.to(self.model_device)
            else:
                moved[key] = value
        return moved

    def _cleanup_generation_tensors(self, *objects) -> None:
        for obj in objects:
            try:
                del obj
            except Exception:
                pass
        gc.collect()
        if self.torch.cuda.is_available():
            self.torch.cuda.empty_cache()

    def _render_audio_prompt(self, prompt_text: str) -> str:
        tokenizer = getattr(self.processor, "tokenizer", None)
        if tokenizer is not None and hasattr(tokenizer, "apply_chat_template"):
            try:
                rendered = tokenizer.apply_chat_template(
                    build_audio_chat_messages(prompt_text),
                    tokenize=False,
                    add_generation_prompt=True,
                )
                if isinstance(rendered, str) and rendered.strip():
                    return rendered
            except Exception:
                pass
        return f"<|audio|>\n{prompt_text.strip()}"

    def _get_generation_stop_ids(self) -> int | list[int] | None:
        tokenizer = getattr(self.processor, "tokenizer", None)
        if tokenizer is None:
            return None

        stop_ids: list[int] = []
        for token_name in ("eos_token_id",):
            token_id = getattr(tokenizer, token_name, None)
            if isinstance(token_id, int) and token_id >= 0:
                stop_ids.append(token_id)

        if hasattr(tokenizer, "convert_tokens_to_ids"):
            for token in ("<turn|>", "<eos>"):
                try:
                    token_id = tokenizer.convert_tokens_to_ids(token)
                except Exception:
                    token_id = None
                if isinstance(token_id, int) and token_id >= 0:
                    stop_ids.append(token_id)

        deduped: list[int] = []
        for token_id in stop_ids:
            if token_id not in deduped:
                deduped.append(token_id)
        if not deduped:
            return None
        if len(deduped) == 1:
            return deduped[0]
        return deduped

    def _get_pad_token_id(self) -> int | None:
        tokenizer = getattr(self.processor, "tokenizer", None)
        if tokenizer is None:
            return None
        pad_token_id = getattr(tokenizer, "pad_token_id", None)
        if isinstance(pad_token_id, int) and pad_token_id >= 0:
            return pad_token_id
        eos_token_id = getattr(tokenizer, "eos_token_id", None)
        if isinstance(eos_token_id, int) and eos_token_id >= 0:
            return eos_token_id
        return None

    def _build_audio_inputs(self, prompt_text: str, chunk, sample_rate: int):
        rendered_prompt = self._render_audio_prompt(prompt_text)
        attempts = (
            {"audio": chunk, "sampling_rate": sample_rate},
            {"audio": [chunk], "sampling_rate": sample_rate},
        )
        last_error = None
        for extra_kwargs in attempts:
            try:
                return (
                    self.processor(
                        text=rendered_prompt,
                        return_tensors="pt",
                        padding=True,
                        **extra_kwargs,
                    ),
                    rendered_prompt,
                )
            except Exception as exc:
                last_error = exc
        if last_error is not None:
            raise last_error
        raise RuntimeError("Failed to build audio inputs.")

    def _run_audio_transcription_attempt(
        self,
        *,
        prompt_text: str,
        chunk,
        sample_rate: int,
        max_new_tokens: int,
        do_sample: bool,
    ) -> tuple[str, str]:
        input_ids, rendered_prompt = self._build_audio_inputs(
            prompt_text, chunk, sample_rate
        )
        input_ids = self._move_inputs(input_ids)
        outputs = None
        response = ""
        fallback = ""
        try:
            generation_kwargs = {
                "max_new_tokens": max_new_tokens,
                "do_sample": do_sample,
                "use_cache": False,
                "repetition_penalty": 1.05,
            }
            eos_token_id = self._get_generation_stop_ids()
            if eos_token_id is not None:
                generation_kwargs["eos_token_id"] = eos_token_id
            pad_token_id = self._get_pad_token_id()
            if pad_token_id is not None:
                generation_kwargs["pad_token_id"] = pad_token_id
            if do_sample:
                generation_kwargs["temperature"] = DEFAULT_TEMPERATURE
                generation_kwargs["top_p"] = DEFAULT_TOP_P
                generation_kwargs["top_k"] = DEFAULT_TOP_K

            with self.torch.inference_mode():
                outputs = self.model.generate(
                    **input_ids,
                    **generation_kwargs,
                )

            input_length = input_ids["input_ids"].shape[-1]
            response = self.processor.decode(
                outputs[0][input_length:],
                skip_special_tokens=False,
            )
            decoded = ""
            if hasattr(self.processor, "parse_response"):
                try:
                    decoded = normalize_parsed_response(
                        self.processor.parse_response(response)
                    )
                except Exception:
                    decoded = ""
            if not decoded.strip():
                decoded = response.strip()
            if not decoded.strip():
                fallback = self.processor.decode(
                    outputs[0][input_length:],
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=True,
                ).strip()
                decoded = fallback
            return strip_special_only_response(decoded), rendered_prompt
        finally:
            self._cleanup_generation_tensors(
                input_ids,
                outputs,
                response,
                fallback,
            )

    def _generate_from_inputs(
        self,
        inputs,
        max_new_tokens: int,
        *,
        do_sample: bool,
        use_parse_response: bool,
    ) -> str:
        with self.torch.inference_mode():
            output_ids = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                temperature=DEFAULT_TEMPERATURE if do_sample else None,
                top_p=DEFAULT_TOP_P if do_sample else None,
                top_k=DEFAULT_TOP_K if do_sample else None,
            )

        prompt_length = inputs["input_ids"].shape[1]
        generated_ids = output_ids[:, prompt_length:]
        raw_text = self.processor.decode(
            generated_ids[0],
            skip_special_tokens=False,
        )

        if use_parse_response and hasattr(self.processor, "parse_response"):
            try:
                parsed = self.processor.parse_response(raw_text)
                parsed_text = normalize_parsed_response(parsed)
                if parsed_text:
                    return parsed_text
            except Exception:
                pass

        return self.processor.decode(
            generated_ids[0],
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True,
        ).strip()

    def _run_chat_template_request(
        self,
        messages: list[dict[str, object]],
        *,
        max_new_tokens: int,
        do_sample: bool = True,
    ) -> str:
        inputs = None
        outputs = None
        response = ""
        fallback = ""
        try:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
            )
            inputs = self._move_inputs(inputs)
            input_len = inputs["input_ids"].shape[-1]

            generation_kwargs = {
                "max_new_tokens": max_new_tokens,
                "do_sample": do_sample,
                "use_cache": False,
            }
            eos_token_id = self._get_generation_stop_ids()
            if eos_token_id is not None:
                generation_kwargs["eos_token_id"] = eos_token_id
            pad_token_id = self._get_pad_token_id()
            if pad_token_id is not None:
                generation_kwargs["pad_token_id"] = pad_token_id
            if do_sample:
                generation_kwargs["temperature"] = DEFAULT_TEMPERATURE
                generation_kwargs["top_p"] = DEFAULT_TOP_P
                generation_kwargs["top_k"] = DEFAULT_TOP_K

            with self.torch.inference_mode():
                outputs = self.model.generate(
                    **inputs,
                    **generation_kwargs,
                )

            response = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=False,
            )
            if hasattr(self.processor, "parse_response"):
                try:
                    parsed = self.processor.parse_response(response)
                    parsed_text = normalize_parsed_response(parsed)
                    if strip_special_only_response(parsed_text):
                        return parsed_text
                except Exception:
                    pass
            stripped_response = response.strip()
            if strip_special_only_response(stripped_response):
                return stripped_response
            fallback = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            ).strip()
            return strip_special_only_response(fallback)
        finally:
            self._cleanup_generation_tensors(inputs, outputs, response, fallback)

    def _run_chat_template_request_debug(
        self,
        messages: list[dict[str, object]],
        *,
        max_new_tokens: int,
        do_sample: bool = False,
    ) -> tuple[str, str]:
        inputs = None
        outputs = None
        response = ""
        fallback = ""
        try:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
            )
            inputs = self._move_inputs(inputs)
            input_len = inputs["input_ids"].shape[-1]

            generation_kwargs = {
                "max_new_tokens": max_new_tokens,
                "do_sample": do_sample,
                "use_cache": False,
            }
            if do_sample:
                generation_kwargs["temperature"] = DEFAULT_TEMPERATURE
                generation_kwargs["top_p"] = DEFAULT_TOP_P
                generation_kwargs["top_k"] = DEFAULT_TOP_K

            with self.torch.inference_mode():
                outputs = self.model.generate(
                    **inputs,
                    **generation_kwargs,
                )

            response = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=False,
            ).strip()
            parsed_text = ""
            if hasattr(self.processor, "parse_response"):
                try:
                    parsed = self.processor.parse_response(response)
                    parsed_text = normalize_parsed_response(parsed)
                except Exception:
                    parsed_text = ""
            if not parsed_text:
                fallback = self.processor.decode(
                    outputs[0][input_len:],
                    skip_special_tokens=True,
                    clean_up_tokenization_spaces=True,
                ).strip()
                parsed_text = fallback
            return response, parsed_text.strip()
        finally:
            self._cleanup_generation_tensors(inputs, outputs, response, fallback)

    def _run_audio_chat_template_request(
        self,
        messages: list[dict[str, object]],
        *,
        max_new_tokens: int,
    ) -> str:
        inputs = None
        outputs = None
        response = ""
        fallback = ""
        try:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
            )
            inputs = self._move_inputs(inputs)
            input_len = inputs["input_ids"].shape[-1]

            with self.torch.inference_mode():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                )

            response = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=False,
            )
            if hasattr(self.processor, "parse_response"):
                try:
                    parsed = self.processor.parse_response(response)
                    parsed_text = normalize_parsed_response(parsed)
                    if parsed_text:
                        return parsed_text
                except Exception:
                    pass
            fallback = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            ).strip()
            return fallback or response.strip()
        finally:
            self._cleanup_generation_tensors(inputs, outputs, response, fallback)

    def _run_text_chat_template_request(
        self,
        messages: list[dict[str, str]],
        *,
        max_new_tokens: int,
        enable_thinking: bool,
    ) -> ChatResult:
        if not hasattr(self.processor, "apply_chat_template"):
            raise RuntimeError("The current processor does not support chat templates.")

        inputs = None
        outputs = None
        response = ""
        try:
            try:
                text = self.processor.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True,
                    enable_thinking=enable_thinking,
                )
            except TypeError:
                text = self.processor.apply_chat_template(
                    messages,
                    tokenize=False,
                    add_generation_prompt=True,
                )
            inputs = self.processor(
                text=text,
                return_tensors="pt",
            )
            inputs = self._move_inputs(inputs)
            input_len = inputs["input_ids"].shape[-1]

            with self.torch.inference_mode():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                )

            response = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=False,
            )
            if hasattr(self.processor, "parse_response"):
                try:
                    parsed = self.processor.parse_response(response)
                    return parse_chat_result(parsed, response)
                except Exception:
                    pass

            fallback = self.processor.decode(
                outputs[0][input_len:],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            ).strip()
            return ChatResult(
                content=fallback or response.strip(),
                thinking="",
                raw_text=response.strip(),
            )
        finally:
            self._cleanup_generation_tensors(inputs, outputs, response)

    def _write_chunk_wav(self, path: Path, waveform, sample_rate: int) -> None:
        pcm = self.np.clip(waveform, -1.0, 1.0)
        pcm = (pcm * 32767).astype(self.np.int16)
        with wave.open(str(path), "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(sample_rate)
            wav_file.writeframes(pcm.tobytes())

    def chat(
        self,
        messages: list[dict[str, str]],
        max_new_tokens: int = 256,
        system_prompt: str = "You are a helpful assistant.",
        enable_thinking: bool = False,
    ) -> ChatResult:
        self._ensure_causal_model()
        if not messages:
            raise RuntimeError("Chat messages are empty.")

        if getattr(self.processor, "chat_template", None) is not None:
            result = self._run_text_chat_template_request(
                build_text_chat_template_messages(system_prompt, messages),
                max_new_tokens=max_new_tokens,
                enable_thinking=enable_thinking,
            )
        else:
            prompt_text = build_chat_prompt(system_prompt, messages)
            inputs = self.processor(
                text=prompt_text,
                add_special_tokens=True,
                return_tensors="pt",
            )
            inputs = self._move_inputs(inputs)
            text = self._generate_from_inputs(
                inputs,
                max_new_tokens,
                do_sample=False,
                use_parse_response=True,
            ).strip()
            result = ChatResult(content=text, thinking="", raw_text=text)
        return ChatResult(
            content=clean_chat_text(result.content),
            thinking=result.thinking.strip(),
            raw_text=result.raw_text,
        )

    def analyze_image(
        self,
        image_path: Path,
        *,
        prompt: str,
        language: str,
        max_new_tokens: int = 256,
    ) -> str:
        self._ensure_multimodal_model()
        if not image_path.is_file():
            raise RuntimeError(f"Image file not found: {image_path}")

        if getattr(self.processor, "chat_template", None) is None:
            raise RuntimeError(
                "This image flow requires a processor with a chat template. "
                "Use a Gemma 4 instruction-tuned checkpoint such as "
                "`google/gemma-4-E2B-it` or a local instruction-tuned export."
            )

        prompt_text = build_image_prompt(language, prompt)
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_path.as_posix()},
                    {"type": "text", "text": prompt_text},
                ],
            }
        ]

        input_ids = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
        )
        input_ids = self._move_inputs(input_ids)
        raw_text = self._generate_from_inputs(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            use_parse_response=True,
        )
        return raw_text.strip()

    def analyze_video(
        self,
        video_path: Path,
        *,
        prompt: str,
        language: str,
        max_new_tokens: int = 256,
    ) -> str:
        self._ensure_multimodal_model()
        if not video_path.is_file():
            raise RuntimeError(f"Video file not found: {video_path}")
        if not is_video_file(video_path):
            raise RuntimeError(f"Unsupported video file: {video_path}")
        chunk_results = self.analyze_video_chunks(
            video_path,
            prompt=prompt,
            language=language,
            max_new_tokens=max_new_tokens,
        )
        chunk_texts = [
            result.cleaned_text or result.raw_text
            for result in chunk_results
            if (result.cleaned_text or result.raw_text).strip()
        ]
        if not chunk_texts:
            return ""

        if getattr(self.processor, "chat_template", None) is None:
            raise RuntimeError(
                "This video flow requires a processor with a chat template. "
                "Use a Gemma 4 instruction-tuned checkpoint such as "
                "`google/gemma-4-E2B-it` or a local instruction-tuned export."
            )

        prompt_text = build_video_summary_prompt(language, chunk_texts)
        messages = [
            {
                "role": "user",
                "content": [{"type": "text", "text": prompt_text}],
            }
        ]
        return self._run_chat_template_request(
            messages,
            max_new_tokens=min(max_new_tokens, 192),
            do_sample=False,
        ).strip()

    def analyze_video_chunks(
        self,
        video_path: Path,
        *,
        prompt: str,
        language: str,
        max_new_tokens: int = 256,
        chunk_seconds: int = DEFAULT_VIDEO_CHUNK_SECONDS,
    ) -> list[VideoAnalysisChunkResult]:
        self._ensure_multimodal_model()
        if not video_path.is_file():
            raise RuntimeError(f"Video file not found: {video_path}")
        if not is_video_file(video_path):
            raise RuntimeError(f"Unsupported video file: {video_path}")

        if getattr(self.processor, "chat_template", None) is None:
            raise RuntimeError(
                "This video flow requires a processor with a chat template. "
                "Use a Gemma 4 instruction-tuned checkpoint such as "
                "`google/gemma-4-E2B-it` or a local instruction-tuned export."
            )

        total_duration = get_video_duration_seconds(
            video_path,
            self.settings.ffmpeg_path,
        )
        if total_duration is None or total_duration <= 0:
            total_duration = float(chunk_seconds)

        results: list[VideoAnalysisChunkResult] = []
        with TemporaryDirectory(prefix="gemma4_video_") as temp_dir:
            chunk_index = 0
            start_seconds = 0.0
            while start_seconds < total_duration:
                end_seconds = min(start_seconds + chunk_seconds, total_duration)
                segment_path = Path(temp_dir) / f"segment_{chunk_index + 1:03d}.mp4"
                split_video_segment(
                    video_path,
                    segment_path,
                    start_time=start_seconds,
                    duration_seconds=max(end_seconds - start_seconds, 1.0),
                    ffmpeg_path=self.settings.ffmpeg_path,
                )
                prompt_text = build_video_chunk_prompt(
                    language,
                    prompt,
                    start_seconds,
                    end_seconds,
                )
                messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "video", "video": segment_path.as_posix()},
                            {"type": "text", "text": prompt_text},
                        ],
                    }
                ]
                raw_text, parsed_text = self._run_chat_template_request_debug(
                    messages,
                    max_new_tokens=min(max_new_tokens, 256),
                    do_sample=False,
                )
                results.append(
                    VideoAnalysisChunkResult(
                        chunk_index=chunk_index,
                        start_seconds=start_seconds,
                        end_seconds=end_seconds,
                        prompt_text=prompt_text,
                        raw_text=raw_text,
                        cleaned_text=parsed_text,
                    )
                )
                chunk_index += 1
                start_seconds = end_seconds

        return results

    def transcribe_chunks(
        self,
        audio_path: Path,
        progress_callback: Callable[[str], None] | None = None,
        chunk_result_callback: Callable[[ChunkTranscriptionResult], None] | None = None,
    ) -> list[ChunkTranscriptionResult]:
        self._ensure_multimodal_model()
        if not audio_path.is_file():
            raise RuntimeError(f"Media file not found: {audio_path}")
        if getattr(self.processor, "chat_template", None) is None:
            raise RuntimeError(
                "This transcription flow requires a processor with a chat template. "
                "Use a Gemma 4 instruction-tuned checkpoint such as "
                "`google/gemma-4-E2B-it` or a local instruction-tuned export."
            )

        with TemporaryDirectory(prefix="gemma4_audio_") as temp_dir:
            source_audio_path = audio_path
            if is_video_file(audio_path):
                if progress_callback is not None:
                    progress_callback("Extracting audio from video...")
                extracted_audio_path = Path(temp_dir) / "extracted_audio.wav"
                extract_audio_from_video(
                    audio_path,
                    extracted_audio_path,
                    self.settings.sample_rate,
                    self.settings.ffmpeg_path,
                )
                source_audio_path = extracted_audio_path

            waveform, sample_rate = load_audio(
                self.librosa, source_audio_path, self.settings.sample_rate
            )
            chunk_size = sample_rate * self.settings.chunk_seconds
            prompt = build_transcription_prompt(
                self.settings.language,
                self.settings.mode,
            )
            chunks = [
                chunk for _, chunk in iter_chunks(waveform, chunk_size) if len(chunk) > 0
            ]
            if not chunks:
                return []

            self._cleanup_generation_tensors()
            chunk_temp_dir = Path(temp_dir) / "chunks"
            chunk_temp_dir.mkdir(parents=True, exist_ok=True)

            if progress_callback is not None:
                progress_callback(f"Preparing {len(chunks)} chunk(s)...")

            results: list[ChunkTranscriptionResult] = []
            for chunk_index, chunk in enumerate(chunks, start=1):
                if progress_callback is not None:
                    progress_callback(
                        f"Generating chunk {chunk_index}/{len(chunks)}..."
                    )

                chunk_wav_path = chunk_temp_dir / f"chunk_{chunk_index:04d}.wav"
                self._write_chunk_wav(chunk_wav_path, chunk, sample_rate)

                prompt_text = build_transcription_prompt(
                    self.settings.language,
                    self.settings.mode,
                )
                messages = build_audio_transcription_messages(
                    chunk_wav_path.as_posix(),
                    prompt_text,
                )
                decoded = self._run_audio_chat_template_request(
                    messages,
                    max_new_tokens=min(self.settings.max_new_tokens, 256),
                )
                cleaned_text = clean_transcript_text(decoded, prompt)
                final_text = cleaned_text or salvage_transcript_text(decoded, prompt)
                final_text = normalize_transcript_output(
                    final_text,
                    self.settings.language,
                    self.settings.mode,
                )
                final_text = dedupe_repeated_phrases(final_text)
                if is_low_quality_transcript(final_text):
                    final_text = ""

                results.append(
                    ChunkTranscriptionResult(
                        chunk_index=chunk_index - 1,
                        prompt_text=prompt_text,
                        raw_text=decoded,
                        cleaned_text=final_text,
                    )
                )
                if chunk_result_callback is not None:
                    chunk_result_callback(results[-1])
                self._cleanup_generation_tensors()

        return results

    def transcribe(
        self,
        audio_path: Path,
        progress_callback: Callable[[str], None] | None = None,
    ) -> str:
        chunk_results = self.transcribe_chunks(
            audio_path,
            progress_callback=progress_callback,
        )
        transcript_parts = [
            result.cleaned_text for result in chunk_results if result.cleaned_text
        ]
        return "\n".join(transcript_parts).strip()

    def unload(self) -> None:
        model = getattr(self, "model", None)
        torch_module = getattr(self, "torch", None)

        if model is not None:
            try:
                model.cpu()
            except Exception:
                pass
            del self.model
        if hasattr(self, "processor"):
            del self.processor

        gc.collect()

        if torch_module is not None and torch_module.cuda.is_available():
            torch_module.cuda.empty_cache()
            try:
                torch_module.cuda.ipc_collect()
            except Exception:
                pass
