from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

import trackio
from lora_data import encode_examples, split_examples
from peft import LoraConfig, TaskType, get_peft_model
from snip_common import ARTIFACT_DIR, parameter_count
from transformers import (
    AutoModelForCausalLM,
    PreTrainedTokenizerFast,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)

ADAPTER_DIR = ARTIFACT_DIR.parent / "snip-0.4m-story-lora"
MERGED_DIR = ARTIFACT_DIR.parent / "snip-0.4m-story-merged"


class AdapterDiagnosticCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        loss = (logs or {}).get("loss")
        if loss is not None and loss != loss:
            trackio.alert(
                title="LoRA produced NaN loss",
                text=f"NaN detected at step {state.global_step}.",
                level=trackio.AlertLevel.ERROR,
            )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-steps", type=int, default=300)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--learning-rate", type=float, default=0.002)
    parser.add_argument("--resume-from-checkpoint")
    args = parser.parse_args()

    set_seed(43)
    os.environ.setdefault("TRACKIO_PROJECT", "snip-model-foundry")
    tokenizer = PreTrainedTokenizerFast.from_pretrained(ARTIFACT_DIR)
    base_model = AutoModelForCausalLM.from_pretrained(ARTIFACT_DIR)
    base_parameters = parameter_count(base_model)
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=4,
        lora_alpha=8,
        lora_dropout=0.05,
        bias="none",
        target_modules=["c_attn", "c_proj", "c_fc"],
        fan_in_fan_out=True,
    )
    model = get_peft_model(base_model, lora_config)
    trainable_parameters = sum(
        parameter.numel() for parameter in model.parameters() if parameter.requires_grad
    )
    train_examples, eval_examples = split_examples()
    train_dataset = encode_examples(train_examples, tokenizer)
    eval_dataset = encode_examples(eval_examples, tokenizer)

    started = time.perf_counter()
    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir=str(ADAPTER_DIR / "checkpoints"),
            max_steps=args.max_steps,
            per_device_train_batch_size=args.batch_size,
            per_device_eval_batch_size=args.batch_size,
            learning_rate=args.learning_rate,
            warmup_steps=max(1, int(args.max_steps * 0.05)),
            weight_decay=0.0,
            lr_scheduler_type="cosine",
            eval_strategy="steps",
            eval_steps=50,
            logging_steps=10,
            save_strategy="steps",
            save_steps=100,
            save_total_limit=2,
            report_to="trackio",
            project="snip-model-foundry",
            run_name="snip-0.4m-story-lora-r4-v1",
            use_cpu=True,
            remove_unused_columns=False,
        ),
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        callbacks=[AdapterDiagnosticCallback()],
    )
    trainer.train(
        resume_from_checkpoint=args.resume_from_checkpoint or None,
    )
    elapsed = time.perf_counter() - started
    model.save_pretrained(ADAPTER_DIR)
    tokenizer.save_pretrained(ADAPTER_DIR)

    evaluations = [entry for entry in trainer.state.log_history if "eval_loss" in entry]
    losses = [
        float(entry["loss"]) for entry in trainer.state.log_history if "loss" in entry
    ]
    summary = {
        "model": "SNIP-0.4M Story LoRA",
        "base_parameters": base_parameters,
        "trainable_parameters": trainable_parameters,
        "trainable_percent": 100 * trainable_parameters / base_parameters,
        "rank": 4,
        "train_examples": len(train_dataset),
        "eval_examples": len(eval_dataset),
        "steps": args.max_steps,
        "mean_logged_loss": sum(losses) / len(losses),
        "eval_loss": float(evaluations[-1]["eval_loss"]),
        "elapsed_seconds": elapsed,
        "resumed_from": args.resume_from_checkpoint,
    }
    (ADAPTER_DIR / "training_summary.json").write_text(
        json.dumps(summary, indent=2),
        encoding="utf-8",
    )

    merged = model.merge_and_unload()
    merged.save_pretrained(MERGED_DIR, safe_serialization=True)
    tokenizer.save_pretrained(MERGED_DIR)
    (MERGED_DIR / "variant.json").write_text(
        json.dumps(
            {
                "base": str(Path(ARTIFACT_DIR).name),
                "adapter": str(ADAPTER_DIR.name),
                "variant": "merged-story-lora-r4",
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()