#!/usr/bin/env python3
"""
Standalone tool to generate VibeVoice tokenizer files from Qwen2 base.

Downloads base tokenizer from Qwen2 and patches it with VibeVoice-specific
audio tokens and chat template modifications.

Usage:
    python generate_tokenizer_files.py --output /path/to/output [--compare /path/to/reference]
"""

import argparse
import json
import os
import shutil
import tempfile
from typing import Optional, Dict, Any


# Qwen2.5 extended tokens (151646-151664)
# These are NOT in base Qwen2-7B but ARE in Qwen2.5 and Qwen2-VL
# VibeVoice uses some of these for speech: object_ref_start/end, box_start
QWEN25_EXTENDED_TOKENS = {
    "<|object_ref_start|>": 151646,  # Used as speech_start_id
    "<|object_ref_end|>": 151647,    # Used as speech_end_id
    "<|box_start|>": 151648,         # Used as speech_pad_id
    "<|box_end|>": 151649,
    "<|quad_start|>": 151650,
    "<|quad_end|>": 151651,
    "<|vision_start|>": 151652,
    "<|vision_end|>": 151653,
    "<|vision_pad|>": 151654,
    "<|image_pad|>": 151655,
    "<|video_pad|>": 151656,
    "<tool_call>": 151657,
    "</tool_call>": 151658,
    "<|fim_prefix|>": 151659,
    "<|fim_middle|>": 151660,
    "<|fim_suffix|>": 151661,
    "<|fim_pad|>": 151662,
    "<|repo_name|>": 151663,
    "<|file_sep|>": 151664,
}

# VibeVoice-specific audio tokens (IDs follow Qwen2.5's last token 151664)
VIBEVOICE_AUDIO_TOKENS = {
    "<|AUDIO|>": 151665,
    "<|audio_bos|>": 151666,
    "<|audio_eos|>": 151667,
}

# All extended tokens (Qwen2.5 + VibeVoice)
ALL_EXTENDED_TOKENS = {**QWEN25_EXTENDED_TOKENS, **VIBEVOICE_AUDIO_TOKENS}

# Chat template with audio support
# Key modification: handles part['type'] == 'audio' or 'audio_url' -> '<|AUDIO|>'
VIBEVOICE_CHAT_TEMPLATE = """{%- if tools %}
    {{- '<|im_start|>system\\n' }}
    {%- if messages[0]['role'] == 'system' %}
        {%- if messages[0]['content'] is string %}
            {{- messages[0]['content'] }}
        {%- else %}
            {%- for part in messages[0]['content'] %}
                {%- if part['type'] == 'text' %}
                    {{- part['text'] }}
                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
                    {{- '<|AUDIO|>' }}
                {%- endif %}
            {%- endfor %}
        {%- endif %}
    {%- else %}
        {{- 'You are a helpful assistant.' }}
    {%- endif %}
    {{- "\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>" }}
    {%- for tool in tools %}
        {{- "\\n" }}
        {{- tool | tojson }}
    {%- endfor %}
    {{- "\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\"name\\": <function-name>, \\"arguments\\": <args-json-object>}\\n</tool_call><|im_end|>\\n" }}
{%- else %}
    {%- if messages[0]['role'] == 'system' %}
        {{- '<|im_start|>system\\n' }}
        {%- if messages[0]['content'] is string %}
            {{- messages[0]['content'] }}
        {%- else %}
            {%- for part in messages[0]['content'] %}
                {%- if part['type'] == 'text' %}
                    {{- part['text'] }}
                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
                    {{- '<|AUDIO|>' }}
                {%- endif %}
            {%- endfor %}
        {%- endif %}
        {{- '<|im_end|>\\n' }}
    {%- else %}
        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}
    {%- endif %}
{%- endif %}
{%- for message in messages %}
    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
        {{- '<|im_start|>' + message.role + '\\n' }}
        {%- if message['content'] is string %}
            {{- message['content'] }}
        {%- else %}
            {%- for part in message['content'] %}
                {%- if part['type'] == 'text' %}
                    {{- part['text'] }}
                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
                    {{- '<|AUDIO|>' }}
                {%- endif %}
            {%- endfor %}
        {%- endif %}
        {{- '<|im_end|>\\n' }}
    {%- elif message.role == "assistant" %}
        {{- '<|im_start|>' + message.role }}
        {%- if message.content %}
            {{- '\\n' + message.content }}
        {%- endif %}
        {%- for tool_call in message.tool_calls %}
            {%- if tool_call.function is defined %}
                {%- set tool_call = tool_call.function %}
            {%- endif %}
            {{- '\\n<tool_call>\\n{"name": "' }}
            {{- tool_call.name }}
            {{- '", "arguments": ' }}
            {{- tool_call.arguments | tojson }}
            {{- '}\\n</tool_call>' }}
        {%- endfor %}
        {{- '<|im_end|>\\n' }}
    {%- elif message.role == "tool" %}
        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
            {{- '<|im_start|>user' }}
        {%- endif %}
        {{- '\\n<tool_response>\\n' }}
        {{- message.content }}
        {{- '\\n</tool_response>' }}
        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
            {{- '<|im_end|>\\n' }}
        {%- endif %}
    {%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
    {{- '<|im_start|>assistant\\n' }}
{%- endif %}"""


