#!/usr/bin/env python3
"""Fine-tune a local LLM with LoRA for Wall-E personalization.

Uses unsloth for efficient LoRA training and exports to GGUF format
for direct use with Ollama.

Progress is reported as JSON lines on stdout so the Node.js runner
can parse them in real time.
"""

import argparse
import json
import os
import sys

# ---------------------------------------------------------------------------
# Model registry — mirrors SUPPORTED_BASES in trainer.js
# ---------------------------------------------------------------------------

MODEL_REGISTRY = {
    "llama3.1-8b": "unsloth/Meta-Llama-3.1-8B-Instruct",
    "phi-4": "unsloth/Phi-4",
    "qwen2.5-7b": "unsloth/Qwen2.5-7B-Instruct",
    "gemma4-4b": "unsloth/gemma-4-E4B-it",
}

# LoRA target modules per architecture
LORA_TARGETS = {
    "llama3.1-8b": [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    "phi-4": [
        "q_proj", "k_proj", "v_proj", "dense",
        "fc1", "fc2",
    ],
    "qwen2.5-7b": [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    "gemma4-4b": [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
}


def report(stage, pct, msg=""):
    """Print a JSON progress line to stdout."""
    print(json.dumps({"stage": stage, "progress": pct, "message": msg}), flush=True)


CHAT_TEMPLATE = """<|im_start|>system
{system}<|im_end|>
<|im_start|>user
{user}<|im_end|>
<|im_start|>assistant
{assistant}<|im_end|>"""


def load_data(data_path, fmt="alpaca"):
    """Load JSONL training data and convert to Alpaca-style dicts.

    Args:
        data_path: Path to JSONL file.
        fmt: 'alpaca' or 'chat-ml'. If 'alpaca', auto-detects chat-ml
             when the first record has a 'messages' key.
    """
    samples = []
    skipped = 0
    raw_lines = []
    with open(data_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                raw_lines.append(json.loads(line))
            except json.JSONDecodeError:
                skipped += 1

    # Auto-detect format from first record
    effective_fmt = fmt
    if raw_lines and fmt == "alpaca":
        first = raw_lines[0]
        if "messages" in first and "instruction" not in first:
            effective_fmt = "chat-ml"

    for obj in raw_lines:
        if effective_fmt == "chat-ml" and "messages" in obj:
            msgs = obj["messages"]
            system_text = ""
            user_text = ""
            assistant_text = ""
            for m in msgs:
                role = m.get("role", "")
                content = m.get("content", "")
                if role == "system":
                    system_text = content
                elif role == "user":
                    user_text = content
                elif role == "assistant":
                    assistant_text = content
            samples.append({
                "instruction": CHAT_TEMPLATE.format(
                    system=system_text, user=user_text, assistant=assistant_text
                ),
                "input": "",
                "output": assistant_text,
            })
        elif "instruction" in obj:
            samples.append({
                "instruction": obj["instruction"],
                "input": obj.get("input", ""),
                "output": obj.get("output", ""),
            })
        elif "messages" in obj:
            # Fallback: convert chat messages to instruction/output
            msgs = obj["messages"]
            instruction_parts = []
            output_text = ""
            for m in msgs:
                role = m.get("role", "")
                content = m.get("content", "")
                if role in ("system", "user"):
                    instruction_parts.append(content)
                elif role == "assistant":
                    output_text = content
            samples.append({
                "instruction": "\n".join(instruction_parts),
                "input": "",
                "output": output_text,
            })
        else:
            # Skip unrecognised formats
            continue

    if skipped:
        report("init", 5, f"Skipped {skipped} malformed JSONL line(s)")
    return samples


def main():
    parser = argparse.ArgumentParser(
        description="Fine-tune a local LLM with LoRA for Wall-E."
    )
    parser.add_argument("--base", required=True, help="Base model key")
    parser.add_argument("--data", required=True, help="Path to JSONL training data")
    parser.add_argument("--output", required=True, help="Output directory")
    parser.add_argument("--name", default="walle-model", help="Output model name")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--lora-rank", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--max-seq-length", type=int, default=2048)
    parser.add_argument("--format", choices=["alpaca", "chat-ml"], default="alpaca",
                        help="Input data format: 'alpaca' or 'chat-ml'")
    args = parser.parse_args()

    # --- Validate base model --------------------------------------------
    if args.base not in MODEL_REGISTRY:
        report("error", 0, f"Unknown base model: {args.base}")
        sys.exit(1)

    hf_model = MODEL_REGISTRY[args.base]

    # --- Load data ------------------------------------------------------
    report("init", 0, "Loading training data...")

    if not os.path.isfile(args.data):
        report("error", 0, f"Data file not found: {args.data}")
        sys.exit(1)

    samples = load_data(args.data, fmt=args.format)
    if not samples:
        report("error", 0, "No valid training samples found in data file")
        sys.exit(1)

    report("init", 10, f"Loaded {len(samples)} training samples")

    # --- Detect device capabilities ------------------------------------
    import torch
    use_fp16 = torch.cuda.is_available()

    # --- Import unsloth -------------------------------------------------
    report("init", 15, "Loading unsloth...")

    try:
        from unsloth import FastLanguageModel
        from trl import SFTTrainer
        from transformers import TrainingArguments
        from datasets import Dataset
    except ImportError as e:
        missing = str(e)
        report(
            "error", 0,
            f"Required package not installed ({missing}). "
            "Install with: pip install 'unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git' "
            "trl datasets"
        )
        sys.exit(1)

    # --- Load model + apply LoRA ----------------------------------------
    report("loading", 20, f"Loading base model: {hf_model}")

    lora_targets = LORA_TARGETS.get(args.base, [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ])

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=hf_model,
        max_seq_length=args.max_seq_length,
        dtype=None,   # auto-detect
        load_in_4bit=True,
    )

    report("loading", 40, "Applying LoRA adapter...")

    model = FastLanguageModel.get_peft_model(
        model,
        r=args.lora_rank,
        target_modules=lora_targets,
        lora_alpha=args.lora_alpha,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
    )

    # --- Prepare dataset ------------------------------------------------
    report("preparing", 50, "Formatting dataset...")

    alpaca_template = (
        "Below is an instruction that describes a task, paired with an input "
        "that provides further context. Write a response that appropriately "
        "completes the request.\n\n"
        "### Instruction:\n{instruction}\n\n"
        "### Input:\n{input}\n\n"
        "### Response:\n{output}"
    )

    def format_sample(sample):
        return alpaca_template.format(**sample) + tokenizer.eos_token

    texts = [format_sample(s) for s in samples]
    dataset = Dataset.from_dict({"text": texts})

    report("preparing", 55, f"Dataset ready: {len(dataset)} samples")

    # --- Train ----------------------------------------------------------
    report("training", 60, "Starting training...")

    output_ckpt = os.path.join(args.output, f"{args.name}-checkpoints")
    os.makedirs(output_ckpt, exist_ok=True)

    training_args = TrainingArguments(
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        num_train_epochs=args.epochs,
        learning_rate=args.lr,
        fp16=use_fp16,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=42,
        output_dir=output_ckpt,
        report_to="none",
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        dataset_text_field="text",
        max_seq_length=args.max_seq_length,
        dataset_num_proc=2,
        packing=False,
        args=training_args,
    )

    train_result = trainer.train()
    final_loss = train_result.training_loss

    report("training", 85, f"Training complete. Final loss: {final_loss:.4f}")

    # --- Save to GGUF ---------------------------------------------------
    report("saving", 90, "Exporting model to GGUF format...")

    model_output = os.path.join(args.output, args.name)
    os.makedirs(model_output, exist_ok=True)

    model.save_pretrained_gguf(
        model_output,
        tokenizer,
        quantization_method="q4_k_m",
    )

    report("saving", 95, "GGUF export complete")

    # --- Final result ---------------------------------------------------
    report("complete", 100, f"Model saved to {model_output}")

    # Print final result object (the runner parses this)
    print(
        json.dumps({
            "success": True,
            "model_path": model_output,
            "samples_trained": len(samples),
            "epochs": args.epochs,
            "final_loss": round(final_loss, 6),
        }),
        flush=True,
    )


if __name__ == "__main__":
    main()
