from io import BytesIO
from pathlib import Path
from typing import cast

import av
import gradio as gr
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from silero_vad import get_speech_timestamps
from silero_vad.utils_vad import OnnxWrapper
from transformers import AutoFeatureExtractor, AutoTokenizer, SequenceFeatureExtractor
from transformers.models.lfm2.configuration_lfm2 import Lfm2Config

from src.modeling.config import ModelConfig
from src.modeling.lambda_audio import LambdaAudioModel
from src.modeling.lfm2_decoder import AudioLfm2ForCausalLM
from src.modeling.loading import ModelComponents
from src.modeling.zipformer import ZipformerConfig, ZipformerModel
from src.shared.transcription import decode_audio


class SpaceConfig(ModelConfig):
    max_audio_seconds: float = 30.0
    max_new_tokens: int = 256


def load_space_components(token: str) -> ModelComponents:
    # ---------------------------------------------------------
    # Build the architecture from pinned configs, then load ASR weights.
    # Keep FP32 computation to match the existing inference path.
    # ---------------------------------------------------------
    config = SpaceConfig()
    encoder_revision = "5d82479188439e0f7a12d18a2e0ec0716785ab78"
    decoder_revision = "473d1e77c805524f4bfca67613146adc084521ec"
    encoder_config = ZipformerConfig.from_pretrained(config.encoder_model_id, revision=encoder_revision, token=token)
    decoder_config = Lfm2Config.from_pretrained(config.decoder_model_id, revision=decoder_revision, token=token)
    decoder_config._attn_implementation = "sdpa"
    model = LambdaAudioModel(ZipformerModel(encoder_config), AudioLfm2ForCausalLM(decoder_config))
    checkpoint = hf_hub_download(
        "KeisukeMiyamoto/lambda-audio-v6",
        "model-fp16.safetensors",
        revision="4746ea39b9454d86fb6f4ae5ddda826d5fb38d6a",
        token=token,
    )
    state = {name.removeprefix("model."): tensor for name, tensor in load_file(checkpoint).items()}
    model.load_state_dict(state, strict=True)
    model.eval()
    return ModelComponents(
        model=model,
        feature_extractor=cast(
            SequenceFeatureExtractor,
            AutoFeatureExtractor.from_pretrained(config.encoder_model_id, revision=encoder_revision, token=token),
        ),
        tokenizer=AutoTokenizer.from_pretrained(config.decoder_model_id, revision=decoder_revision, token=token),
    )


def read_recording(audio_path: str | None, vad_model: OnnxWrapper) -> list[bytes]:
    # ---------------------------------------------------------
    # Decode and validate audio before requesting a GPU.
    # ---------------------------------------------------------
    if audio_path is None:
        raise gr.Error("音声をアップロードするか、録音してください。")
    path = Path(audio_path)
    if path.stat().st_size > 200 * 1024 * 1024:
        raise gr.Error("200 MB以下の音声ファイルを使用してください。")
    audio = path.read_bytes()
    config = SpaceConfig()
    try:
        waveform = decode_audio(audio, config.sampling_rate)
    except (av.FFmpegError, IndexError, ValueError) as error:
        raise gr.Error("再生できる音声ファイルを使用してください。") from error
    if waveform.size <= config.sampling_rate * config.max_audio_seconds:
        return [audio]

    # ---------------------------------------------------------
    # Split longer recordings at speech boundaries, at most 30 seconds.
    # Encode lossless clips in chronological order for batched ASR.
    # ---------------------------------------------------------
    timestamps = get_speech_timestamps(
        torch.from_numpy(waveform),
        vad_model,
        sampling_rate=config.sampling_rate,
        max_speech_duration_s=config.max_audio_seconds,
        min_silence_duration_ms=3000,
        speech_pad_ms=100,
    )
    segments: list[bytes] = []
    segment_samples = int(config.max_audio_seconds * config.sampling_rate)
    for timestamp in timestamps:
        for start in range(timestamp["start"], timestamp["end"], segment_samples):
            end = min(start + segment_samples, timestamp["end"])
            buffer = BytesIO()
            sf.write(buffer, waveform[start:end], config.sampling_rate, format="WAV", subtype="FLOAT")
            segments.append(buffer.getvalue())
    return segments