import base64
import io
import os
import random
import re
from collections.abc import Callable

import numpy as np
import torch
import torchaudio
from process_speechdata import InputSegment, StreamingInputSegment
from torchaudio.transforms import MelSpectrogram
from vllm.multimodal.media.audio import load_audio

speech_zeroemb_idx = 151667
empty_token = "<|empty|>"
mimo_audio_tokenizer = None
device = "cpu"

asr_zh_templates = [
    "请将这段语音转换为文字",
    "帮我识别这个音频文件中的内容",
    "把这段录音转成文本",
    "请转录这段语音",
    "将音频内容转换成文字格式",
    "识别并转写这段语音",
    "把语音内容写成文字",
    "转录这个音频片段",
    "将这段对话转换为文本",
    "麻烦帮我把这段录音整理成详细的文字记录",
]

asr_en_templates = [
    "Please transcribe this audio file",
    "Convert this speech recording to text",
    "Transcribe the following voice message",
    "Turn this audio into readable text",
    "Please convert the recording to written format",
    "Transcribe what you hear in this audio",
    "Convert this spoken content to text",
    "Please write down what is said in this recording",
    "Transcribe this voice recording",
    "Could you please help me transcribe this important recording?",
    "Would you mind converting this voice message into a readable text format?",
    "I'd really appreciate it if you could turn this audio file into a written document",
]

tts_zh_templates = [
    "请将这段文字转换为语音",
    "帮我把这个文本读出来",
    "将这些文字生成音频",
    "请朗读这段内容",
    "把这段话转换成语音文件",
    "生成这段文字的语音版本",
    "请用语音播报这些内容",
    "将文本转换为可听的音频",
    "帮我朗读这段文字",
    "把这些内容念出来",
]

tts_en_templates = [
    "Please convert this text to speech",
    "Turn this writing into audio",
    "Generate speech from this text",
    "Read this content out loud",
    "Convert these words to voice",
    "Create an audio version of this text",
    "Please vocalize this content",
    "Turn this text into audible format",
    "Help me convert this writing to speech",
    "Make this text into spoken audio",
]


def detect_language(text):
    if re.search(r"[\u4e00-\u9fff]", text):
        return "zh"
    else:
        return "en"


# ============================================
# Common Helper Functions - InputSegment Creation
# ============================================


def create_segment(text: str = "", audio=None) -> InputSegment:
    """Create a standard InputSegment with default zeroemb parameters"""
    return InputSegment(
        text=text,
        audio=audio,
        speech_zeroemb_idx=speech_zeroemb_idx,
        text_zeroemb_idx=empty_token,
    )


def create_streaming_segment(
    text: str, audio, tokenizer, group_size: int, audio_channels: int
) -> StreamingInputSegment:
    """Create a StreamingInputSegment"""
    return StreamingInputSegment(
        text=text,
        audio=audio,
        tokenizer=tokenizer,
        group_size=group_size,
        audio_channels=audio_channels,
        speech_zeroemb_idx=speech_zeroemb_idx,
        text_zeroemb_idx=empty_token,
    )


# ============================================
# Common Helper Functions - Common Token Segment Creation
# ============================================


def create_user_start() -> InputSegment:
    """Create user role start token"""
    return create_segment(text="<|im_start|>user\n")


def create_user_end() -> InputSegment:
    """Create role end token"""
    return create_segment(text="<|im_end|>\n")


def create_assistant_start() -> InputSegment:
    """Create assistant role start token"""
    return create_segment(text="<|im_start|>assistant\n")


def create_system_start() -> InputSegment:
    """Create system role start token"""
    return create_segment(text="<|im_start|>system\n")


def create_thinking_segment(thinking: bool = False) -> InputSegment:
    """Create thinking token, closed or open based on thinking parameter"""
    if thinking:
        return create_segment(text="<think>\n")
    else:
        return create_segment(text="<think>\n\n</think>\n")


def create_sostm_segment() -> InputSegment:
    """Create streaming output start token"""
    return create_segment(text="<|sostm|>")


