#!/usr/bin/env python3
"""
Add qwen3_5_moe architecture support to Transformers
"""
import os
import sys

def add_qwen3_5_moe_support():
    """Add qwen3_5_moe to Transformers configuration"""
    
    # Find transformers installation path
    try:
        import transformers
        transformers_path = transformers.__path__[0]
    except ImportError:
        print("Transformers not installed")
        return False
    
    print(f"Transformers path: {transformers_path}")
    
    # Modify configuration_auto.py
    config_auto_file = os.path.join(transformers_path, "models", "auto", "configuration_auto.py")
    
    if not os.path.exists(config_auto_file):
        print(f"Config file not found: {config_auto_file}")
        return False
    
    with open(config_auto_file, "r") as f:
        content = f.read()
    
    # Check if already exists
    if "qwen3_5_moe" in content:
        print("qwen3_5_moe already exists in configuration")
        return True
    
    # Add to MODEL_NAMES_MAPPING
    # Find the line with "qwen2_moe": "Qwen2Moe" and add qwen3_5_moe after it
    if '"qwen2_moe": "Qwen2Moe"' in content:
        content = content.replace(
            '"qwen2_moe": "Qwen2Moe"',
            '"qwen2_moe": "Qwen2Moe",\n        "qwen3_5_moe": "Qwen3_5_MoE"'
        )
        print("Added qwen3_5_moe to MODEL_NAMES_MAPPING")
    
    # Add to MODEL_MAPPING (mapping to Qwen2MoeConfig as base)
    if '"qwen2_moe": Qwen2MoeConfig' in content:
        content = content.replace(
            '"qwen2_moe": Qwen2MoeConfig',
            '"qwen2_moe": Qwen2MoeConfig,\n        "qwen3_5_moe": Qwen2MoeConfig'
        )
        print("Added qwen3_5_moe to MODEL_MAPPING")
    
    with open(config_auto_file, "w") as f:
        f.write(content)
    
    print("Successfully added qwen3_5_moe support")
    return True

if __name__ == "__main__":
    success = add_qwen3_5_moe_support()
    sys.exit(0 if success else 1)