import os
import sys

__package__ = "scripts"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import argparse
import base64
import io
import tempfile
import warnings
import logging
import contextlib
import librosa
import soundfile as sf
import numpy as np
import torch
import gradio as gr
from queue import Queue
from threading import Thread, Lock
from PIL import Image
from pydub import AudioSegment
from transformers import AutoTokenizer, AutoModelForCausalLM, MimiModel, TextStreamer
from model.model_omni import MiniMindOmni, OmniConfig
from dataset.omni_dataset import OmniDataset
from trainer.trainer_utils import setup_seed, log_model_params
logging.getLogger().setLevel(logging.ERROR)
with contextlib.redirect_stdout(io.StringIO()):
    from funasr import AutoModel
    from funasr.utils.postprocess_utils import rich_transcription_postprocess

warnings.filterwarnings('ignore')

model_lock = Lock()

model, tokenizer, device, mimi_model, asr_model = None, None, None, None, None
voices_data = {}
builtin_voices, clone_voices = set(), set()


def default_device():
    if torch.cuda.is_available():
        return 'cuda'
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return 'mps'
    return 'cpu'


def resolve_aux_device(value, fallback):
    if value and value != 'auto':
        return value
    return fallback


def get_device_type(value):
    return torch.device(value).type


def normalize_audio(samples, sr, target_sr=16000):
    if len(samples.shape) > 1:
        samples = samples.mean(axis=1)
    if np.issubdtype(samples.dtype, np.integer):
        max_value = max(abs(np.iinfo(samples.dtype).min), np.iinfo(samples.dtype).max)
        samples = samples.astype(np.float32) / max_value
    else:
        samples = samples.astype(np.float32)
        peak = np.abs(samples).max() if samples.size else 0
        if peak > 1.0:
            samples = samples / peak
    if sr != target_sr:
        samples = librosa.resample(samples.astype(float), orig_sr=sr, target_sr=target_sr).astype(np.float32)
    return np.ascontiguousarray(samples, dtype=np.float32)


def frames_to_mimi(frames):
    codes = [f for f in frames if f and len(f) == 8]
    if not codes:
        return None
    return torch.tensor(codes, dtype=torch.long).T.unsqueeze(0)


def decode_mimi(mimi_codes):
    if mimi_codes is None or mimi_codes.numel() == 0:
        return None
    filtered = torch.where(mimi_codes >= 2049, torch.zeros_like(mimi_codes), mimi_codes).to(next(mimi_model.parameters()).device)
    with torch.no_grad():
        audio = mimi_model.decode(filtered).audio_values
    return audio.squeeze().float().cpu().numpy()


def scan_hf_models(base_dir):
    models = {}
    base_dir = os.path.abspath(base_dir)
    for d in sorted(os.listdir(base_dir), reverse=True):
        full_path = os.path.join(base_dir, d)
        if not os.path.isdir(full_path) or d.startswith('.') or d.startswith('_'):
            continue
        files = set(os.listdir(full_path))
        has_model = bool(files & {'pytorch_model.bin', 'model.safetensors', 'pytorch_model.bin.index.json', 'model.safetensors.index.json'})
        if has_model:
            models[d] = full_path
    return models


def load_hf_model(model_path):
    global model, tokenizer
    with model_lock:
        [sys.modules.pop(k) for k in list(sys.modules) if 'transformers_modules' in k]
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
        vision_encoder, vision_processor = MiniMindOmni.load_vision(args.vision_model)
        audio_encoder, audio_processor = MiniMindOmni.load_sensevoice(args.audio_encoder)
        object.__setattr__(model, 'vision_encoder', vision_encoder)
        object.__setattr__(model, 'vision_processor', vision_processor)
        object.__setattr__(model, 'audio_encoder', audio_encoder)
        object.__setattr__(model, 'audio_processor', audio_processor)
        model = model.half().eval().to(device)
        if model.audio_encoder:
            model.audio_encoder.to(device)
        if model.vision_encoder:
            model.vision_encoder.to(device)
        name = os.path.basename(model_path)
        params = sum(p.numel() for p in model.parameters()) / 1e6
        print(f'已加载 {name}，参数量：{params:.2f}M')
        return f"已加载: {name} ({params:.2f}M)"


