import modal

# Define the environment with Unsloth pre-installed
image = (
    modal.Image.from_registry("nvidia/cuda:12.1.1-devel-ubuntu22.04", add_python="3.10")
    .apt_install("git")
    .pip_install(
        "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git",
        "xformers",
        "trl",
        "peft",
        "accelerate",
        "bitsandbytes"
    )
    .add_local_file("verbs.jsonl", "/root/verbs.jsonl")
)

app = modal.App("step-zero-finetune")
vol = modal.Volume.from_name("step-zero-volume", create_if_missing=True)

@app.function(
    image=image, 
    gpu="A10G", 
    timeout=1800,
    volumes={"/vol": vol}
)
def train_model():
    from unsloth import FastLanguageModel, is_bfloat16_supported
    from datasets import load_dataset
    from trl import SFTTrainer
    from transformers import TrainingArguments
    import json

    print("Loading Base Model: Nemotron-Mini-4B...")
    max_seq_length = 1024
    
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = "nvidia/Nemotron-Mini-4B-Instruct",
        max_seq_length = max_seq_length,
        dtype = None,
        load_in_4bit = False,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r = 16,
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                          "gate_proj", "up_proj", "down_proj",],
        lora_alpha = 16,
        lora_dropout = 0,
        bias = "none",
        use_gradient_checkpointing = "unsloth",
        random_state = 3407,
        use_rslora = False,
    )

    print("Loading and formatting verbs.jsonl dataset...")
    dataset = load_dataset("json", data_files="/root/verbs.jsonl", split="train")
    
    def format_prompt(examples):
        texts = []
        for instruction, output in zip(examples["instruction"], examples["output"]):
            try:
                task = json.loads(output)["task"]
            except:
                task = output
                
            parts = instruction.split(". Failures: ")
            goal = parts[0].replace("Goal: ", "").strip()
            failures = parts[1].split(".")[0].strip()
            
            system_msg = "You are a cognitive pacemaker. Break down goals into extremely tiny, atomic physical actions under 8 words."
            prompt = f"<extra_id_0>System\n{system_msg}\n\n"
            prompt += f"<extra_id_1>User\nGoal: {goal}\nCompleted Tasks: None\nFailures: {failures}\nOutput the NEXT step.\n"
            prompt += f"<extra_id_1>Assistant\n{task}\n"
            texts.append(prompt)
        return {"text": texts}

    dataset = dataset.map(format_prompt, batched=True)

    print("Starting LoRA fine-tuning for Syntactic Enforcement...")
    trainer = SFTTrainer(
        model = model,
        tokenizer = tokenizer,
        train_dataset = dataset,
        dataset_text_field = "text",
        max_seq_length = max_seq_length,
        dataset_num_proc = 2,
        packing = False,
        args = TrainingArguments(
            per_device_train_batch_size = 2,
            gradient_accumulation_steps = 4,
            warmup_steps = 5,
            max_steps = 70,
            learning_rate = 1e-4,
            fp16 = not is_bfloat16_supported(),
            bf16 = is_bfloat16_supported(),
            logging_steps = 1,
            optim = "adamw_8bit",
            weight_decay = 0.01,
            lr_scheduler_type = "cosine",
            seed = 3407,
            output_dir = "outputs",
            save_strategy = "no",
        ),
    )
    trainer.train()
    
    print("Merging LoRA adapter into base model...")
    import os, subprocess
    hf_dir = "/root/step-zero-nemotron-hf"
    os.makedirs(hf_dir, exist_ok=True)
    
    # Merge the LoRA weights into the base model using PEFT directly
    merged_model = model.merge_and_unload()
    merged_model.save_pretrained(hf_dir)
    tokenizer.save_pretrained(hf_dir)
    
    print(f"HF export contents: {os.listdir(hf_dir)}")
    
    print("Cloning llama.cpp for manual GGUF conversion...")
    subprocess.run(["git", "clone", "--depth", "1", "https://github.com/ggerganov/llama.cpp.git", "/root/llama.cpp"], check=True)
    subprocess.run(["pip", "install", "-r", "/root/llama.cpp/requirements.txt"], check=True)
    
    print("Converting HF model to GGUF...")
    result = subprocess.run([
        "python3", "/root/llama.cpp/convert_hf_to_gguf.py", 
        hf_dir, 
        "--outfile", "/vol/step-zero-nemotron-finetuned.gguf", 
        "--outtype", "q8_0"
    ], capture_output=True, text=True)
    print("STDOUT:", result.stdout[-2000:] if result.stdout else "")
    print("STDERR:", result.stderr[-2000:] if result.stderr else "")
    if result.returncode != 0:
        print(f"convert_hf_to_gguf.py failed with code {result.returncode}")
        print("Trying f16 fallback...")
        result2 = subprocess.run([
            "python3", "/root/llama.cpp/convert_hf_to_gguf.py", 
            hf_dir, 
            "--outfile", "/vol/step-zero-nemotron-finetuned.gguf", 
            "--outtype", "f16"
        ], capture_output=True, text=True)
        print("STDOUT:", result2.stdout[-2000:] if result2.stdout else "")
        print("STDERR:", result2.stderr[-2000:] if result2.stderr else "")
    
    try:
        vol.commit()
    except Exception as e:
        print("vol.commit() skipped:", e)
    print("Training complete. GGUF artifact generated in /vol.")

@app.local_entrypoint()
def main():
    print("Submitting training job to Modal infrastructure...")
    train_model.remote()
    print("Done! To download your weights, run:")
    print("modal volume get step-zero-volume step-zero-nemotron.gguf models/")