"""Bounded HF GPU LoRA smoke run. Existing successful traces, no live teacher."""
import json, os, time
from collections.abc import Mapping
from pathlib import Path
from huggingface_hub import HfApi,hf_hub_download

def tokenize_rows(rows, tokenizer, max_length=8192):
    output=[]; skipped=0
    for original in rows:
        row=json.loads(json.dumps(original))
        for message in row['prompt']+row['completion']:
            for call in message.get('tool_calls') or []:
                fn=call.get('function',{})
                if isinstance(fn.get('arguments'),str):
                    try: fn['arguments']=json.loads(fn['arguments'])
                    except ValueError: fn['arguments']={"raw":fn['arguments']}
        kwargs={'tools':row['tools'],'tokenize':True,'enable_thinking':True}
        prompt=tokenizer.apply_chat_template(row['prompt'],add_generation_prompt=True,**kwargs)
        full=tokenizer.apply_chat_template(row['prompt']+row['completion'],**kwargs)
        if isinstance(prompt,Mapping): prompt=prompt['input_ids']
        if isinstance(full,Mapping): full=full['input_ids']
        common=0
        for a,b in zip(prompt,full):
            if a!=b: break
            common+=1
        if len(prompt)-common>8 or common==0 or len(full)>max_length or len(full)<=common:
            skipped+=1;continue
        output.append({'input_ids':full,'labels':[-100]*common+full[common:]})
    return output,skipped

def main():
    import torch
    from transformers import AutoTokenizer,AutoModelForImageTextToText,Trainer,TrainingArguments,DataCollatorForSeq2Seq,BitsAndBytesConfig
    from peft import LoraConfig,get_peft_model,prepare_model_for_kbit_training
    from datasets import Dataset
    cfg=json.loads(os.environ['TRAIN_CONFIG']); api=HfApi(); run=os.environ['RUN_ID']
    out=Path('/tmp/adapter'); out.mkdir()
    data=hf_hub_download(cfg['data_repo'],'training/hf-corpus/exchanges.jsonl',repo_type='dataset',revision=cfg['data_revision'])
    rows=[json.loads(x) for x in Path(data).read_text().splitlines()]
    tokenizer=AutoTokenizer.from_pretrained(cfg['model'],revision=cfg['model_revision'])
    tokenizer.pad_token=tokenizer.eos_token
    examples,skipped=tokenize_rows(rows,tokenizer)
    if not examples: raise RuntimeError('No valid full-context training examples')
    print(json.dumps({'stage':'data-ready','examples':len(examples),'skipped_overlength_or_template':skipped,'gpu':torch.cuda.get_device_name(0)}),flush=True)
    model=AutoModelForImageTextToText.from_pretrained(cfg['model'],revision=cfg['model_revision'],dtype=torch.bfloat16,attn_implementation='sdpa',device_map={'':'cuda:0'},quantization_config=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type='nf4',bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True))
    model=prepare_model_for_kbit_training(model,use_gradient_checkpointing=True)
    model.config.use_cache=False
    model=get_peft_model(model,LoraConfig(r=cfg['rank'],lora_alpha=cfg['rank']*2,lora_dropout=0.05,bias='none',task_type='CAUSAL_LM',target_modules=['q_proj','k_proj','v_proj','o_proj']))
    model.print_trainable_parameters()
    args=TrainingArguments(output_dir=str(out),max_steps=cfg['steps'],learning_rate=cfg['rate'],per_device_train_batch_size=1,gradient_accumulation_steps=1,gradient_checkpointing=True,gradient_checkpointing_kwargs={'use_reentrant':False},bf16=True,logging_steps=1,save_strategy='no',report_to='none',seed=42,remove_unused_columns=False)
    trainer=Trainer(model=model,args=args,train_dataset=Dataset.from_list(examples),data_collator=DataCollatorForSeq2Seq(tokenizer,padding=True,label_pad_token_id=-100))
    result=trainer.train()
    if trainer.state.global_step<1: raise RuntimeError('No optimizer steps completed')
    trainer.save_model(str(out));tokenizer.save_pretrained(str(out))
    metrics={'status':'trained','global_step':trainer.state.global_step,'metrics':result.metrics,'examples':len(examples),'skipped':skipped,'config':cfg,'evaluation':'not run','log_history':trainer.state.log_history}
    (out/'training-metrics.json').write_text(json.dumps(metrics,indent=2))
    api.upload_folder(repo_id=cfg['data_repo'],repo_type='dataset',folder_path=out,path_in_repo=f'training/{run}/adapter',commit_message=f'Save HF GPU trained adapter {run}')
    print(json.dumps({'stage':'adapter-uploaded','global_step':trainer.state.global_step,'artifact':f'https://huggingface.co/datasets/{cfg["data_repo"]}/tree/main/training/{run}/adapter'}),flush=True)
if __name__=='__main__': main()