import torch
import numpy as np
import soundfile as sf
import librosa
import re
import logging
from functools import partial
from dataclasses import asdict
from typing import Optional
from transformers import (
    AutoConfig,
    AutoModel,
    AutoTokenizer,
    WhisperProcessor,
)

import transformers.models.llama.modeling_llama as _llama_mod

if not hasattr(_llama_mod, "LLAMA_ATTENTION_CLASSES"):
    _llama_mod.LLAMA_ATTENTION_CLASSES = _llama_mod.ALL_ATTENTION_FUNCTIONS

from boson_multimodal.data_collator.higgs_audio_collator import HiggsAudioSampleCollator
from boson_multimodal.data_types import ChatMLSample, AudioContent, Message
from boson_multimodal.dataset.chatml_dataset import (
    ChatMLDatasetSample,
    prepare_chatml_sample_qwen,
)

from ..config import settings

logger = logging.getLogger(__name__)

_model: Optional[AutoModel] = None
_tokenizer: Optional[AutoTokenizer] = None
_collator: Optional[HiggsAudioSampleCollator] = None
_config: Optional[AutoConfig] = None


def _load_model():
    global _model, _tokenizer, _collator, _config
    if _model is not None:
        return

    logger.info("Loading Higgs STT on %s", settings.stt_device)
    dtype = torch.bfloat16 if settings.stt_dtype == "bfloat16" else torch.float16

    _config = AutoConfig.from_pretrained(settings.stt_model_id, trust_remote_code=True)
    _model = AutoModel.from_pretrained(
        settings.stt_model_id,
        torch_dtype=dtype,
        trust_remote_code=True,
        attn_implementation="eager",
        device_map=settings.stt_device,
    )
    _model.eval()
    _tokenizer = AutoTokenizer.from_pretrained(settings.stt_model_id)
    _model.audio_out_bos_token_id = _tokenizer.convert_tokens_to_ids(
        "<|audio_out_bos|>"
    )
    _model.audio_eos_token_id = _tokenizer.convert_tokens_to_ids("<|audio_eos|>")

    whisper_proc = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
    _collator = HiggsAudioSampleCollator(
        whisper_processor=whisper_proc,
        audio_in_token_id=_config.audio_in_token_idx,
        audio_out_token_id=_config.audio_out_token_idx,
        audio_stream_bos_id=_config.audio_stream_bos_id,
        audio_stream_eos_id=_config.audio_stream_eos_id,
        encode_whisper_embed=_config.encode_whisper_embed,
        pad_token_id=_config.pad_token_id,
        return_audio_in_tokens=_config.encode_audio_in_tokens,
        use_delay_pattern=_config.use_delay_pattern,
        round_to=1,
        audio_num_codebooks=_config.audio_num_codebooks,
        chunk_size_seconds=getattr(_config, "chunk_size_seconds", 30),
        encoder_padding_method=getattr(_config, "encoder_padding_method", "max_length"),
    )


def transcribe_segment(audio_path: str, start_time: float, end_time: float) -> str:
    _load_model()

    audio_np, sr = sf.read(audio_path)
    if sr != 16000:
        audio_np = librosa.resample(audio_np, orig_sr=sr, target_sr=16000)

    start_sample = int(start_time * 16000)
    end_sample = int(end_time * 16000)
    segment_audio = audio_np[start_sample:end_sample]

    if len(segment_audio) == 0:
        return ""

    prompt = "Transcribe the speech. Output only the spoken words in lowercase with no punctuation."
    messages = [
        Message(role="user", content=[prompt, AudioContent(audio_url="placeholder")])
    ]
    chatml = ChatMLSample(messages=messages)
    prep_fn = partial(prepare_chatml_sample_qwen, enable_thinking=True)
    input_tokens, _, _, _ = prep_fn(chatml, _tokenizer, add_generation_prompt=True)

    sample = ChatMLDatasetSample(
        input_ids=torch.LongTensor(input_tokens),
        label_ids=None,
        audio_ids_concat=None,
        audio_ids_start=None,
        audio_waveforms_concat=torch.tensor(segment_audio, dtype=torch.float32),
        audio_waveforms_start=torch.tensor([0]),
        audio_sample_rate=torch.tensor([16000]),
        audio_speaker_indices=torch.tensor([0]),
    )

    batch = asdict(_collator([sample]))
    device = next(_model.parameters()).device
    batch = {
        k: v.to(device).contiguous() if isinstance(v, torch.Tensor) else v
        for k, v in batch.items()
    }

    with torch.inference_mode():
        outputs = _model.generate(
            **batch,
            max_new_tokens=1024,
            use_cache=True,
            do_sample=False,
            stop_strings=["<|im_end|>", "<|endoftext|>"],
            tokenizer=_tokenizer,
        )

    output_ids = outputs[0] if isinstance(outputs, tuple) else outputs
    full_text = _tokenizer.decode(output_ids[0], skip_special_tokens=False)

    parts = full_text.split("assistant\n")
    hyp = parts[-1] if len(parts) > 1 else full_text
    hyp = re.sub(r"<think>.*?</think>", "", hyp, flags=re.DOTALL)
    hyp = re.sub(r"<\|.*?\|>", "", hyp).strip()

    return hyp


def unload_model():
    global _model, _tokenizer, _collator, _config
    _model = None
    _tokenizer = None
    _collator = None
    _config = None
    torch.cuda.empty_cache()