def load_voices():
    global voices_data, builtin_voices, clone_voices
    voices_data = {}
    builtin_voices, clone_voices = set(), set()
    for name, is_builtin in [('voices.pt', True), ('voices_unseen.pt', False)]:
        path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker', name)
        if os.path.exists(path):
            data = torch.load(path, map_location=device)
            for speaker, v in data.items():
                if speaker not in voices_data:
                    voices_data[speaker] = v
                    (builtin_voices if is_builtin else clone_voices).add(speaker)


def chat_stream(prompt, audio_input=None, image_input=None, voice_name="default", history=None, temperature=0.85, max_tokens=512):
    audio_inputs, audio_lens, pixel_values, ref_codes, spk_emb = None, None, None, None, None
    asr_result = [None]
    asr_thread = None

    if audio_input is not None:
        sr, samples = audio_input
        samples = normalize_audio(samples, sr)
        inputs = model.audio_processor(samples, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
        mel = inputs.input_features.squeeze(0)
        valid_len = inputs.attention_mask.sum().item()
        audio_inputs = mel.unsqueeze(0).to(device)
        audio_lens = torch.tensor([valid_len], device=device)
        audio_token_len = valid_len or 1
        prompt = model.config.audio_special_token * audio_token_len
        samples_for_asr = samples.copy()
        def _do_asr():
            r = asr_model.generate(input=samples_for_asr, cache={}, language='auto', use_itn=True)
            asr_result[0] = rich_transcription_postprocess(r[0]['text']).strip() if r else ''
            print(f'[ASR] {asr_result[0]}')
        asr_thread = Thread(target=_do_asr)
        asr_thread.start()

    if image_input is not None:
        image = Image.open(image_input).convert('RGB') if isinstance(image_input, str) else image_input.convert('RGB')
        pixel_values = {k: v.to(device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
        prompt = (prompt + "\n\n" if prompt else "") + model.config.image_special_token * model.config.image_token_len

    if voice_name != "default" and voice_name in voices_data:
        v = voices_data[voice_name]
        ref_codes = v['ref_codes'].unsqueeze(0).to(device)
        spk_emb = v['spk_emb'].half().unsqueeze(0).to(device) if 'spk_emb' in v else None

    messages = (history or []) + [{"role": "user", "content": prompt}]
    open_thinking = bool(args.open_thinking)
    inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=open_thinking)
    x = torch.tensor(tokenizer(inputs_text).data['input_ids'], dtype=torch.long, device=device)[None, ...]

    audio_frames = []
    with torch.no_grad():
        with model_lock:
            res = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_tokens,
                                 temperature=temperature, top_p=args.top_p, stream=True,
                                 return_audio_codes=True, open_thinking=open_thinking,
                                 audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values,
                                 ref_codes=ref_codes, spk_emb=spk_emb)
            history_idx = 0
            for y, audio_frame in res:
                text_chunk = None
                if y is not None:
                    answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
                    if answer and answer[-1] != '�' and len(answer) > history_idx:
                        text_chunk = answer[history_idx:]
                        history_idx = len(answer)
                if audio_frame:
                    audio_frames.append(audio_frame)
                yield text_chunk, None, None

    if asr_thread:
        asr_thread.join()

    if audio_frames and mimi_model:
        yield None, "loading_audio", None
        mimi_codes = frames_to_mimi(audio_frames)
        audio_np = decode_mimi(mimi_codes)
        if audio_np is not None:
            yield None, (24000, audio_np), asr_result[0]
            return
    yield None, None, asr_result[0]


def launch_gradio(server_name="0.0.0.0", server_port=8888):
    voice_choices = [("default", "default")]
    for s in sorted(builtin_voices): voice_choices.append((f"[内置] {s}", s))
    for s in sorted(clone_voices): voice_choices.append((f"[克隆] {s}", s))

    def respond(message, audio, voice, chat_history, model_history, max_turns):
        text = message.get("text", "") if isinstance(message, dict) else (message or "")
        files = message.get("files", []) if isinstance(message, dict) else []
        img_path = next((f for f in files if any(f.lower().endswith(e) for e in ('.png','.jpg','.jpeg','.gif','.bmp','.webp'))), None)

        if not text and audio is None and img_path is None:
            yield chat_history + [{"role": "assistant", "content": "请输入文本、上传图片或录制音频"}], gr.update(), gr.update(), model_history, ""
            return

        if audio is not None:
            sr, samples = audio
            display_samples = samples.mean(axis=1) if len(samples.shape) > 1 else samples
            wav_path = tempfile.NamedTemporaryFile(suffix='.wav', delete=False).name
            sf.write(wav_path, display_samples, sr)
            chat_history = chat_history + [{"role": "user", "content": {"path": wav_path}}]
        elif img_path:
            chat_history = chat_history + [{"role": "user", "content": {"path": img_path}}, {"role": "user", "content": text or "请描述这张图片"}]
        else:
            chat_history = chat_history + [{"role": "user", "content": text}]

        response_text = ""
        final_audio = None
        asr_text = None
        chat_history = chat_history + [{"role": "assistant", "content": ""}]
        yield chat_history, gr.update(value=None), None, model_history, ""

        hist = model_history[-(max_turns * 2):] if max_turns > 0 else []
        for text_chunk, audio_data, asr in chat_stream(
            text or "", audio_input=audio, image_input=img_path,
            voice_name=voice, history=hist, temperature=0.7, max_tokens=512
        ):
            if text_chunk:
                response_text += text_chunk
                chat_history[-1]["content"] = response_text
                yield chat_history, gr.update(), None, model_history, gr.update()
            if audio_data == "loading_audio":
                yield chat_history, gr.update(), None, model_history, '<div style="text-align:center;padding:10px 0;color:#aaa;font-size:13px;animation:pulse 1.5s ease-in-out infinite">正在组装语音（受限于 Gradio 特性，此处为非流式解码）...</div>'
            elif audio_data:
                final_audio = audio_data
            if asr is not None:
                asr_text = asr

        user_text = asr_text if asr_text else (text or "")
        if user_text:
            model_history = model_history + [{"role": "user", "content": user_text}]
        if response_text:
            model_history = model_history + [{"role": "assistant", "content": response_text}]

        yield chat_history, final_audio if final_audio else gr.update(), None, model_history, ""

    with gr.Blocks(title="MiniMind-O", js="()=>{new MutationObserver(()=>{const m=document.getElementById('mic-box');if(!m)return;const h=!!m.querySelector('audio');document.body.classList.toggle('has-audio',h);const t=document.querySelector('textarea');if(t){t.placeholder=h?'已加载语音，点击发送':'输入文本';t.disabled=h}}).observe(document.body,{childList:true,subtree:true})}", css=".app{padding-top:6px!important} #component-0{gap:6px!important} #component-1{padding:2px 0!important;margin:0!important;min-height:0!important;border:none!important} #component-1 .padding{padding:0!important} #chatbox img{max-width:120px!important;max-height:120px!important;border-radius:8px} textarea{overflow-y:hidden!important;height:auto!important;min-height:30px!important;max-height:60px!important} #mic-box{max-height:150px!important;overflow:hidden!important} #mic-box .wrap span.or,#mic-box .wrap span:first-child{display:none!important} #mic-box .wrap{font-size:0!important;min-height:40px!important;padding:8px!important} #mic-box .wrap::after{content:'上传/录音';font-size:14px!important} #mic-box .mic-select{display:none!important} .has-audio textarea{opacity:0.4!important;pointer-events:none!important} .has-audio .upload-button,.has-audio [data-testid='upload-button']{display:none!important} @keyframes pulse{0%,100%{opacity:1}50%{opacity:0.4}}") as demo:
        gr.HTML('<div style="text-align:center;margin:2px 0"><span style="font-size:1.2rem;font-weight:bold;font-style:italic">MiniMind-O</span> <span style="color:#999;font-size:0.8rem">text / image / audio → text + audio</span></div>')

        chatbot = gr.Chatbot(label="", height=380, elem_id="chatbox", type="messages")
        model_history = gr.State([])
        audio_status = gr.HTML("", elem_id="audio-status")
        audio_out = gr.Audio(label="语音回复", autoplay=True, elem_id="audio-out")
        with gr.Row(equal_height=True):
            with gr.Column(scale=0, min_width=160):
                aud = gr.Audio(sources=["upload", "microphone"], type="numpy", show_label=False, elem_id="mic-box")
            with gr.Column(scale=4):
                msg = gr.MultimodalTextbox(placeholder="输入文本", show_label=False, submit_btn="发送")
        with gr.Row():
            voice_dd = gr.Dropdown(choices=voice_choices, value="default", label="音色选择", scale=0, min_width=140)
            turns_dd = gr.Dropdown(choices=[0, 2, 4, 6, 8], value=0, label="多轮记忆", scale=0, min_width=120)
            if model_dict:
                model_dd = gr.Dropdown(choices=list(model_dict.keys()), value=current_model_name, label="模型选择", scale=1, min_width=180)
                status = gr.Textbox(value=f"已加载: {current_model_name}", label="状态", interactive=False, scale=2)
                model_dd.change(lambda n: load_hf_model(model_dict[n]), [model_dd], [status])

        msg.submit(respond, [msg, aud, voice_dd, chatbot, model_history, turns_dd], [chatbot, audio_out, aud, model_history, audio_status])

    demo.queue().launch(server_name=server_name, server_port=server_port)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="MiniMind-O Gradio Demo")
    parser.add_argument('--load_from', default='./', type=str, help="transformers模型扫描目录")
    parser.add_argument('--audio_encoder', default='../model/SenseVoiceSmall', type=str)
    parser.add_argument('--vision_model', default='../model/siglip2-base-p32-256-ve', type=str)
    parser.add_argument('--mimi_path', default='../model/mimi', type=str)
    parser.add_argument('--device', default=default_device(), type=str)
    parser.add_argument('--asr_device', default='auto', type=str, help='ASR设备；auto 在 MPS 主设备下使用 cpu，其它情况跟随 --device')
    parser.add_argument('--mimi_device', default='auto', type=str, help='Mimi 解码设备；auto 跟随 --device')
    parser.add_argument('--open_thinking', default=0, type=int, choices=[0, 1])
    parser.add_argument('--top_p', default=0.85, type=float)
    parser.add_argument('--port', default=8888, type=int)
    args = parser.parse_args()

    device = args.device
    asr_device = resolve_aux_device(args.asr_device, 'cpu' if get_device_type(device) == 'mps' else device)
    mimi_device = resolve_aux_device(args.mimi_device, device)
    model_dict = scan_hf_models(args.load_from)
    if not model_dict:
        print(f"未在 {os.path.abspath(args.load_from)} 找到 transformers 模型")
        exit(1)
    current_model_name = list(model_dict.keys())[0]
    load_hf_model(model_dict[current_model_name])

    try:
        mimi_model = MimiModel.from_pretrained(args.mimi_path).eval().to(mimi_device)
        if get_device_type(mimi_device) != 'cpu':
            mimi_model = mimi_model.half()
        print(f'Mimi model loaded on {mimi_device}')
    except Exception:
        mimi_model = None
        print('Mimi model not found, audio output disabled')

    with contextlib.redirect_stdout(io.StringIO()):
        asr_model = AutoModel(model=args.audio_encoder, trust_remote_code=True, device=asr_device, disable_update=True)
    load_voices()
    print(f'Voices loaded: {list(voices_data.keys()) or "none"}')
    launch_gradio(server_port=args.port)