"""
🧬 HIVEMIND GPU WORKER v5
Training LoRA adapters (trl 0.7.x compatible)
"""
import gradio as gr
import os
import traceback

HF_TOKEN = os.environ.get("HF_TOKEN", "")

def train_agent(agent_name: str, skill: str, epochs: int = 2):
    """Trenira LoRA i uploaduje na HF"""
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
        from peft import LoraConfig, get_peft_model
        from trl import SFTTrainer
        from datasets import Dataset
        from huggingface_hub import HfApi, login
        from datetime import datetime
        
        log = [f"🚀 Starting: {agent_name} - {skill}"]
        
        if not HF_TOKEN:
            return "❌ HF_TOKEN not set!"
        
        login(token=HF_TOKEN)
        api = HfApi(token=HF_TOKEN)
        task_id = f"{agent_name[:8].lower().replace(' ','')}-{datetime.now().strftime('%m%d%H%M%S')}"
        
        log.append("📦 Loading TinyLlama...")
        model = AutoModelForCausalLM.from_pretrained(
            "TinyLlama/TinyLlama-1.1B-Chat-v1.0", 
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True
        )
        tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
        tokenizer.pad_token = tokenizer.eos_token
        
        log.append("🔧 Setting up LoRA...")
        lora = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj","v_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM")
        model = get_peft_model(model, lora)
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        log.append(f"   Trainable: {trainable:,}")
        
        texts = [
            f"<|user|>\nWhat is {skill}?</s>\n<|assistant|>\n{skill} is fundamental.</s>",
            f"<|user|>\nExplain {skill}</s>\n<|assistant|>\n{skill} optimizes models.</s>",
            f"<|user|>\nHow to {skill}?</s>\n<|assistant|>\nApply techniques.</s>",
        ]
        dataset = Dataset.from_dict({"text": texts})
        log.append(f"📊 Dataset: {len(dataset)} samples")
        
        log.append(f"🏋️ Training {epochs} epoch(s)...")
        trainer = SFTTrainer(
            model=model,
            train_dataset=dataset,
            dataset_text_field="text",
            max_seq_length=128,
            tokenizer=tokenizer,
            args=TrainingArguments(
                output_dir="./out",
                num_train_epochs=epochs,
                per_device_train_batch_size=1,
                learning_rate=2e-4,
                save_strategy="no",
                report_to="none",
                logging_steps=1,
                fp16=False,
            ),
        )
        result = trainer.train()
        log.append(f"✅ Done! Loss: {result.training_loss:.4f}")
        
        model.save_pretrained("./lora")
        tokenizer.save_pretrained("./lora")
        
        repo_id = f"Pista1981/hivemind-hf-{task_id}"
        log.append(f"📤 Uploading to {repo_id}...")
        api.create_repo(repo_id=repo_id, exist_ok=True, private=False)
        api.upload_folder(folder_path="./lora", repo_id=repo_id)
        log.append(f"✅ https://huggingface.co/{repo_id}")
        
        return "\n".join(log)
        
    except Exception as e:
        return f"❌ {e}\n\n{traceback.format_exc()}"


with gr.Blocks(title="Hivemind Worker") as demo:
    gr.Markdown("# 🧬 Hivemind GPU Worker v5")
    with gr.Row():
        agent_input = gr.Textbox(label="Agent", value="TestAgent")
        skill_input = gr.Textbox(label="Skill", value="optimization")
        epochs_input = gr.Slider(1, 3, value=1, step=1, label="Epochs")
    train_btn = gr.Button("🚀 Train", variant="primary")
    output = gr.Textbox(label="Output", lines=12)
    train_btn.click(fn=train_agent, inputs=[agent_input, skill_input, epochs_input], outputs=output)

demo.launch(show_error=True)