"""FikoYT Whisper STT — faster-whisper veya openai-whisper, otomatik fallback."""
import os
import tempfile

import torch


WHISPER_LANGS = ["auto", "tr", "en", "de", "fr", "es", "it", "pt", "ru", "ja", "ko", "zh", "ar"]
WHISPER_SIZES = ["tiny", "base", "small", "medium", "large-v2", "large-v3"]


def _audio_to_wav(audio: dict, dst: str):
    import torchaudio
    wav = audio["waveform"]
    sr = int(audio["sample_rate"])
    if wav.dim() == 3:
        wav = wav.squeeze(0)
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    torchaudio.save(dst, wav.cpu().float(), sr)


def _fmt_srt_time(t: float) -> str:
    if t < 0: t = 0.0
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = int(t % 60)
    ms = int((t - int(t)) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def _segments_to_srt(segments) -> str:
    lines = []
    for i, seg in enumerate(segments, 1):
        start = float(seg.get("start", 0.0) if isinstance(seg, dict) else seg.start)
        end = float(seg.get("end", 0.0) if isinstance(seg, dict) else seg.end)
        text = (seg.get("text", "") if isinstance(seg, dict) else seg.text).strip()
        lines.append(f"{i}\n{_fmt_srt_time(start)} --> {_fmt_srt_time(end)}\n{text}\n")
    return "\n".join(lines)


class FikoYTWhisperSTT:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO",),
                "language": (WHISPER_LANGS, {"default": "auto"}),
                "model_size": (WHISPER_SIZES, {"default": "large-v3"}),
            }
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("transcript", "srt")
    FUNCTION = "transcribe"
    CATEGORY = "FikoYT/audio"

    def transcribe(self, audio, language, model_size):
        tmpdir = tempfile.mkdtemp(prefix="fikoyt_whisper_")
        wav_path = os.path.join(tmpdir, "in.wav")
        _audio_to_wav(audio, wav_path)

        lang = None if language == "auto" else language
        device = "cuda" if torch.cuda.is_available() else "cpu"

        # Önce faster-whisper dene
        try:
            from faster_whisper import WhisperModel
            compute = "float16" if device == "cuda" else "int8"
            print(f"[FikoYT Whisper] faster-whisper {model_size} ({device}/{compute})")
            model = WhisperModel(model_size, device=device, compute_type=compute)
            segs_iter, _info = model.transcribe(wav_path, language=lang, vad_filter=True)
            segs = []
            for s in segs_iter:
                segs.append({"start": s.start, "end": s.end, "text": s.text})
            del model
        except ImportError:
            # openai-whisper fallback
            try:
                import whisper as _wh
            except ImportError:
                raise RuntimeError(
                    "Whisper kurulu değil.\n"
                    "Kurulum: pip install faster-whisper  (önerilen) veya pip install openai-whisper"
                )
            print(f"[FikoYT Whisper] openai-whisper {model_size}")
            model = _wh.load_model(model_size, device=device)
            result = model.transcribe(wav_path, language=lang)
            segs = result.get("segments", []) or []
            del model

        transcript = " ".join((s["text"] if isinstance(s, dict) else s.text).strip() for s in segs).strip()
        srt = _segments_to_srt(segs)

        try:
            torch.cuda.empty_cache()
        except Exception:
            pass
        try:
            os.remove(wav_path)
            os.rmdir(tmpdir)
        except Exception:
            pass

        print(f"[FikoYT Whisper] {len(segs)} segment, {len(transcript)} char")
        return (transcript, srt)


NODE_CLASS_MAPPINGS = {"FikoYTWhisperSTT": FikoYTWhisperSTT}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTWhisperSTT": "FikoYT • Whisper STT (+SRT)"}