import math
import sys
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
from scipy.signal import resample_poly
from transformers import WhisperFeatureExtractor


REPO_ROOT = Path(__file__).resolve().parents[2]
QWEN_ASR_CORE = REPO_ROOT / "Qwen3-ASR" / "qwen_asr" / "core"
if str(QWEN_ASR_CORE) not in sys.path:
    sys.path.insert(0, str(QWEN_ASR_CORE))

from transformers_backend.modeling_qwen3_asr import (  # noqa: E402
    Qwen3ASRForConditionalGeneration,
)


TORCH_DTYPES = {
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}


def get_torch_dtype(dtype: str) -> torch.dtype:
    if dtype not in TORCH_DTYPES:
        raise ValueError(f"Unsupported dtype: {dtype}")
    return TORCH_DTYPES[dtype]


def load_waveform(audio_path: str, target_sr: int = 16000) -> np.ndarray:
    audio, sr = sf.read(audio_path, dtype="float32", always_2d=False)
    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim == 2:
        audio = audio.mean(axis=-1)
    if sr != target_sr:
        divisor = math.gcd(int(sr), int(target_sr))
        up = int(target_sr // divisor)
        down = int(sr // divisor)
        audio = resample_poly(audio, up=up, down=down).astype(np.float32)
    return audio


def configure_feature_extractor_for_audio(feature_extractor: WhisperFeatureExtractor, waveform: np.ndarray) -> None:
    required_seconds = max(1, math.ceil(waveform.shape[0] / float(feature_extractor.sampling_rate)))
    if required_seconds <= feature_extractor.chunk_length:
        return
    feature_extractor.chunk_length = required_seconds
    feature_extractor.n_samples = int(required_seconds * feature_extractor.sampling_rate)
    feature_extractor.nb_max_frames = feature_extractor.n_samples // feature_extractor.hop_length


def extract_mel_features(model_path: str, audio_path: str) -> tuple[np.ndarray, int]:
    feature_extractor = WhisperFeatureExtractor.from_pretrained(model_path)
    waveform = load_waveform(audio_path)
    configure_feature_extractor_for_audio(feature_extractor, waveform)
    outputs = feature_extractor(
        waveform,
        sampling_rate=16000,
        return_attention_mask=True,
        return_tensors="np",
    )
    input_features = outputs["input_features"][0].astype(np.float32)
    feature_len = int(outputs["attention_mask"][0].sum())
    return input_features, feature_len


def split_mel_features(input_features: np.ndarray, feature_len: int, chunk_frames: int) -> list[tuple[np.ndarray, int]]:
    chunks = []
    start = 0
    while start < feature_len:
        cur_len = min(chunk_frames, feature_len - start)
        chunk = np.zeros((input_features.shape[0], chunk_frames), dtype=np.float32)
        chunk[:, :cur_len] = input_features[:, start : start + cur_len]
        chunks.append((chunk, cur_len))
        start += cur_len
    return chunks


def load_audio_encoder(model_path: str, dtype: str = "float32", device: str = "cpu") -> torch.nn.Module:
    model = Qwen3ASRForConditionalGeneration.from_pretrained(
        model_path,
        dtype=get_torch_dtype(dtype),
    )
    model = model.to(device)
    model.eval()
    tower = model.thinker.audio_tower
    tower.config._attn_implementation = "eager"
    for layer in tower.layers:
        layer.self_attn.config._attn_implementation = "eager"
    tower.eval()
    return tower


def get_chunk_output_length(length: torch.Tensor) -> torch.Tensor:
    length = length.to(torch.int64)
    length = torch.div(length + 1, 2, rounding_mode="floor")
    length = torch.div(length + 1, 2, rounding_mode="floor")
    length = torch.div(length + 1, 2, rounding_mode="floor")
    return length


def get_chunk_output_length_value(length: int) -> int:
    value = int(length)
    value = (value + 1) // 2
    value = (value + 1) // 2
    value = (value + 1) // 2
    return value


class StaticChunkAudioEncoder(torch.nn.Module):
    def __init__(self, tower: torch.nn.Module, chunk_frames: int = 100):
        super().__init__()
        self.tower = tower
        self.chunk_frames = int(chunk_frames)
        self.max_aftercnn_len = int(get_chunk_output_length(torch.tensor(self.chunk_frames)).item())

    def forward(self, input_features: torch.Tensor, feature_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size = input_features.shape[0]
        outputs = []
        valid_lens = []

        for batch_idx in range(batch_size):
            current_feature_len = feature_len.reshape(-1)[batch_idx].to(torch.int64)
            padded = input_features[batch_idx, :, : self.chunk_frames].unsqueeze(0).unsqueeze(0)

            padded_embed = torch.nn.functional.gelu(self.tower.conv2d1(padded))
            padded_embed = torch.nn.functional.gelu(self.tower.conv2d2(padded_embed))
            padded_embed = torch.nn.functional.gelu(self.tower.conv2d3(padded_embed))

            batch, channels, freq, time_steps = padded_embed.size()
            padded_embed = self.tower.conv_out(
                padded_embed.permute(0, 3, 1, 2).contiguous().view(batch, time_steps, channels * freq)
            )
            positional_embedding = self.tower.positional_embedding.positional_embedding[:time_steps, :]
            padded_embed = padded_embed + positional_embedding.unsqueeze(0).to(padded_embed.dtype)
            hidden_states = padded_embed.squeeze(0)

            valid_len = get_chunk_output_length(current_feature_len.unsqueeze(0))[0].to(torch.int32)
            valid_positions = torch.arange(time_steps, device=hidden_states.device, dtype=torch.int32) < valid_len
            allowed = valid_positions[:, None] & valid_positions[None, :]

            zeros = torch.zeros((1, 1, time_steps, time_steps), dtype=hidden_states.dtype, device=hidden_states.device)
            minus_inf = torch.full(
                (1, 1, time_steps, time_steps),
                torch.finfo(hidden_states.dtype).min,
                dtype=hidden_states.dtype,
                device=hidden_states.device,
            )
            attention_mask = torch.where(allowed.unsqueeze(0).unsqueeze(0), zeros, minus_inf)
            cu_seqlens = torch.stack(
                (
                    torch.zeros((), dtype=torch.int32, device=hidden_states.device),
                    valid_len,
                )
            )

            for encoder_layer in self.tower.layers:
                hidden_states = encoder_layer(
                    hidden_states,
                    cu_seqlens=cu_seqlens,
                    attention_mask=attention_mask,
                )[0]

            hidden_states = self.tower.ln_post(hidden_states)
            hidden_states = self.tower.proj1(hidden_states)
            hidden_states = self.tower.act(hidden_states)
            hidden_states = self.tower.proj2(hidden_states)

            outputs.append(hidden_states.unsqueeze(0))
            valid_lens.append(valid_len)

        return torch.cat(outputs, dim=0), torch.stack(valid_lens, dim=0)


def run_chunked_torch(
    model_path: str,
    input_features: np.ndarray,
    feature_len: int,
    chunk_frames: int = 100,
    dtype: str = "float32",
    device: str = "cpu",
) -> np.ndarray:
    tower = load_audio_encoder(model_path=model_path, dtype=dtype, device=device)
    wrapper = StaticChunkAudioEncoder(tower=tower, chunk_frames=chunk_frames).to(device).eval()

    outputs = []
    for chunk, chunk_len in split_mel_features(input_features, feature_len, chunk_frames):
        chunk_tensor = torch.from_numpy(chunk).unsqueeze(0).to(device=device, dtype=get_torch_dtype(dtype))
        chunk_len_tensor = torch.tensor([chunk_len], dtype=torch.int32, device=device)
        with torch.no_grad():
            features, valid_len = wrapper(chunk_tensor, chunk_len_tensor)
        outputs.append(features[0, : valid_len[0].item()].detach().cpu().numpy())

    return np.concatenate(outputs, axis=0)