from __future__ import annotations

import argparse
import json
import os
import time

import trackio
from snip_common import (
    ARTIFACT_DIR,
    DATA_DIR,
    build_tokenizer,
    make_model,
    parameter_count,
    read_texts,
    texts_to_blocks,
)
from transformers import (
    DataCollatorForLanguageModeling,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)


class DiagnosticCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if not logs:
            return
        loss = logs.get("loss")
        if loss is not None and loss != loss:
            trackio.alert(
                title="NaN loss",
                text=f"Training produced NaN at step {state.global_step}.",
                level=trackio.AlertLevel.ERROR,
            )
        if loss is not None and state.global_step >= 200 and loss > 7:
            trackio.alert(
                title="High loss",
                text=f"Loss is {loss:.4f} at step {state.global_step}.",
                level=trackio.AlertLevel.WARN,
            )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-steps", type=int, default=800)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--learning-rate", type=float, default=8e-4)
    parser.add_argument("--resume-from-checkpoint")
    args = parser.parse_args()

    set_seed(42)
    os.environ.setdefault("TRACKIO_PROJECT", "snip-model-foundry")

    train_texts = read_texts(DATA_DIR / "train.jsonl")
    eval_texts = read_texts(DATA_DIR / "eval.jsonl")
    tokenizer = build_tokenizer(train_texts)
    train_dataset = texts_to_blocks(train_texts, tokenizer)
    eval_dataset = texts_to_blocks(eval_texts, tokenizer)
    model = make_model(tokenizer)
    parameters = parameter_count(model)

    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    started = time.perf_counter()
    training_args = TrainingArguments(
        output_dir=str(ARTIFACT_DIR / "checkpoints"),
        max_steps=args.max_steps,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=1,
        learning_rate=args.learning_rate,
        warmup_steps=max(1, int(args.max_steps * 0.05)),
        weight_decay=0.01,
        lr_scheduler_type="cosine",
        eval_strategy="steps",
        eval_steps=100,
        logging_steps=20,
        save_strategy="steps",
        save_steps=200,
        save_total_limit=2,
        report_to="trackio",
        project="snip-model-foundry",
        run_name="snip-0.4m-pretrain-v1",
        use_cpu=True,
        dataloader_num_workers=0,
        remove_unused_columns=False,
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
        processing_class=tokenizer,
        callbacks=[DiagnosticCallback()],
    )
    result = trainer.train(
        resume_from_checkpoint=args.resume_from_checkpoint or None,
    )
    elapsed = time.perf_counter() - started
    trainer.save_model(ARTIFACT_DIR)
    tokenizer.save_pretrained(ARTIFACT_DIR)
    logged_losses = [
        float(entry["loss"]) for entry in trainer.state.log_history if "loss" in entry
    ]
    evaluations = [entry for entry in trainer.state.log_history if "eval_loss" in entry]
    if not evaluations:
        raise RuntimeError("Training completed without a recorded evaluation.")
    final_evaluation = evaluations[-1]

    summary = {
        "model": "SNIP-0.4M",
        "parameters": parameters,
        "train_examples": len(train_dataset),
        "eval_examples": len(eval_dataset),
        "max_steps": args.max_steps,
        "train_loss": sum(logged_losses) / len(logged_losses),
        "trainer_reported_loss": float(result.training_loss),
        "eval_loss": float(final_evaluation["eval_loss"]),
        "continuation_elapsed_seconds": elapsed,
        "resumed_from": args.resume_from_checkpoint,
        "tokens_seen": args.max_steps * args.batch_size * 128,
    }
    (ARTIFACT_DIR / "training_summary.json").write_text(
        json.dumps(summary, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()