"""
Flint-1.2B: Convert JAX checkpoint to HuggingFace format
Converts trained JAX parameters to a HuggingFace-compatible safetensors model.

Usage:
    python convert_to_hf.py --checkpoint /kaggle/working/checkpoints/flint-1b/step_00005700
    python convert_to_hf.py --checkpoint /path/to/checkpoint --push --repo tekkmaven/flint-1.2B
"""

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import save_file
from transformers import LlamaConfig, LlamaForCausalLM, AutoTokenizer
from huggingface_hub import HfApi


def convert_jax_to_hf(checkpoint_path: str, output_dir: str = "./hf_model"):
    """Convert JAX checkpoint to HuggingFace LlamaForCausalLM format."""
    
    print(f"[Convert] Loading JAX checkpoint from: {checkpoint_path}")
    
    # Load JAX checkpoint
    import orbax.checkpoint as ocp
    checkpointer = ocp.PyTreeCheckpointer()
    state = checkpointer.restore(checkpoint_path)
    params = state['params']
    
    # Create HF config
    config = LlamaConfig(
        vocab_size=49216,
        hidden_size=2048,
        intermediate_size=5504,
        num_hidden_layers=22,
        num_attention_heads=32,
        num_key_value_heads=4,
        hidden_act="silu",
        max_position_embeddings=2048,
        rms_norm_eps=1e-5,
        rope_theta=500000.0,
        tie_word_embeddings=True,
        attention_bias=False,
        mlp_bias=False,
        torch_dtype="bfloat16",
    )
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Convert parameters
    print("[Convert] Mapping JAX parameters to HuggingFace format...")
    
    state_dict = {}
    
    # Embedding
    state_dict['model.embed_tokens.weight'] = torch.from_numpy(
        np.array(params['embed'])
    ).to(torch.bfloat16)
    
    # Transformer layers
    for i in range(22):
        layer = params['layers'][i]
        prefix = f'model.layers.{i}'
        
        # Attention
        state_dict[f'{prefix}.self_attn.q_proj.weight'] = torch.from_numpy(
            np.array(layer['wq']).T
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.self_attn.k_proj.weight'] = torch.from_numpy(
            np.array(layer['wk']).T
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.self_attn.v_proj.weight'] = torch.from_numpy(
            np.array(layer['wv']).T
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.self_attn.o_proj.weight'] = torch.from_numpy(
            np.array(layer['wo']).T
        ).to(torch.bfloat16)
        
        # FFN (SwiGLU: gate_proj, up_proj, down_proj)
        state_dict[f'{prefix}.mlp.gate_proj.weight'] = torch.from_numpy(
            np.array(layer['w_gate']).T
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.mlp.up_proj.weight'] = torch.from_numpy(
            np.array(layer['w_up']).T
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.mlp.down_proj.weight'] = torch.from_numpy(
            np.array(layer['w_down']).T
        ).to(torch.bfloat16)
        
        # Norms
        state_dict[f'{prefix}.input_layernorm.weight'] = torch.from_numpy(
            np.array(layer['attn_norm'])
        ).to(torch.bfloat16)
        state_dict[f'{prefix}.post_attention_layernorm.weight'] = torch.from_numpy(
            np.array(layer['ffn_norm'])
        ).to(torch.bfloat16)
    
    # Final norm
    state_dict['model.norm.weight'] = torch.from_numpy(
        np.array(params['final_norm'])
    ).to(torch.bfloat16)
    
    # LM head (tied with embeddings)
    state_dict['lm_head.weight'] = state_dict['model.embed_tokens.weight']
    
    # Save
    print(f"[Convert] Saving to {output_dir}")
    
    # Save config
    config.save_pretrained(output_dir)
    
    # Save model weights
    save_file(state_dict, os.path.join(output_dir, "model.safetensors"))
    
    # Save tokenizer
    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
    tokenizer.add_special_tokens({
        "additional_special_tokens": [
            "<think>", "</think>",
            "<tool_call>", "</tool_call>",
            "<tool_response>", "</tool_response>",
            "<observe>", "</observe>",
            "<act>", "</act>",
            "<|pad|>"
        ]
    })
    tokenizer.pad_token = "<|pad|>"
    tokenizer.save_pretrained(output_dir)
    
    # Verify
    print("[Convert] Verifying model loads correctly...")
    model = LlamaForCausalLM.from_pretrained(output_dir, torch_dtype=torch.bfloat16)
    print(f"[Convert] ✓ Model loaded: {sum(p.numel() for p in model.parameters()):,} params")
    
    return output_dir


def push_to_hub(model_dir: str, repo_id: str):
    """Push converted model to HuggingFace Hub."""
    print(f"[Hub] Pushing to {repo_id}...")
    
    api = HfApi()
    api.upload_folder(
        folder_path=model_dir,
        repo_id=repo_id,
        repo_type="model",
        commit_message="Upload trained Flint-1.2B weights",
    )
    
    print(f"[Hub] ✓ Pushed to https://huggingface.co/{repo_id}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert Flint JAX checkpoint to HF format")
    parser.add_argument("--checkpoint", type=str, required=True,
                       help="Path to JAX checkpoint directory")
    parser.add_argument("--output", type=str, default="./hf_model",
                       help="Output directory for HF model")
    parser.add_argument("--push", action="store_true",
                       help="Push to HuggingFace Hub")
    parser.add_argument("--repo", type=str, default="tekkmaven/flint-1.2B",
                       help="HF repo ID for pushing")
    
    args = parser.parse_args()
    
    output_dir = convert_jax_to_hf(args.checkpoint, args.output)
    
    if args.push:
        push_to_hub(output_dir, args.repo)