"""Create a standard Hugging Face Llama-format Eather-500M model directory.

This initializes NEW RANDOM weights. It does not convert Eather custom checkpoints.
Requires transformers, tokenizers, torch, safetensors.
"""
import argparse
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--tokenizer', required=True, help='Trained Eather tokenizer.json')
    parser.add_argument('--output', default='Eather-500M-Unsloth', help='Output HF model directory')
    parser.add_argument('--config', default='config.json')
    parser.add_argument('--smoke-test', action='store_true', help='Use a tiny architecture to verify files and loading')
    args = parser.parse_args()

    import torch
    from tokenizers import Tokenizer
    from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast, AutoTokenizer, AutoModelForCausalLM

    src_cfg = json.loads(Path(args.config).read_text(encoding='utf-8'))
    raw = Path(args.tokenizer).resolve()
    if not raw.is_file():
        raise FileNotFoundError(f'Tokenizer not found: {raw}. Run tokenizer_train.py first.')
    trained = Tokenizer.from_file(str(raw))
    actual_vocab = trained.get_vocab_size(with_added_tokens=True)
    expected_vocab = int(src_cfg['vocab_size'])
    if actual_vocab != expected_vocab:
        raise ValueError(f'Tokenizer has {actual_vocab} tokens but config expects {expected_vocab}; do not silently change the model vocabulary.')
    specials = ('<|pad|>', '<|unk|>', '<|bos|>', '<|eos|>', '<|system|>', '<|user|>', '<|assistant|>', '<|end|>')
    missing = [s for s in specials if trained.token_to_id(s) is None]
    if missing:
        raise ValueError(f'Tokenizer is missing special tokens: {missing}')

    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=str(raw),
        pad_token='<|pad|>', unk_token='<|unk|>', bos_token='<|bos|>', eos_token='<|eos|>',
        additional_special_tokens=['<|system|>', '<|user|>', '<|assistant|>', '<|end|>'],
        clean_up_tokenization_spaces=False,
        model_max_length=int(src_cfg['max_seq_len']),
    )
    if len(tokenizer) != expected_vocab:
        raise AssertionError(f'HF tokenizer changed vocab: {len(tokenizer)} != {expected_vocab}')
    # A basic template for SFT training. It adds special role delimiter tokens.
    tokenizer.chat_template = (
        "{% for message in messages %}"
        "{% if message['role'] == 'system' %}{{ '<|system|>' }}"
        "{% elif message['role'] == 'user' %}{{ '<|user|>' }}"
        "{% elif message['role'] == 'assistant' %}{{ '<|assistant|>' }}"
        "{% else %}{{ '<|user|>' }}{% endif %}"
        "{{ message['content'] }}{{ '<|end|>' }}"
        "{% endfor %}"
        "{% if add_generation_prompt %}{{ '<|assistant|>' }}{% endif %}"
    )
    tokenizer.save_pretrained(out)

    d_model = int(src_cfg['d_model'])
    layers = int(src_cfg['n_layers'])
    n_heads = int(src_cfg['n_heads'])
    kv_heads = int(src_cfg['n_kv_heads'])
    ff_hidden = int(src_cfg['ffn_hidden'])
    if args.smoke_test:
        d_model, layers, n_heads, kv_heads, ff_hidden = 128, 2, 4, 2, 256
    cfg = LlamaConfig(
        vocab_size=expected_vocab,
        hidden_size=d_model,
        intermediate_size=ff_hidden,
        num_hidden_layers=layers,
        num_attention_heads=n_heads,
        num_key_value_heads=kv_heads,
        max_position_embeddings=int(src_cfg['max_seq_len']),
        rope_theta=float(src_cfg['rope_theta']),
        attention_dropout=float(src_cfg['dropout']),
        rms_norm_eps=1e-6,
        hidden_act='silu',
        attention_bias=False,
        mlp_bias=False,
        tie_word_embeddings=bool(src_cfg['tie_embeddings']),
        pad_token_id=tokenizer.pad_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=True,
        torch_dtype='float32',
    )
    model = LlamaForCausalLM(cfg)  # ~2GB RAM of random fp32 weights for full size
    count = sum(p.numel() for p in model.parameters())
    print(f'Created standard Llama-format Eather model with {count:,} parameters')
    model.save_pretrained(out, safe_serialization=True, max_shard_size='1GB')
    # Prevent loss of original config, and mark fresh random weights clearly.
    (out / 'EATHER_MODEL_INFO.json').write_text(json.dumps({
        'status': 'fresh_random_weights_NOT_PRETRAINED',
        'implementation': 'standard_transformers_llama',
        'original_custom_weights_migrated': False,
        'parameters': count,
        'source_config': src_cfg,
        'smoke_test': args.smoke_test,
    }, indent=2), encoding='utf-8')

    tok2 = AutoTokenizer.from_pretrained(out, use_fast=True)
    if len(tok2) != expected_vocab or tok2.encode('Hello world', add_special_tokens=False) != tokenizer.encode('Hello world', add_special_tokens=False):
        raise AssertionError('Tokenizer round trip failed')
    del model
    # Small smoke-test only, to keep memory/time within a predictable budget.
    if args.smoke_test:
        m = AutoModelForCausalLM.from_pretrained(out)
        with torch.no_grad():
            x = torch.tensor([[tokenizer.bos_token_id, tokenizer.eos_token_id]])
            logits = m(input_ids=x).logits
        assert logits.shape == (1, 2, expected_vocab)
        print('Small model load and forward pass passed.')
    print(f'OUTPUT: {out.resolve()}')
    print('WARNING: The weights are random. You must pretrain this model before it can chat or code.')


if __name__ == '__main__':
    main()