from __future__ import annotations

import gc
import json
import os
import sys
import uuid
from pathlib import Path
from typing import Any

import torch
import whisper
from safetensors.torch import load_file as load_safetensors_file

from shared.deepy import video_tools as deepy_video_tools
from shared.deepy.assets import (
    WHISPER_MEDIUM_CONFIG_FILENAME,
    WHISPER_MEDIUM_FOLDER,
    WHISPER_MEDIUM_REPO,
    WHISPER_MEDIUM_REQUIRED_FILES,
    WHISPER_MEDIUM_WEIGHTS_FILENAME,
    query_deepy_download_defs,
)
from shared.ffmpeg_setup import download_ffmpeg
from shared.utils import files_locator as fl


_TEMP_ROOT = Path(__file__).resolve().parents[2] / "_temp_codex" / "deepy_transcribe"
_TIMESTAMP_TYPE_ALIASES = {
    "none": None,
    "off": None,
    "disabled": None,
    "segment": "segment",
    "segments": "segment",
    "word": "word",
    "words": "word",
}
_WHISPER_MEDIUM_REQUIRED_FILES = WHISPER_MEDIUM_REQUIRED_FILES


def normalize_timestamp_type(value: Any) -> str | None:
    normalized = str(value or "").strip().lower()
    if len(normalized) == 0:
        return "segment"
    if normalized not in _TIMESTAMP_TYPE_ALIASES:
        raise ValueError("timestamp_type must be 'segment', 'word', or 'none'.")
    return _TIMESTAMP_TYPE_ALIASES[normalized]


def _get_main_callable(name: str) -> Any:
    main_module = sys.modules.get("__main__")
    return None if main_module is None else getattr(main_module, str(name or "").strip(), None)


def _whisper_medium_files_present(model_dir: Path | None) -> bool:
    if model_dir is None or not model_dir.is_dir():
        return False
    return all((model_dir / filename).is_file() for filename in _WHISPER_MEDIUM_REQUIRED_FILES)


def _ensure_whisper_medium_assets(model_dir: Path | None = None) -> None:
    if _whisper_medium_files_present(model_dir):
        return
    process_files_def = _get_main_callable("process_files_def")
    if callable(process_files_def):
        for download_def in query_deepy_download_defs():
            process_files_def(**download_def)


def _whisper_medium_dir() -> Path:
    located = fl.locate_folder(WHISPER_MEDIUM_FOLDER, error_if_none=False)
    located_path = None if located is None else Path(located).resolve()
    _ensure_whisper_medium_assets(located_path)
    located = fl.locate_folder(WHISPER_MEDIUM_FOLDER, error_if_none=False)
    if located is not None:
        resolved = Path(located).resolve()
        if _whisper_medium_files_present(resolved):
            return resolved
    fallback = Path("e:/ml/wan2gp/ckpts") / WHISPER_MEDIUM_FOLDER
    if fallback.is_dir() and _whisper_medium_files_present(fallback):
        return fallback.resolve()
    raise FileNotFoundError(
        f"Unable to locate the Whisper medium folder '{WHISPER_MEDIUM_FOLDER}' in the configured checkpoints paths."
    )


def _load_whisper_medium(device: torch.device) -> whisper.Whisper:
    model_dir = _whisper_medium_dir()
    config_path = model_dir / WHISPER_MEDIUM_CONFIG_FILENAME
    weights_path = model_dir / WHISPER_MEDIUM_WEIGHTS_FILENAME
    if not config_path.is_file():
        raise FileNotFoundError(f"Whisper config file not found: {config_path}")
    if not weights_path.is_file():
        raise FileNotFoundError(f"Whisper weights file not found: {weights_path}")
    with config_path.open("r", encoding="utf-8") as reader:
        config = json.load(reader)
    dims = whisper.model.ModelDimensions(**dict(config.get("dims", {}) or {}))
    model = whisper.model.Whisper(dims)
    model.load_state_dict(load_safetensors_file(str(weights_path), device="cpu"))
    alignment_heads = str(config.get("alignment_heads", "") or "").strip()
    if len(alignment_heads) > 0:
        model.set_alignment_heads(alignment_heads.encode("ascii"))
    model.eval()
    if device.type == "cuda":
        return model.to(device=device)
    return model.to(device=device, dtype=torch.float32)


def _make_temp_audio_path() -> Path:
    _TEMP_ROOT.mkdir(parents=True, exist_ok=True)
    return (_TEMP_ROOT / f"{uuid.uuid4().hex}.wav").resolve()


def _prepare_audio_input(source_path: str, audio_track_no: int | None = None) -> tuple[str, list[Path]]:
    download_ffmpeg()
    temp_audio_path = _make_temp_audio_path()
    deepy_video_tools.extract_audio(source_path, str(temp_audio_path), audio_track_no=audio_track_no, audio_codec="wav")
    return str(temp_audio_path), [temp_audio_path]


def _round_timestamp(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return round(float(value), 3)
    except Exception:
        return None


def _serialize_segments(segments: list[dict[str, Any]], timestamp_type: str | None) -> list[dict[str, Any]]:
    serialized = []
    include_words = timestamp_type == "word"
    for segment in segments:
        item = {
            "start": _round_timestamp(segment.get("start", None)),
            "end": _round_timestamp(segment.get("end", None)),
            "text": str(segment.get("text", "") or "").strip(),
        }
        if include_words:
            words = []
            for word in list(segment.get("words", []) or []):
                words.append(
                    {
                        "word": str(word.get("word", "") or ""),
                        "start": _round_timestamp(word.get("start", None)),
                        "end": _round_timestamp(word.get("end", None)),
                        "probability": None if word.get("probability", None) is None else round(float(word["probability"]), 4),
                    }
                )
            if len(words) > 0:
                item["words"] = words
        serialized.append(item)
    return serialized


def transcribe_media(source_path: str, *, timestamp_type: str | None = None, audio_track_no: int | None = None) -> dict[str, Any]:
    normalized_timestamp_type = normalize_timestamp_type(timestamp_type)
    source_path = str(source_path or "").strip()
    if len(source_path) == 0 or not os.path.isfile(source_path):
        raise FileNotFoundError(f"Media file not found: {source_path}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    audio_path, temporary_paths = _prepare_audio_input(source_path, audio_track_no=audio_track_no)
    model = None
    try:
        model = _load_whisper_medium(device)
        raw_result = model.transcribe(
            audio_path,
            verbose=None,
            fp16=device.type == "cuda",
            word_timestamps=normalized_timestamp_type == "word",
        )
    finally:
        for temporary_path in temporary_paths:
            try:
                temporary_path.unlink(missing_ok=True)
            except Exception:
                pass
        if model is not None:
            del model
        gc.collect()
        if device.type == "cuda":
            torch.cuda.empty_cache()

    segments = list(raw_result.get("segments", []) or [])
    payload = {
        "text": str(raw_result.get("text", "") or "").strip(),
        "language": str(raw_result.get("language", "") or "").strip(),
        "segment_count": int(len(segments)),
    }
    if normalized_timestamp_type is not None:
        payload["timestamp_type"] = normalized_timestamp_type
        payload["segments"] = _serialize_segments(segments, normalized_timestamp_type)
    return payload