def create_assistant_start_with_sostm() -> InputSegment:
    """Create assistant start token with sostm"""
    return create_segment(text="<|im_start|>assistant\n<|sostm|>")


def create_assistant_start_with_think() -> InputSegment:
    """Create assistant start token with think opening"""
    return create_segment(text="<|im_start|>assistant\n<think>\n")


# ============================================
# Common Helper Functions - Composite Segment Creation
# ============================================


def create_user_turn_with_audio(audio_tokenized, extra_text: str = None) -> list[InputSegment]:
    """Create a user turn containing audio"""
    segments = [
        create_user_start(),
        create_segment(audio=audio_tokenized),
    ]
    if extra_text:
        segments.append(create_segment(text=extra_text))
    segments.append(create_user_end())
    return segments


def create_user_turn_with_text(text: str) -> list[InputSegment]:
    """Create a text-only user turn"""
    return [
        create_user_start(),
        create_segment(text=text),
        create_user_end(),
    ]


def create_system_turn_with_voice_prompt(prompt_text: str, audio_token) -> list[InputSegment]:
    """Create a system turn with voice prompt"""
    return [
        create_system_start(),
        create_segment(text=prompt_text),
        create_segment(text="", audio=audio_token),
        create_user_end(),
    ]


def create_system_turn_text_only(system_text: str) -> list[InputSegment]:
    """Create a text-only system turn"""
    return [
        create_system_start(),
        create_segment(text=system_text),
        create_user_end(),
    ]


# ============================================
# Common Helper Functions - Multi-turn Dialogue Processing
# ============================================


def process_multiturn_messages(
    message_list: list[dict],
    user_processor: Callable[[dict], list[InputSegment]],
    assistant_processor: Callable[[dict], list[InputSegment]],
) -> list[InputSegment]:
    """
    Generic multi-turn dialogue message processing function

    Args:
        message_list: List of messages, each containing 'role' and 'content'
        user_processor: Function to process user messages
        assistant_processor: Function to process assistant messages

    Returns:
        Processed list of InputSegments
    """
    lm_prompt = []
    for message in message_list:
        role = message["role"]
        if role == "user":
            lm_prompt.extend(user_processor(message))
        elif role == "assistant":
            lm_prompt.extend(assistant_processor(message))
        else:
            raise ValueError(f"Invalid role: {role}")
    return lm_prompt


def create_text_user_message(message: dict) -> list[InputSegment]:
    """Process a text-only user message"""
    return [
        create_user_start(),
        create_segment(text=message["content"]),
        create_user_end(),
    ]


def create_text_assistant_message(message: dict) -> list[InputSegment]:
    """Process a text-only assistant message"""
    return [
        create_assistant_start(),
        create_segment(text=message["content"]),
        create_user_end(),
    ]


def create_audio_user_message(message: dict) -> list[InputSegment]:
    """Process an audio user message"""
    return [
        create_user_start(),
        create_segment(audio=preprocess_input(message["content"])),
        create_user_end(),
    ]


def append_assistant_ending(
    lm_prompt: list[InputSegment], thinking: bool = False, use_sostm: bool = False
) -> list[InputSegment]:
    """
    Append assistant ending to the prompt

    Args:
        lm_prompt: Existing prompt list
        thinking: Whether to use open thinking token
        use_sostm: Whether to use sostm token (for speech output)
    """
    if use_sostm:
        lm_prompt.append(create_assistant_start_with_sostm())
    else:
        lm_prompt.append(create_assistant_start())
        lm_prompt.append(create_thinking_segment(thinking))
    return lm_prompt


def get_asr_sft_prompt(
    input: None | str = None,
):
    """Build prompt for ASR (Automatic Speech Recognition) task"""
    audio_tokenized = preprocess_input(input)
    template = random.choice(asr_zh_templates + asr_en_templates)

    lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=template)
    lm_prompt = append_assistant_ending(lm_prompt, thinking=False)
    return lm_prompt


def resample_audio_if_needed(wav_tensor: torch.Tensor, original_sr: int):
    target_sr = 24000
    if original_sr != target_sr:
        wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
    return wav_tensor


