# Prompt: 弄成input parameters
# Version 2: Uses QwenVLMDMModel class from qwen_vl_mdm_model.py
import argparse
import os
import sys
from typing import Optional

# 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 qwen_vl_mdm_model import QwenVLMDMModel


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run Qwen-VL chatbot generation with configurable inputs (using QwenVLMDMModel).")
    parser.add_argument(
        "--model-name",
        default="output_models/debug_alpaca",
        help="Model identifier or local path passed to AutoModel/AutoTokenizer.",
    )
    parser.add_argument(
        "--processor-name",
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Processor identifier passed to AutoProcessor.",
    )
    parser.add_argument(
        "--image",
        default="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
        help="Image URL or path used in the user message. Leave empty to skip image input.",
    )
    parser.add_argument(
        "--prompt",
        default="Describe this image.",
        help="Text prompt included in the user message.",
    )
    parser.add_argument(
        "--block-size",
        type=int,
        default=8,
        help="Block size passed to model.mdm_sample.",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=128,
        help="Maximum number of tokens to generate.",
    )
    parser.add_argument(
        "--mask-token",
        default="|<MASK>|",
        help="Mask token string used to compute mask_id.",
    )
    parser.add_argument(
        "--device",
        default="cuda",
        help="Device to move processor outputs to (e.g., 'cuda' or 'cpu').",
    )
    return parser


def main() -> None:
    args = build_parser().parse_args()

    # Initialize model using QwenVLMDMModel class
    model = QwenVLMDMModel(
        model_name=args.model_name,
        processor_name=args.processor_name,
        device=args.device,
        block_size=args.block_size,
        max_tokens=args.max_tokens,
        mask_token=args.mask_token,
    )

    # Generate using the model's generate method
    output_text = model.generate(
        prompt=args.prompt,
        image=args.image if args.image else None,
        max_tokens=args.max_tokens,
        block_size=args.block_size,
    )
    
    # Print as list to match original chatbot_mdm.py output format
    print([output_text])


if __name__ == "__main__":
    main()

