import os, sys, json, argparse, shutil, torch, warnings
warnings.filterwarnings('ignore')

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from models.vam import VAM, VAMConfig
from transformers import AutoTokenizer


def infer_config(state_dict):
    hs = state_dict['model.embed_tokens.weight'].shape[1]
    nl = max(int(k.split('.')[2]) for k in state_dict if k.startswith('model.layers.')) + 1
    use_moe = any('expert' in k for k in state_dict if 'mlp' in k)
    return dict(hidden_size=hs, num_hidden_layers=nl, use_moe=use_moe,
                num_attention_heads=hs // 96, num_key_value_heads=hs // 192)


def convert(torch_path, output_dir, tokenizer_dir='checkpoint/omni/native_hf',
            sensevoice_dir='checkpoint/sensevoice', siglip_dir='checkpoint/siglip',
            dtype=torch.float16, device='cpu'):
    os.makedirs(output_dir, exist_ok=True)

    print(f'Loading checkpoint: {torch_path}')
    state = torch.load(torch_path, map_location=device, weights_only=True)
    cfg = infer_config(state)
    print(f'  hidden_size={cfg["hidden_size"]}, num_hidden_layers={cfg["num_hidden_layers"]}, use_moe={cfg["use_moe"]}')

    print('Creating model...')
    config = VAMConfig(**cfg)
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    model = VAM(config,
                audio_encoder_path=os.path.join(root, sensevoice_dir),
                vision_model_path=os.path.join(root, siglip_dir))

    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        print(f'  Missing keys (expected for encoders): {len(missing)}')
    if unexpected:
        print(f'  Unexpected keys: {len(unexpected)}')

    del state
    model = model.to(dtype).half()

    VAM.register_for_auto_class("AutoModelForCausalLM")
    VAMConfig.register_for_auto_class()

    print(f'Saving to {output_dir}...')
    model.save_pretrained(output_dir, safe_serialization=True)

    tokenizer_path = os.path.join(root, tokenizer_dir)
    if os.path.exists(tokenizer_path):
        for fn in ['tokenizer.json', 'tokenizer_config.json', 'generation_config.json', 'chat_template.jinja']:
            src = os.path.join(tokenizer_path, fn)
            if os.path.exists(src):
                shutil.copy2(src, os.path.join(output_dir, fn))
        print('Tokenzier files copied')
    else:
        print(f'Tokenzier not found at {tokenizer_path}, skipping')

    modeling_code = '''import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
from models.vam import VAM, VAMConfig
'''
    with open(os.path.join(output_dir, 'modeling_omni_o.py'), 'w') as f:
        f.write(modeling_code)

    config_path = os.path.join(output_dir, 'config.json')
    if os.path.exists(config_path):
        with open(config_path, 'r') as f:
            cf = json.load(f)
        cf['auto_map'] = {
            "AutoConfig": "modeling_omni_o.VAMConfig",
            "AutoModelForCausalLM": "modeling_omni_o.VAM",
        }
        cf['model_type'] = 'omni-o'
        with open(config_path, 'w') as f:
            json.dump(cf, f, indent=2, ensure_ascii=False)

    params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f'Done! Model params: {params:.2f}M')
    print(f'HF model saved to: {output_dir}')
    print(f'  VAM.from_pretrained("{output_dir}", audio_encoder_path=..., vision_model_path=...)')
    print(f'  AutoModelForCausalLM.from_pretrained("{output_dir}", trust_remote_code=True)')


if __name__ == '__main__':
    p = argparse.ArgumentParser(description='Convert omni-o .pth to HuggingFace format')
    p.add_argument('torch_path', help='Path to .pth checkpoint')
    p.add_argument('output_dir', help='Output directory for HF model')
    p.add_argument('--tokenizer_dir', default='checkpoint/omni/native_hf')
    p.add_argument('--sensevoice_dir', default='checkpoint/sensevoice')
    p.add_argument('--siglip_dir', default='checkpoint/siglip')
    p.add_argument('--dtype', default='float16')
    p.add_argument('--device', default='cpu')
    args = p.parse_args()

    convert(args.torch_path, args.output_dir, args.tokenizer_dir,
            args.sensevoice_dir, args.siglip_dir,
            getattr(torch, args.dtype), args.device)