def wav2mel(wav, device="cpu"):
    mel_transform = MelSpectrogram(
        sample_rate=mimo_audio_tokenizer.config.sampling_rate,
        n_fft=mimo_audio_tokenizer.config.nfft,
        hop_length=mimo_audio_tokenizer.config.hop_length,
        win_length=mimo_audio_tokenizer.config.window_size,
        f_min=mimo_audio_tokenizer.config.fmin,
        f_max=mimo_audio_tokenizer.config.fmax,
        n_mels=mimo_audio_tokenizer.config.n_mels,
        power=1.0,
        center=True,
    ).to(device)
    spec = mel_transform(wav[None, :])
    return torch.log(torch.clip(spec, min=1e-7)).squeeze()


def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):
    if features.size(0) != lengths.sum().item():
        raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")

    split_points = []
    current_sum = 0

    for i, seq_len in enumerate(lengths):
        if current_sum + seq_len > max_length and current_sum > 0:
            split_points.append(i)
            current_sum = seq_len.item()
        else:
            current_sum += seq_len.item()

    # Convert split points to group sizes
    group_sizes = []
    prev = 0
    for point in split_points:
        group_sizes.append(point - prev)
        prev = point
    if prev < len(lengths):
        group_sizes.append(len(lengths) - prev)

    len_groups = torch.split(lengths, group_sizes)
    feature_sizes = [group.sum().item() for group in len_groups]
    feature_groups = torch.split(features, feature_sizes)

    return feature_groups, len_groups


def encode_batch(input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000):
    feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)

    encoded_parts = []
    for features, lengths in zip(feature_groups, len_groups):
        with torch.no_grad():
            codes, _ = mimo_audio_tokenizer.encoder.encode(
                input_features=features.to(device), input_lens=lengths.to(device), return_codes_only=True
            )
            encoded_parts.append(codes)

    return torch.cat(encoded_parts, dim=-1)


def preprocess_input(input: None | str | torch.Tensor = None, device="cpu", audio_channels=4, group_size=8):
    if isinstance(input, torch.Tensor) or (isinstance(input, str) and os.path.isfile(input)):
        return "<|sosp|><|empty|><|eosp|>"

    else:
        text = input
        if (
            text.isupper() or text.islower()
        ):  # If the text only contains upper-case or lower-case letters, capitalize it.
            text = text.capitalize()
        return text


def _build_tts_system_prompt(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
    """Build system prompt for TTS task"""
    if has_voice_prompt and voice_audio_token is not None:
        return [
            create_system_start(),
            create_segment(
                # text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be: "
                text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
            ),
            create_segment(text="", audio=voice_audio_token),
            create_user_end(),
        ]
    else:
        return create_system_turn_text_only("你需要根据指定的风格指令和文本内容来生成语音。")
        # return create_system_turn_text_only(
        #     "You need to generate speech based on the specified style instructions and text content."
        # )


def _build_tts_system_prompt_no_instruct(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
    """Build system prompt for TTS task"""
    if has_voice_prompt and voice_audio_token is not None:
        return [
            create_system_start(),
            create_segment(
                # text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be:"
                text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
            ),
            create_segment(text="", audio=voice_audio_token),
            create_user_end(),
        ]
    else:
        return []


def get_tts_sft_prompt(
    input: None | str = None,
    instruct=None,
    read_text_only=True,
    prompt_speech=None,
):
    """
    Build prompt for TTS (Text-to-Speech) task

    Args:
        input: Input text
        instruct: Style instruction (e.g., "speak happily in a child's voice")
        read_text_only: Whether to read only plain text (False means text contains template)
        prompt_speech: Reference audio (for voice cloning)
    """
    assistant_prompt_audio_token = preprocess_input(prompt_speech) if prompt_speech is not None else None

    if not read_text_only:
        # Not just reading text, text contains template (template:text format)
        text = preprocess_input(input)
        lm_prompt = _build_tts_system_prompt(
            has_voice_prompt=assistant_prompt_audio_token is not None,
            voice_audio_token=assistant_prompt_audio_token,
        )
        lm_prompt.append(create_segment(text=f"<|im_start|>user\n{text}<|im_end|>\n"))
        lm_prompt.append(create_assistant_start_with_think())
    else:
        # Plain text (no instruction inside)
        language = detect_language(input)
        template = random.choice(tts_zh_templates if language == "zh" else tts_en_templates)
        text = preprocess_input(input)

        if instruct is None:
            # No instruct instruction
            lm_prompt = _build_tts_system_prompt_no_instruct(
                has_voice_prompt=assistant_prompt_audio_token is not None,
                voice_audio_token=assistant_prompt_audio_token,
            )
            lm_prompt.extend(
                [
                    create_segment(text=f"<|im_start|>user\n{template}: {text}<|im_end|>\n"),
                    create_assistant_start_with_sostm(),
                ]
            )
        else:
            # Has instruct instruction
            lm_prompt = _build_tts_system_prompt(
                has_voice_prompt=assistant_prompt_audio_token is not None,
                voice_audio_token=assistant_prompt_audio_token,
            )
            lm_prompt.append(create_segment(text=f"<|im_start|>user\n{template}: {text}({instruct})<|im_end|>\n"))
            lm_prompt.append(create_assistant_start_with_think())

    return lm_prompt


def get_audio_understanding_sft_prompt(
    input_speech,
    input_text,
    thinking=False,
    use_sostm=False,
):
    """Build prompt for audio understanding task"""
    audio_tokenized = preprocess_input(input_speech)

    lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=input_text)
    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking, use_sostm=use_sostm)
    return lm_prompt


