import shutil
import subprocess
import tempfile
import unittest
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from fireredtts3.reference_asr import (
    QWEN3_ASR_MODEL_ID,
    transcribe_audio,
    transcribe_reference_audio,
)

REFERENCE_TRANSCRIPT = "The reference transcript is filled automatically."


class _FakeBatch(dict):
    def to(self, device, dtype):
        self.to_args = (device, dtype)
        return self


class _FakeProcessor:
    def __init__(self, decoded_text=REFERENCE_TRANSCRIPT):
        self.decoded_text = decoded_text
        self.audio = None
        self.sampling_rate = None
        self.decode_args = None

    def apply_transcription_request(self, audio, processor_kwargs):
        self.audio = audio
        self.sampling_rate = processor_kwargs["sampling_rate"]
        return _FakeBatch({"input_ids": torch.tensor([[10, 11, 12]])})

    def decode(self, generated_ids, return_format):
        self.decode_args = (generated_ids, return_format)
        return [self.decoded_text]


class _FakeModel:
    device = torch.device("cpu")
    dtype = torch.float32

    def __init__(self):
        self.generate_kwargs = None

    def generate(self, **kwargs):
        self.generate_kwargs = kwargs
        return torch.tensor([[10, 11, 12, 20, 21]])


class TranscribeAudioUnitTests(unittest.TestCase):
    def test_returns_only_the_stripped_reference_transcript(self):
        processor = _FakeProcessor(f"  {REFERENCE_TRANSCRIPT}  ")
        model = _FakeModel()

        transcript = transcribe_audio(
            np.zeros(16000, dtype=np.float32),
            16000,
            processor=processor,
            model=model,
        )

        self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
        self.assertEqual(processor.decode_args[1], "transcription_only")
        self.assertEqual(processor.decode_args[0].tolist(), [[20, 21]])
        self.assertEqual(model.generate_kwargs["max_new_tokens"], 256)
        self.assertFalse(model.generate_kwargs["do_sample"])

    def test_rejects_an_empty_model_transcript(self):
        with self.assertRaisesRegex(RuntimeError, "empty transcript"):
            transcribe_audio(
                np.zeros(16000, dtype=np.float32),
                16000,
                processor=_FakeProcessor("  "),
                model=_FakeModel(),
            )

    def test_rejects_invalid_audio_input(self):
        with self.assertRaisesRegex(ValueError, "mono waveform"):
            transcribe_audio(
                np.zeros((1, 16000), dtype=np.float32),
                16000,
                processor=_FakeProcessor(),
                model=_FakeModel(),
            )


class ReferenceTranscriptIntegrationTests(unittest.TestCase):
    def test_uploaded_audio_is_transcribed_into_reference_transcript(self):
        processor = _FakeProcessor()
        model = _FakeModel()

        with tempfile.TemporaryDirectory() as tmpdir:
            audio_path = Path(tmpdir) / "reference.wav"
            sf.write(audio_path, np.zeros(8000, dtype=np.float32), 8000)

            transcript = transcribe_reference_audio(
                audio_path,
                processor=processor,
                model=model,
            )

        self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
        self.assertEqual(processor.sampling_rate, 16000)
        self.assertEqual(processor.audio.shape, (16000,))

    def test_reference_audio_is_limited_to_the_clone_prompt_duration(self):
        processor = _FakeProcessor()

        with tempfile.TemporaryDirectory() as tmpdir:
            audio_path = Path(tmpdir) / "long-reference.wav"
            sf.write(audio_path, np.zeros(25 * 8000, dtype=np.float32), 8000)

            transcribe_reference_audio(
                audio_path,
                processor=processor,
                model=_FakeModel(),
                max_seconds=20.0,
            )

        self.assertEqual(processor.audio.shape, (20 * 16000,))

    @unittest.skipUnless(shutil.which("ffmpeg"), "FFmpeg is required for M4A coverage")
    def test_m4a_upload_is_transcribed_into_reference_transcript(self):
        processor = _FakeProcessor()
        model = _FakeModel()

        with tempfile.TemporaryDirectory() as tmpdir:
            wav_path = Path(tmpdir) / "reference.wav"
            m4a_path = Path(tmpdir) / "reference clip.m4a"
            time = np.arange(8000, dtype=np.float32) / 8000
            sf.write(wav_path, 0.25 * np.sin(2 * np.pi * 440 * time), 8000)
            subprocess.run(
                [
                    "ffmpeg",
                    "-nostdin",
                    "-loglevel",
                    "error",
                    "-y",
                    "-i",
                    str(wav_path),
                    "-c:a",
                    "aac",
                    str(m4a_path),
                ],
                check=True,
            )

            transcript = transcribe_reference_audio(
                m4a_path,
                processor=processor,
                model=model,
            )

        self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
        self.assertEqual(processor.sampling_rate, 16000)
        self.assertGreaterEqual(processor.audio.shape[0], 16000)
        self.assertLess(processor.audio.shape[0], 17000)

    def test_uses_the_official_hugging_face_qwen3_asr_checkpoint(self):
        self.assertEqual(QWEN3_ASR_MODEL_ID, "Qwen/Qwen3-ASR-0.6B-hf")


if __name__ == "__main__":
    unittest.main()