# Default to Qwen2.5-7B which has all the extended tokens (151646-151664)
DEFAULT_QWEN_MODEL = "Qwen/Qwen2.5-7B"


def download_qwen_tokenizer_files(output_dir: str, qwen_model: str = DEFAULT_QWEN_MODEL) -> None:
    """Download base tokenizer files from Qwen2.5 (which includes extended tokens)."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise ImportError("Please install huggingface_hub: pip install huggingface_hub")
    
    files_to_download = [
        "vocab.json",
        "merges.txt",
        "tokenizer.json",
        "tokenizer_config.json",
    ]
    
    os.makedirs(output_dir, exist_ok=True)
    
    for filename in files_to_download:
        print(f"Downloading {filename} from {qwen_model}...")
        hf_hub_download(
            repo_id=qwen_model,
            filename=filename,
            local_dir=output_dir,
            local_dir_use_symlinks=False,
        )


def patch_tokenizer_config(output_dir: str) -> None:
    """
    Patch tokenizer_config.json with VibeVoice audio tokens and chat template.
    """
    config_path = os.path.join(output_dir, "tokenizer_config.json")
    
    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)
    
    # 1. Add ALL extended tokens to added_tokens_decoder (Qwen2.5 + VibeVoice audio)
    if "added_tokens_decoder" not in config:
        config["added_tokens_decoder"] = {}
    
    for token, token_id in ALL_EXTENDED_TOKENS.items():
        if str(token_id) not in config["added_tokens_decoder"]:
            # Determine if token should be marked as "special" 
            # tool_call tokens are NOT special in Qwen2.5
            is_special = token not in ("<tool_call>", "</tool_call>", "<|fim_prefix|>", 
                                       "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>",
                                       "<|repo_name|>", "<|file_sep|>")
            config["added_tokens_decoder"][str(token_id)] = {
                "content": token,
                "lstrip": False,
                "normalized": False,
                "rstrip": False,
                "single_word": False,
                "special": is_special,
            }
    
    # 2. Add audio tokens to additional_special_tokens
    if "additional_special_tokens" not in config:
        config["additional_special_tokens"] = []
    
    for token in VIBEVOICE_AUDIO_TOKENS.keys():
        if token not in config["additional_special_tokens"]:
            config["additional_special_tokens"].append(token)
    
    # 3. Modify chat_template to support audio
    # Instead of replacing entirely, we patch the existing template to handle audio
    chat_template = config.get("chat_template", "")
    if chat_template and "<|AUDIO|>" not in chat_template:
        # Insert audio handling into the template
        # Find patterns like: {%- if part['type'] == 'text' %}
        # Add after: {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}\n                    {{- '<|AUDIO|>' }}
        audio_handler = """{%- elif part['type'] == 'audio' or part['type'] == 'audio_url' %}
                    {{- '<|AUDIO|>' }}"""
        
        # Pattern to find: after handling 'text' type, before endif
        import re
        # Look for the pattern where we handle text type and add audio handling
        pattern = r"(\{\%- if part\['type'\] == 'text' \%\}\s*\n\s*\{\{- part\['text'\] \}\})"
        replacement = r"\1\n                " + audio_handler.replace("\n", r"\n")
        
        modified_template = re.sub(pattern, replacement, chat_template)
        
        if modified_template != chat_template:
            config["chat_template"] = modified_template
            print("  - Added audio support to existing chat_template")
        else:
            # Fallback: use our predefined template
            print("  - Warning: Could not patch existing template, using predefined template")
            config["chat_template"] = VIBEVOICE_CHAT_TEMPLATE
    
    # 4. Update model_max_length for long audio support
    config["model_max_length"] = 131072
    
    # 5. Add add_bos_token if not present
    if "add_bos_token" not in config:
        config["add_bos_token"] = False
    
    # Write back
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    
    print(f"Patched {config_path}")


def patch_tokenizer_json(output_dir: str) -> None:
    """
    Patch tokenizer.json with VibeVoice audio tokens.
    """
    tokenizer_path = os.path.join(output_dir, "tokenizer.json")
    
    with open(tokenizer_path, "r", encoding="utf-8") as f:
        tokenizer = json.load(f)
    
    # Find existing token IDs to avoid duplicates
    existing_ids = set()
    if "added_tokens" in tokenizer:
        for token_entry in tokenizer["added_tokens"]:
            existing_ids.add(token_entry.get("id"))
    
    # Add ALL extended tokens (Qwen2.5 + VibeVoice audio)
    for token, token_id in ALL_EXTENDED_TOKENS.items():
        if token_id not in existing_ids:
            # Determine if token should be marked as "special"
            is_special = token not in ("<tool_call>", "</tool_call>", "<|fim_prefix|>", 
                                       "<|fim_middle|>", "<|fim_suffix|>", "<|fim_pad|>",
                                       "<|repo_name|>", "<|file_sep|>")
            tokenizer["added_tokens"].append({
                "id": token_id,
                "content": token,
                "single_word": False,
                "lstrip": False,
                "rstrip": False,
                "normalized": False,
                "special": is_special,
            })
    
    # Write back
    with open(tokenizer_path, "w", encoding="utf-8") as f:
        json.dump(tokenizer, f, indent=2, ensure_ascii=False)
    
    print(f"Patched {tokenizer_path}")


def generate_added_tokens_json(output_dir: str) -> None:
    """
    Generate added_tokens.json from tokenizer_config.json.
    """
    config_path = os.path.join(output_dir, "tokenizer_config.json")
    
    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)
    
    added_tokens = {}
    for token_id, token_info in config.get("added_tokens_decoder", {}).items():
        content = token_info.get("content")
        if content:
            added_tokens[content] = int(token_id)
    
    output_path = os.path.join(output_dir, "added_tokens.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(added_tokens, f, indent=2, ensure_ascii=False)
    
    print(f"Generated {output_path}")


def generate_special_tokens_map_json(output_dir: str) -> None:
    """
    Generate special_tokens_map.json with VibeVoice special tokens.
    """
    # Build the special tokens map
    special_tokens_map = {
        "additional_special_tokens": [],
        "eos_token": "<|endoftext|>",
        "pad_token": "<|endoftext|>",
        "unk_token": "<|endoftext|>",
    }
    
    # Add audio tokens as additional_special_tokens
    for token in VIBEVOICE_AUDIO_TOKENS.keys():
        special_tokens_map["additional_special_tokens"].append({
            "content": token,
            "lstrip": False,
            "normalized": False,
            "rstrip": False,
            "single_word": False,
        })
    
    # Add some commonly used special tokens
    common_special = ["<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>"]
    for token in common_special:
        special_tokens_map["additional_special_tokens"].append({
            "content": token,
            "lstrip": False,
            "normalized": False,
            "rstrip": False,
            "single_word": False,
        })
    
    output_path = os.path.join(output_dir, "special_tokens_map.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(special_tokens_map, f, indent=2, ensure_ascii=False)
    
    print(f"Generated {output_path}")


def generate_vibevoice_tokenizer_files(output_dir: str, qwen_model: str = DEFAULT_QWEN_MODEL) -> None:
    """
    Generate all 6 VibeVoice tokenizer files.
    
    Files generated:
    1. vocab.json - from Qwen2.5 (unchanged)
    2. merges.txt - from Qwen2.5 (unchanged)
    3. tokenizer.json - from Qwen2.5 + audio tokens
    4. tokenizer_config.json - from Qwen2.5 + audio tokens + chat_template
    5. added_tokens.json - generated from tokenizer_config.json
    6. special_tokens_map.json - generated with VibeVoice tokens
    """
    print(f"=== Generating VibeVoice tokenizer files to {output_dir} ===\n")
    
    # Step 1: Download base files from Qwen2
    download_qwen_tokenizer_files(output_dir, qwen_model)
    
    # Step 2: Patch tokenizer_config.json
    patch_tokenizer_config(output_dir)
    
    # Step 3: Patch tokenizer.json
    patch_tokenizer_json(output_dir)
    
    # Step 4: Generate added_tokens.json
    generate_added_tokens_json(output_dir)
    
    # Step 5: Generate special_tokens_map.json
    generate_special_tokens_map_json(output_dir)
    
    print(f"\n✅ All 6 tokenizer files generated in {output_dir}")


def compare_json_files(file1: str, file2: str, name: str) -> Dict[str, Any]:
    """Compare two JSON files and return differences."""
    result = {
        "name": name,
        "identical": False,
        "differences": [],
    }
    
    if not os.path.exists(file1):
        result["differences"].append(f"File 1 not found: {file1}")
        return result
    
    if not os.path.exists(file2):
        result["differences"].append(f"File 2 not found: {file2}")
        return result
    
    with open(file1, "r", encoding="utf-8") as f:
        data1 = json.load(f)
    
    with open(file2, "r", encoding="utf-8") as f:
        data2 = json.load(f)
    
    if data1 == data2:
        result["identical"] = True
        return result
    
    # Find specific differences
    def find_diff(d1, d2, path=""):
        diffs = []
        if isinstance(d1, dict) and isinstance(d2, dict):
            all_keys = set(d1.keys()) | set(d2.keys())
            for k in all_keys:
                new_path = f"{path}.{k}" if path else k
                if k not in d1:
                    diffs.append(f"Missing in generated: {new_path}")
                elif k not in d2:
                    diffs.append(f"Extra in generated: {new_path}")
                else:
                    diffs.extend(find_diff(d1[k], d2[k], new_path))
        elif isinstance(d1, list) and isinstance(d2, list):
            if len(d1) != len(d2):
                diffs.append(f"{path}: list length differs ({len(d1)} vs {len(d2)})")
            # For lists, just check if they're equal (detailed diff is complex)
            if d1 != d2:
                diffs.append(f"{path}: list content differs")
        elif d1 != d2:
            # Truncate long values for readability
            v1 = str(d1)[:100] + "..." if len(str(d1)) > 100 else str(d1)
            v2 = str(d2)[:100] + "..." if len(str(d2)) > 100 else str(d2)
            diffs.append(f"{path}: '{v1}' vs '{v2}'")
        return diffs
    
    result["differences"] = find_diff(data1, data2)
    return result


def compare_text_files(file1: str, file2: str, name: str) -> Dict[str, Any]:
    """Compare two text files."""
    result = {
        "name": name,
        "identical": False,
        "differences": [],
    }
    
    if not os.path.exists(file1):
        result["differences"].append(f"File 1 not found: {file1}")
        return result
    
    if not os.path.exists(file2):
        result["differences"].append(f"File 2 not found: {file2}")
        return result
    
    with open(file1, "r", encoding="utf-8") as f:
        content1 = f.read()
    
    with open(file2, "r", encoding="utf-8") as f:
        content2 = f.read()
    
    if content1 == content2:
        result["identical"] = True
    else:
        lines1 = content1.splitlines()
        lines2 = content2.splitlines()
        result["differences"].append(f"Line count: {len(lines1)} vs {len(lines2)}")
        
        # Find first difference
        for i, (l1, l2) in enumerate(zip(lines1, lines2)):
            if l1 != l2:
                result["differences"].append(f"First diff at line {i+1}")
                break
    
    return result


def compare_with_reference(generated_dir: str, reference_dir: str) -> None:
    """Compare generated files with reference files."""
    print(f"\n=== Comparing generated files with reference ===")
    print(f"Generated: {generated_dir}")
    print(f"Reference: {reference_dir}\n")
    
    files_to_compare = [
        ("vocab.json", "json"),
        ("merges.txt", "text"),
        ("tokenizer.json", "json"),
        ("tokenizer_config.json", "json"),
        ("added_tokens.json", "json"),
        ("special_tokens_map.json", "json"),
    ]
    
    all_identical = True
    
    for filename, file_type in files_to_compare:
        gen_file = os.path.join(generated_dir, filename)
        ref_file = os.path.join(reference_dir, filename)
        
        if file_type == "json":
            result = compare_json_files(gen_file, ref_file, filename)
        else:
            result = compare_text_files(gen_file, ref_file, filename)
        
        if result["identical"]:
            print(f"✅ {filename}: IDENTICAL")
        else:
            print(f"❌ {filename}: DIFFERENT")
            for diff in result["differences"][:5]:  # Show first 5 differences
                print(f"   - {diff}")
            if len(result["differences"]) > 5:
                print(f"   ... and {len(result['differences']) - 5} more differences")
            all_identical = False
    
    print()
    if all_identical:
        print("🎉 All files are identical!")
    else:
        print("⚠️  Some files have differences. See details above.")


def main():
    parser = argparse.ArgumentParser(
        description="Generate VibeVoice tokenizer files from Qwen2 base"
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output directory for generated files (default: temp directory)"
    )
    parser.add_argument(
        "--compare", "-c",
        type=str,
        default=None,
        help="Reference directory to compare generated files against"
    )
    parser.add_argument(
        "--qwen-model",
        type=str,
        default=DEFAULT_QWEN_MODEL,
        help=f"Qwen model to download base tokenizer from (default: {DEFAULT_QWEN_MODEL})"
    )
    
    args = parser.parse_args()
    
    # Determine output directory
    if args.output:
        output_dir = args.output
        cleanup = False
    else:
        output_dir = tempfile.mkdtemp(prefix="vibevoice_tokenizer_")
        cleanup = not args.compare  # Only cleanup if not comparing
    
    try:
        # Generate files
        generate_vibevoice_tokenizer_files(output_dir, args.qwen_model)
        
        # Compare if requested
        if args.compare:
            compare_with_reference(output_dir, args.compare)
        
        if not args.output:
            print(f"\nGenerated files are in: {output_dir}")
            
    finally:
        if cleanup and not args.output:
            print(f"\nCleaning up temporary directory: {output_dir}")
            shutil.rmtree(output_dir, ignore_errors=True)


if __name__ == "__main__":
    main()