def _build_voice_prompt_system(prompt_speech) -> list[InputSegment]:
    """Build system prompt with voice prompt"""
    return create_system_turn_with_voice_prompt(
        prompt_text="Your voice should be：", audio_token=preprocess_input(prompt_speech)
    )


def get_spoken_dialogue_sft_prompt(
    input_speech,
    system_prompt=None,
    prompt_speech=None,
    add_history=False,
):
    """
    Build prompt for spoken dialogue task

    Args:
        input_speech: Input speech
        system_prompt: System prompt text
        prompt_speech: Reference audio (for voice cloning)
        add_history: Whether to add history (Note: history variable is undefined in original code)
    """
    audio_tokenized = preprocess_input(input_speech)
    lm_prompt = []

    # Note: history variable is undefined in original code, this branch may never execute
    # To use history feature, history should be passed as a parameter
    if add_history:
        # Simplified form of adding history
        lm_prompt = create_user_turn_with_audio(audio_tokenized)
        lm_prompt.append(create_assistant_start_with_sostm())
    else:
        # Add voice prompt (if available)
        if prompt_speech:
            lm_prompt.extend(_build_voice_prompt_system(prompt_speech))

        # Add user turn
        lm_prompt.append(create_user_start())
        if system_prompt:
            lm_prompt.append(create_segment(text=system_prompt))
        lm_prompt.append(create_segment(audio=audio_tokenized))
        lm_prompt.append(create_user_end())
        lm_prompt.append(create_assistant_start_with_sostm())

    return lm_prompt


def get_spoken_dialogue_sft_multiturn_prompt(
    message_list,
    system_prompt=None,
    prompt_speech=None,
    tokenizer=None,
    group_size=8,
    audio_channels=4,
):
    """
    Build prompt for multi-turn spoken dialogue task

    Args:
        message_list: List of messages containing role and content
        system_prompt: System prompt text
        prompt_speech: Reference audio (for voice cloning)
        tokenizer: Tokenizer
        group_size: Group size
        audio_channels: Number of audio channels
    """
    lm_prompt = []

    # Add voice prompt (if available)
    if prompt_speech:
        lm_prompt.extend(
            create_system_turn_with_voice_prompt(
                prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
            )
        )

    # Define message processors
    def user_processor(msg):
        segments = [create_user_start()]
        if system_prompt:
            segments.append(create_segment(text=system_prompt))
        segments.append(create_segment(audio=preprocess_input(msg["content"])))
        segments.append(create_user_end())
        return segments

    def assistant_processor(msg):
        return [
            create_assistant_start(),
            create_streaming_segment(
                text=msg["content"]["text"],
                audio=preprocess_input(msg["content"]["audio"]),
                tokenizer=tokenizer,
                group_size=group_size,
                audio_channels=audio_channels,
            ),
            create_user_end(),
        ]

    # Process message list
    lm_prompt.extend(process_multiturn_messages(message_list, user_processor, assistant_processor))
    lm_prompt.append(create_assistant_start_with_sostm())

    return lm_prompt


