import os

import torch
import torchaudio
from transformers import AutoTokenizer, AutoModelForCausalLM

from tokenizer import StepAudioTokenizer
from tts import StepAudioTTS
from utils import load_audio, speech_adjust, volumn_adjust
from transformers import BitsAndBytesConfig
from collections import OrderedDict

class StepAudio:
    def __init__(self, tokenizer_path: str, tts_path: str, llm_path: str, device_map="4x24"):
        # load optimus_ths for flash attention, make sure LD_LIBRARY_PATH has `nvidia/cuda_nvrtc/lib`
        # if not, please manually set LD_LIBRARY_PATH=xxx/python3.10/site-packages/nvidia/cuda_nvrtc/lib
        try:
            if torch.__version__ >= "2.5":
                torch.ops.load_library(os.path.join(llm_path, 'lib/liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so'))
            elif torch.__version__ >= "2.3":
                torch.ops.load_library(os.path.join(llm_path, 'lib/liboptimus_ths-torch2.3-cu121.cpython-310-x86_64-linux-gnu.so'))
            elif torch.__version__ >= "2.2":
                torch.ops.load_library(os.path.join(llm_path, 'lib/liboptimus_ths-torch2.2-cu121.cpython-310-x86_64-linux-gnu.so'))
            print("Load optimus_ths successfully and flash attn would be enabled")
        except Exception as err:
            print(f"Fail to load optimus_ths and flash attn is disabled: {err}")

        self.llm_tokenizer = AutoTokenizer.from_pretrained(
            llm_path, trust_remote_code=True
        )
        self.encoder = StepAudioTokenizer(tokenizer_path)
        self.decoder = StepAudioTTS(tts_path, self.encoder)
        device_map_4x24 = OrderedDict([('model.embed_tokens', 0), ('model.layers.0', 0), ('model.layers.1', 0), ('model.layers.2', 0), ('model.layers.3', 0), ('model.layers.4', 0), ('model.layers.5', 0), ('model.layers.6', 0), ('model.layers.7', 0), ('model.layers.8', 0), ('model.layers.9', 0), ('model.layers.10', 0), ('model.layers.11', 1), ('model.layers.12', 1), ('model.layers.13', 1), ('model.layers.14', 1), ('model.layers.15', 1), ('model.layers.16', 1), ('model.layers.17', 1), ('model.layers.18', 1), ('model.layers.19', 1), ('model.layers.20', 1), ('model.layers.21', 1), ('model.layers.22', 1), ('model.layers.23', 1), ('model.layers.24', 1), ('model.layers.25', 1), ('model.layers.26', 1), ('model.layers.27', 1), ('model.layers.28', 1), ('model.layers.29', 1), ('model.layers.30', 1), ('model.layers.31', 1), ('model.layers.32', 1), ('model.layers.33', 1), ('model.layers.34', 1), ('model.layers.35', 1), ('model.layers.36', 1), ('model.layers.37', 2), ('model.layers.38', 2), ('model.layers.39', 2), ('model.layers.40', 2), ('model.layers.41', 2), ('model.layers.42', 2), ('model.layers.43', 2), ('model.layers.44', 2), ('model.layers.45', 2), ('model.layers.46', 2), ('model.layers.47', 2), ('model.layers.48', 2), ('model.layers.49', 2), ('model.layers.50', 2), ('model.layers.51', 2), ('model.layers.52', 2), ('model.layers.53', 2), ('model.layers.54', 2), ('model.layers.55', 2), ('model.layers.56', 2), ('model.layers.57', 2), ('model.layers.58', 2), ('model.layers.59', 2), ('model.layers.60', 2), ('model.layers.61', 2), ('model.layers.62', 2), ('model.layers.63', 3), ('model.layers.64', 3), ('model.layers.65', 3), ('model.layers.66', 3), ('model.layers.67', 3), ('model.layers.68', 3), ('model.layers.69', 3), ('model.layers.70', 3), ('model.layers.71', 3), ('model.layers.72', 3), ('model.layers.73', 3), ('model.layers.74', 3), ('model.layers.75', 3), ('model.layers.76', 3), ('model.layers.77', 3), ('model.layers.78', 3), ('model.layers.79', 3), ('model.layers.80', 3), ('model.layers.81', 3), ('model.layers.82', 3), ('model.layers.83', 3), ('model.layers.84', 3), ('model.layers.85', 3), ('model.layers.86', 3), ('model.layers.87', 3), ('model.norm', 3), ('lm_head', 3)])
        nf4_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
        if device_map == "4x24":
            device_map = device_map_4x24
        self.llm = AutoModelForCausalLM.from_pretrained(
            llm_path,
            quantization_config=nf4_config,
            device_map=device_map,
            trust_remote_code=True,
            #torch_dtype=torch.bfloat16,
        )

    def __call__(
        self,
        messages: list,
        speaker_id: str,
        speed_ratio: float = 1.0,
        volumn_ratio: float = 1.0,
    ):
        text_with_audio = self.apply_chat_template(messages)
        token_ids = self.llm_tokenizer.encode(text_with_audio, return_tensors="pt")
        outputs = self.llm.generate(
            token_ids, max_new_tokens=2048, temperature=0.7, top_p=0.9, do_sample=True
        )
        output_token_ids = outputs[:, token_ids.shape[-1] : -1].tolist()[0]
        output_text = self.llm_tokenizer.decode(output_token_ids)
        output_audio, sr = self.decoder(output_text, speaker_id)
        if speed_ratio != 1.0:
            output_audio = speech_adjust(output_audio, sr, speed_ratio)
        if volumn_ratio != 1.0:
            output_audio = volumn_adjust(output_audio, volumn_ratio)
        return output_text, output_audio, sr

    def encode_audio(self, audio_path):
        audio_wav, sr = load_audio(audio_path)
        audio_tokens = self.encoder(audio_wav, sr)
        return audio_tokens

    def apply_chat_template(self, messages: list):
        text_with_audio = ""
        if messages[-1]["role"] == "assistant" and isinstance(messages[-1]["content"], str):
            prefill = messages.pop()["content"]
        else:
            prefill = None
        for msg in messages:
            role = msg["role"]
            content = msg["content"]
            if role == "user":
                role = "human"
            if isinstance(content, str):
                text_with_audio += f"<|BOT|>{role}\n{content}<|EOT|>"
            elif isinstance(content, dict):
                if content["type"] == "text":
                    text_with_audio += f"<|BOT|>{role}\n{content['text']}<|EOT|>"
                elif content["type"] == "audio":
                    audio_tokens = self.encode_audio(content["audio"])
                    text_with_audio += f"<|BOT|>{role}\n{audio_tokens}<|EOT|>"
            elif content is None:
                text_with_audio += f"<|BOT|>{role}\n"
            else:
                raise ValueError(f"Unsupported content type: {type(content)}")
        if not text_with_audio.endswith("<|BOT|>assistant\n") and prefill is None:
            text_with_audio += "<|BOT|>assistant\n"
        elif prefill is not None:
            text_with_audio += "<|BOT|>assistant\n" + prefill
        print("DEBUG: text_with_audio", text_with_audio)
        return text_with_audio


if __name__ == "__main__":
    model = StepAudio(
        encoder_path="/mnt/ys-shai-jfs/open-step1o-audio/step1o-audio-encoder",
        decoder_path="/mnt/ys-shai-jfs/open-step1o-audio/step1o-audio-decoder",
        llm_path="/mnt/ys-shai-jfs/open-step1o-audio/step1o-audio-v18",
    )

    text, audio, sr = model(
        [{"role": "user", "content": "你好，我是你的朋友，我叫小明，你叫什么名字？"}],
        "闫雨婷",
    )
    torchaudio.save("output/output_e2e_tqta.wav", audio, sr)
    text, audio, sr = model(
        [
            {
                "role": "user",
                "content": {"type": "audio", "audio": "output/output_e2e_tqta.wav"},
            }
        ],
        "闫雨婷",
    )
    torchaudio.save("output/output_e2e_aqta.wav", audio, sr)
