# -*- coding: utf-8 -*-
"""
QLoRA fine-tune Qwen3.5-9B base on Solidity security data via Unsloth.
Targets: severity calibration + tool calling (Qwythos-9B base).
Trains on A100 via HF Jobs.
"""

import sys, os, subprocess

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
sys.stderr.reconfigure(encoding="utf-8", errors="replace")
os.environ.setdefault("PYTHONIOENCODING", "utf-8")

WORKDIR = "/workspace"
os.makedirs(WORKDIR, exist_ok=True)


def log(msg):
    print(msg)


def run_cmd(cmd, timeout=None):
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    return r


def main():
    log("=" * 50)
    log("QLoRA Fine-Tune: Qwen3.5-9B (Unsloth)")
    log("Dataset: qwythos-sec-training-data")
    log("Target: severity calibration + tool calling")
    log("=" * 50)

    # ================================================================
    # Step 1/4: Install Unsloth
    # ================================================================
    log("\n[1/4] Installing Unsloth ...")
    r1 = run_cmd(["pip", "install", "--quiet", "--no-cache-dir", "unsloth"], timeout=600)
    log(f"  pip install unsloth exit={r1.returncode}")
    if r1.returncode != 0:
        log(f"  stderr (last 1000): {r1.stderr[-1000:]}")
        sys.exit(1)
    log("  [OK] Unsloth installed")

    # hf_transfer for faster model upload
    r2 = run_cmd(["pip", "install", "--quiet", "--no-cache-dir", "hf_transfer"], timeout=60)
    log(f"  hf_transfer exit={r2.returncode}")

    # ================================================================
    # Step 2/4: Load dataset
    # ================================================================
    log("\n[2/4] Loading security training dataset ...")
    from datasets import load_dataset, Dataset

    ds = load_dataset("mxguru1/qwythos-sec-training-data", split="train")
    val_ds = load_dataset("mxguru1/qwythos-sec-training-data", split="validation")
    log(f"  train: {len(ds)} rows, val: {len(val_ds)} rows")

    # Format as chat templates for Unsloth SFT
    def format_prompt(row):
        text = (
            "<|im_start|>user\n" + row["prompt"] + "<|im_end|>\n"
            "<|im_start|>assistant\n" + row["completion"] + "<|im_end|>"
        )
        return {"text": text}

    train_ds = ds.map(format_prompt, remove_columns=ds.column_names)
    val_ds_out = val_ds.map(format_prompt, remove_columns=val_ds.column_names)
    log(f"  formatted {len(train_ds)} train / {len(val_ds_out)} val samples")

    # ================================================================
    # Step 3/4: Train with Unsloth
    # ================================================================
    log("\n[3/4] Loading Qwen3.5-9B + tokenizer (Unsloth 4-bit) ...")
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name="Qwen/Qwen3.5-9B",
        max_seq_length=2048,
        load_in_4bit=True,
        load_in_8bit=False,
        fast_inference=False,
        token=os.environ.get("HF_TOKEN", ""),
    )
    log("  model loaded (4-bit QLoRA)")

    # Add LoRA adapters - all linear modules for full coverage
    model = FastLanguageModel.get_peft_model(
        model,
        r=32,
        lora_alpha=64,
        lora_dropout=0.05,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
            "embed_tokens", "lm_head",
        ],
        bias="none",
        use_gradient_checkpointing="unsloth",
    )
    log("  LoRA adapters attached (r=32, all linear modules)")

    log("  Starting training ...")
    from unsloth import is_bf16_supported
    from trl import SFTTrainer
    from transformers import TrainingArguments, DataCollatorForSeq2Seq

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_ds,
        eval_dataset=val_ds_out,
        dataset_text_field="text",
        max_seq_length=2048,
        data_collator=DataCollatorForSeq2Seq(tokenizer, model=model, padding=True),
        args=TrainingArguments(
            output_dir="/workspace/checkpoints",
            per_device_train_batch_size=2,
            gradient_accumulation_steps=8,
            num_train_epochs=3,
            warmup_steps=10,
            learning_rate=2e-4,
            weight_decay=0.0,
            lr_scheduler_type="cosine",
            optim="adamw_8bit",
            bf16=is_bf16_supported(),
            fp16=not is_bf16_supported(),
            logging_steps=5,
            save_steps=50,
            eval_steps=50,
            save_total_limit=3,
            report_to="none",
        ),
    )
    log("  trainer initialized - calling train() ...")
    trainer.train()
    log("  [OK] training complete")

    # ================================================================
    # Step 4/4: Push adapter to HF
    # ================================================================
    log("\n[4/4] Saving and pushing adapters to HuggingFace ...")
    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

    adapter_dir = "/workspace/qwythos-9b-security-adapter"
    model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    log(f"  adapters saved to {adapter_dir}")

    from huggingface_hub import HfApi, create_repo

    org_repo = "mxguru1/qwythos-9b-security-unsloth"
    try:
        create_repo(org_repo, repo_type="model", private=True, exist_ok=True)
        log(f"  repo ready: {org_repo}")
    except Exception as e:
        log(f"  [WARN] create_repo: {e}")

    api = HfApi(token=os.environ.get("HF_TOKEN", ""))
    try:
        api.upload_folder(
            folder_path=adapter_dir,
            repo_id=org_repo,
            repo_type="model",
        )
        log("  [OK] adapter pushed to HF")
    except Exception as e:
        log(f"  [FAIL] upload: {e}")
        sys.exit(1)

    log("")
    log("=" * 50)
    log("COMPLETE")
    log(f"Adapter: https://huggingface.co/{org_repo}")
    log("=" * 50)


if __name__ == "__main__":
    main()