def get_s2t_dialogue_sft_prompt(
    input_speech,
    thinking=False,
):
    """Build prompt for speech-to-text dialogue task"""
    audio_tokenized = preprocess_input(input_speech)

    lm_prompt = create_user_turn_with_audio(audio_tokenized)
    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
    return lm_prompt


def get_s2t_dialogue_sft_multiturn_prompt(message_list, thinking=False):
    """Build prompt for multi-turn speech-to-text dialogue task"""
    lm_prompt = process_multiturn_messages(
        message_list, user_processor=create_audio_user_message, assistant_processor=create_text_assistant_message
    )
    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
    return lm_prompt


def get_text_dialogue_sft_prompt(
    input_text,
    thinking=False,
):
    """Build prompt for text-only dialogue task"""
    lm_prompt = create_user_turn_with_text(input_text)
    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
    return lm_prompt


def get_text_dialogue_sft_multiturn_prompt(
    message_list,
    thinking=False,
):
    """Build prompt for multi-turn text-only dialogue task"""
    lm_prompt = process_multiturn_messages(
        message_list, user_processor=create_text_user_message, assistant_processor=create_text_assistant_message
    )
    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
    return lm_prompt


def get_in_context_learning_s2s_prompt(
    instruction,
    prompt_examples,
    audio,
    tokenizer=None,
    group_size=8,
    audio_channels=4,
):
    """
    Build prompt for In-Context Learning speech-to-speech task

    Args:
        instruction: Instruction text
        prompt_examples: List of examples, each containing input_audio, output_transcription, output_audio
        audio: Input audio to be processed
        tokenizer: Tokenizer
        group_size: Group size
        audio_channels: Number of audio channels
    """
    prompt = [create_segment(text=f"[Int]:{instruction}\n")]

    # Add examples
    for example in prompt_examples:
        prompt.extend(
            [
                create_segment(audio=preprocess_input(example["input_audio"])),
                create_segment(text="\n"),
                create_streaming_segment(
                    text=example["output_transcription"],
                    audio=preprocess_input(example["output_audio"]),
                    tokenizer=tokenizer,
                    group_size=group_size,
                    audio_channels=audio_channels,
                ),
                create_segment(text=" \n\n"),
            ]
        )

    # Add input audio to be processed
    prompt.extend(
        [
            create_segment(audio=preprocess_input(audio)),
            create_segment(text="\n"),
            create_sostm_segment(),
        ]
    )

    return prompt


def get_audio_data(audio_url):
    if audio_url.startswith("data:"):
        header, b64_data = audio_url.split(",", 1)
        audio_bytes = base64.b64decode(b64_data.strip())
        audio_file = io.BytesIO(audio_bytes)
    else:
        # File path
        audio_file = audio_url

    audio_signal, sr = load_audio(audio_file, sr=24000)
    audio_data = (audio_signal.astype(np.float32), sr)
    return audio_data


def to_prompt(input_segs):
    out_put = []

    for input_seg in input_segs:
        if isinstance(input_seg, StreamingInputSegment) and input_seg.text:
            out_put.append("<|sostm|>")
            if input_seg.audio is not None and isinstance(input_seg.audio, str):
                out_put.append(input_seg.text)
                out_put.append("<|eot|>")
                out_put.append("<|empty|>")
            else:
                out_put.append(input_seg.text)
                out_put.append("<|eot|>")
            out_put.append("<|eostm|>")

        else:
            out_put.append(input_seg.text)
            if input_seg.audio is not None:
                out_put.append(input_seg.audio)

    prompt = "".join(out_put)
    return prompt