"""Extract text weights, add six tokens and initialize a decision head. No training."""
import argparse
import hashlib
import json
import shutil
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
from safetensors import safe_open
from transformers import AutoTokenizer
from transformers.dynamic_module_utils import get_class_from_dynamic_module
_REPO = str(Path(__file__).resolve().parents[1])
MugiDecisionConfig = get_class_from_dynamic_module('configuration_mugi.MugiDecisionConfig', _REPO, local_files_only=True)
MugiDecisionModel = get_class_from_dynamic_module('modeling_mugi.MugiDecisionModel', _REPO, local_files_only=True)
from mugi_decision.encoding import TOKENS
from mugi_decision.runtime import save_json

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--source', type=Path, default=Path('Qwen3.5-0.8B'))
    p.add_argument('--output', type=Path, default=Path('models/mugi-decision-0.8b-base'))
    p.add_argument('--no-enum', action='store_true')
    p.add_argument('--seed', type=int, default=42)
    args = p.parse_args()
    if args.output.exists() and any(args.output.iterdir()):
        raise SystemExit('Output is not empty; select a new path.')
    config = json.loads((args.source / 'config.json').read_text(encoding='utf-8'))
    if config['model_type'] != 'qwen3_5':
        raise ValueError('Expected an original Qwen3.5 multimodal checkpoint')
    torch.manual_seed(args.seed)
    tokenizer = AutoTokenizer.from_pretrained(args.source, local_files_only=True)
    tokenizer.add_special_tokens({'extra_special_tokens': TOKENS}, replace_extra_special_tokens=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'right'
    token_ids = tokenizer.convert_tokens_to_ids(TOKENS)
    cfg = MugiDecisionConfig(text_config=config['text_config'], include_enum=not args.no_enum,
                             decision_token_ids=token_ids, pad_token_id=tokenizer.pad_token_id)
    cfg.text_config.pad_token_id = tokenizer.pad_token_id
    cfg._attn_implementation = 'sdpa'
    previous_dtype = torch.get_default_dtype()
    try:
        torch.set_default_dtype(torch.bfloat16)
        model = MugiDecisionModel(cfg)
    finally:
        torch.set_default_dtype(previous_dtype)
    weights = {}
    prefix = 'model.language_model.'
    index = args.source / 'model.safetensors.index.json'
    filenames = sorted(set(json.loads(index.read_text())['weight_map'].values())) if index.exists() else [x.name for x in args.source.glob('*.safetensors')]
    if not filenames: raise FileNotFoundError('No safetensors weights found')
    for filename in filenames:
        with safe_open(args.source / filename, framework='pt', device='cpu') as handle:
            for key in handle.keys():
                if key.startswith(prefix):
                    weights[key[len(prefix):]] = handle.get_tensor(key)
    model.backbone.load_state_dict(weights, strict=True, assign=True)
    old_size = model.get_input_embeddings().num_embeddings
    size = max(old_size, len(tokenizer))
    # Preserve existing vocabulary rows; initialize new markers in unused padded rows.
    model.resize_token_embeddings(size, mean_resizing=False)
    cfg.text_config.vocab_size = size
    model.backbone.config.vocab_size = size
    with torch.no_grad():
        for token_id in token_ids:
            torch.nn.init.normal_(model.get_input_embeddings().weight[token_id], std=cfg.initializer_range)
    MugiDecisionConfig.register_for_auto_class()
    MugiDecisionModel.register_for_auto_class('AutoModel')
    args.output.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(args.output, safe_serialization=True, max_shard_size='4GB')
    tokenizer.save_pretrained(args.output)
    if (args.source / 'LICENSE').exists(): shutil.copyfile(args.source / 'LICENSE', args.output / 'LICENSE')
    metadata = dict(source=str(args.source.resolve()), source_config_sha256=hashlib.sha256((args.source/'config.json').read_bytes()).hexdigest(), source_weight_files=filenames,
                    seed=args.seed, parameters=sum(x.numel() for x in model.parameters()),
                    original_vocab_size=old_size, vocab_size=size, special_tokens=dict(zip(TOKENS, token_ids)),
                    include_enum=cfg.include_enum, state='initialized_not_finetuned',
                    note='Text backbone copied strictly. Vision encoder and vocabulary output projection omitted. Decision head and marker embeddings require training.')
    save_json(args.output / 'conversion.json', metadata)
    print(json.dumps(metadata, ensure_ascii=False, indent=2))

if __name__ == '__main__': main()