from transformers import AutoModel, AutoProcessor, AutoTokenizer
import argparse
import os
import sys
import torch
import numpy as np
import random
import time
import types

# Add models directory to path for imports
_script_dir = os.path.dirname(os.path.abspath(__file__))
_models_dir = os.path.join(os.path.dirname(_script_dir), "models")
if _models_dir not in sys.path:
    sys.path.insert(0, _models_dir)

from generation_functions import Fast_dLLM_QwenForCausalLM

parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="output_models/j63_mdm_8node_no_resize_2048_always_mask_eos_bsz256/checkpoint-3400")
args = parser.parse_args()

model_name = args.model_name

# 加载VL模型和处理器
model = AutoModel.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="cuda:0",
    trust_remote_code=True
)

# 使用本地模型路径对应的处理器，如果不存在则使用官方版本
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
tokenizer = AutoTokenizer.from_pretrained(model_name)
processor.tokenizer = tokenizer

model.mdm_sample = types.MethodType(Fast_dLLM_QwenForCausalLM.mdm_sample_cache, model)

bd_size = model.bd_size
# 初始化对话历史
messages = [
]

def fix_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

fix_seed(42)
print("VL聊天机器人已启动！输入 'exit' 退出对话。")
print("输入 'clear' 清空对话记录。")
print("注意：这是纯文本模式，不支持图像输入。")
print("-" * 50)

while True:
    # 获取用户输入
    user_input = input("用户: ").strip()
    threshold = float(input("阈值: ").strip())

    # 检查是否退出
    if user_input.lower() == "exit":
        print("再见！")
        break

    # 检查是否清空对话记录
    if user_input.lower() == "clear":
        messages = messages[:1] if messages[0]["role"] == "system" else messages[:0]
        print("对话记录已清空！")
        continue

    # 如果用户输入为空，继续下一轮
    if not user_input:
        continue

    # 添加用户消息到对话历史（纯文本格式）
    messages.append({
        "role": "user", 
        "content": [
            {"type": "text", "text": user_input}
        ]
    })

    # 准备推理输入
    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    print(f"模板化文本: {text}")
    
    # 处理输入（纯文本，无图像和视频）
    inputs = processor(
        text=[text],
        images=None,  # 纯文本模式，无图像
        videos=None,  # 纯文本模式，无视频
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(model.device)

    # 开始生成
    start_time = time.time()
    
    with torch.no_grad():
        generated_ids = model.mdm_sample(
            input_ids=inputs["input_ids"],
            tokenizer=processor.tokenizer,
            block_size=bd_size,
            max_tokens=512,
            mask_id=processor.tokenizer.encode("|<MASK>|")[0], 
            threshold=threshold
        )
    
    end_time = time.time()
    
    # 提取生成的新token
    generated_ids_trimmed = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    
    # 解码输出
    response = processor.batch_decode(
        generated_ids_trimmed, 
        skip_special_tokens=True, 
        clean_up_tokenization_spaces=False
    )[0]

    print(f"生成时间: {end_time - start_time:.2f} 秒")
    print(f"生成速度: {len(generated_ids_trimmed[0]) / (end_time - start_time):.2f} tokens/s")
    print(f"AI: {response}")

    # 添加AI回复到对话历史
    messages.append({
        "role": "assistant", 
        "content": response
    })

    print("-" * 50)
