"""BoccaccioAI - Main training entrypoint.

Invoked as ``python -m src.training.train`` from the project root.
Supports both pre-training (on binary memmap data) and instruction
fine-tuning (on JSONL data).

De Lauretis Tech
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

import torch
import yaml
import lightning as L
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
from lightning.pytorch.loggers import WandbLogger, TensorBoardLogger
from torch.utils.data import DataLoader

from src.model.config import BoccaccioConfig
from src.training.lightning_module import BoccaccioLightningModule
from src.training.callbacks import TrainingStatsCallback, SaveConfigCallback
from src.data.dataset import PreTokenizedDataset, InstructionDataset

log = logging.getLogger(__name__)


# ------------------------------------------------------------------
# CLI
# ------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="BoccaccioAI training entrypoint")

    parser.add_argument(
        "--mode",
        type=str,
        required=True,
        choices=["pretrain", "finetune"],
        help="Training mode: 'pretrain' or 'finetune'.",
    )
    parser.add_argument(
        "--model-config",
        type=str,
        default="configs/model.yaml",
        help="Path to model architecture YAML.",
    )
    parser.add_argument(
        "--model-variant",
        type=str,
        default="model",
        choices=["model", "nano"],
        help="Model variant to load from the config.",
    )
    parser.add_argument(
        "--training-config",
        type=str,
        default="configs/training.yaml",
        help="Path to training hyper-parameters YAML.",
    )
    parser.add_argument(
        "--data-path",
        type=str,
        required=True,
        help="Path to the tokenized data directory.",
    )
    parser.add_argument(
        "--tokenizer-path",
        type=str,
        default="tokenizer/boccaccio-32k.json",
        help="Path to the tokenizer JSON file.",
    )
    parser.add_argument(
        "--resume-from",
        type=str,
        default=None,
        help="Path to a Lightning checkpoint to load weights from (resets step/optimizer).",
    )
    parser.add_argument(
        "--resume-ckpt",
        type=str,
        default=None,
        help="Path to a Lightning checkpoint for full resume (step, optimizer, LR, etc.).",
    )
    parser.add_argument(
        "--wandb-project",
        type=str,
        default="boccaccio-ai",
        help="Weights & Biases project name.",
    )
    parser.add_argument(
        "--config-section",
        type=str,
        default=None,
        help="Override config section name (e.g. 'pretrain_phase2'). Defaults to --mode value.",
    )
    parser.add_argument(
        "--wandb-offline",
        action="store_true",
        help="Run Weights & Biases in offline mode.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility.",
    )

    return parser.parse_args()


# ------------------------------------------------------------------
# Dataset / DataLoader helpers
# ------------------------------------------------------------------

def _build_dataloaders(
    mode: str,
    data_path: str,
    tokenizer_path: str,
    train_config: dict,
) -> tuple[DataLoader, DataLoader | None]:
    """Create train and (optional) validation DataLoaders."""

    sequence_length: int = train_config["sequence_length"]
    batch_size: int = train_config["micro_batch_size"]
    num_workers: int = train_config.get("num_workers", 4)

    if mode == "pretrain":
        train_file = str(Path(data_path) / "train.bin")
        val_file = str(Path(data_path) / "val.bin")

        train_dataset = PreTokenizedDataset(train_file, sequence_length)
        val_dataset = PreTokenizedDataset(val_file, sequence_length)

        log.info(
            "Pre-training datasets -- train: %d sequences, val: %d sequences",
            len(train_dataset),
            len(val_dataset),
        )
    else:
        # Fine-tuning: usa InstructionDataset con formato JSONL conversations.
        # Loss mascherata (-100) su user/system, solo assistant + </s> contribuiscono.
        train_file = Path(data_path) / "train.jsonl"
        val_file = Path(data_path) / "val.jsonl"

        if not train_file.exists():
            raise FileNotFoundError(
                f"Fine-tuning requires train.jsonl in {data_path}. "
                f"Run 'python -m src.data.prepare_finetune' first."
            )

        train_dataset = InstructionDataset(str(train_file), tokenizer_path, sequence_length)

        if val_file.exists():
            val_dataset = InstructionDataset(str(val_file), tokenizer_path, sequence_length)
            log.info(
                "Fine-tuning datasets (JSONL) -- train: %d examples, val: %d examples",
                len(train_dataset),
                len(val_dataset),
            )
        else:
            val_dataset = None
            log.info(
                "Fine-tuning dataset (JSONL) -- train: %d examples (no validation set)",
                len(train_dataset),
            )

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
    )

    val_loader = None
    if val_dataset is not None:
        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            pin_memory=True,
        )

    return train_loader, val_loader


# ------------------------------------------------------------------
# Callbacks
# ------------------------------------------------------------------

def _build_callbacks(
    train_config: dict,
    model_config: BoccaccioConfig,
    num_params: int,
) -> list[L.Callback]:
    """Assemble the list of Lightning callbacks."""

    checkpoint_dir = train_config.get("checkpoint_dir", "checkpoints")
    checkpoint_every = train_config.get("checkpoint_every_n_steps", 1000)
    save_top_k = train_config.get("save_top_k", 3)

    callbacks: list[L.Callback] = [
        ModelCheckpoint(
            dirpath=checkpoint_dir,
            every_n_train_steps=checkpoint_every,
            save_top_k=-1,
            save_last=True,
        ),
        LearningRateMonitor(logging_interval="step"),
        TrainingStatsCallback(
            num_params=num_params,
            sequence_length=train_config["sequence_length"],
            gradient_accumulation_steps=train_config.get("gradient_accumulation_steps", 1),
        ),
        SaveConfigCallback(model_config, train_config),
    ]

    return callbacks


# ------------------------------------------------------------------
# Logger
# ------------------------------------------------------------------

def _build_logger(
    wandb_project: str,
    wandb_offline: bool,
    mode: str,
) -> WandbLogger | TensorBoardLogger:
    """Try to create a WandB logger; fall back to TensorBoard on failure."""

    try:
        logger = WandbLogger(
            project=wandb_project,
            name=mode,
            offline=wandb_offline,
        )
        log.info("Using Weights & Biases logger (project=%s)", wandb_project)
        return logger
    except Exception:
        log.warning(
            "Failed to initialise WandB logger; falling back to TensorBoard."
        )
        return TensorBoardLogger(save_dir="logs", name=mode)


# ------------------------------------------------------------------
# Main
# ------------------------------------------------------------------

def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )

    args = parse_args()

    if args.resume_from and args.resume_ckpt:
        raise ValueError(
            "Cannot use both --resume-from and --resume-ckpt. "
            "Use --resume-from to load weights only (resets step), "
            "or --resume-ckpt for full resume (step, optimizer, LR)."
        )

    # Seed everything for reproducibility.
    L.seed_everything(args.seed)

    # ── Load configs ──────────────────────────────────────────────
    model_config = BoccaccioConfig.from_yaml(args.model_config, args.model_variant)
    log.info("Model config loaded: %s (variant=%s)", args.model_config, args.model_variant)

    with open(args.training_config, "r", encoding="utf-8") as f:
        full_training_config = yaml.safe_load(f)
    config_section = args.config_section or args.mode
    train_config: dict = full_training_config[config_section]
    log.info("Training config loaded: %s [%s]", args.training_config, config_section)

    # ── Datasets & DataLoaders ────────────────────────────────────
    train_loader, val_loader = _build_dataloaders(
        mode=args.mode,
        data_path=args.data_path,
        tokenizer_path=args.tokenizer_path,
        train_config=train_config,
    )

    # Inject num_examples into train_config for the scheduler calculation.
    # This is needed for fine-tuning where total_steps depends on the number
    # of examples (not num_tokens, since examples are padded to sequence_length).
    if args.mode == "finetune":
        train_config["num_examples"] = len(train_loader.dataset)
        log.info("Fine-tuning: %d examples in training set", train_config["num_examples"])

    # ── Lightning module ──────────────────────────────────────────
    model = BoccaccioLightningModule(model_config, train_config)
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    # ── Load pre-trained weights (for fine-tuning) ───────────────
    # Lightning's ckpt_path mechanism doesn't handle torch.compile prefix
    # remapping (model._orig_mod. -> model.), so we load weights manually.
    ckpt_path = None
    if args.resume_from:
        log.info("Loading pre-trained weights from %s", args.resume_from)
        ckpt = torch.load(args.resume_from, map_location="cpu", weights_only=False)
        state_dict = ckpt.get("state_dict", ckpt)

        # Remap torch.compile prefix: model._orig_mod.X -> model.model.X
        # The checkpoint was saved with torch.compile which wraps the model
        # as model._orig_mod.<inner>. The LightningModule has structure:
        #   self.model = BoccaccioForCausalLM  (model.<X>)
        #   BoccaccioForCausalLM.model = Transformer  (model.model.<X>)
        # So checkpoint key "model._orig_mod.model.embed_tokens" must become
        # "model.model.embed_tokens".
        cleaned = {}
        for key, value in state_dict.items():
            new_key = key
            if new_key.startswith("model._orig_mod."):
                # Strip "model._orig_mod." and prepend "model."
                new_key = "model." + new_key[len("model._orig_mod."):]
            elif new_key.startswith("_orig_mod."):
                new_key = new_key[len("_orig_mod."):]
            cleaned[new_key] = value

        missing, unexpected = model.load_state_dict(cleaned, strict=False)
        if missing:
            log.warning("Missing keys (%d): %s", len(missing), missing[:5])
        if unexpected:
            log.warning("Unexpected keys (%d): %s", len(unexpected), unexpected[:5])
        log.info("Pre-trained weights loaded successfully.")

    # ── Callbacks ─────────────────────────────────────────────────
    callbacks = _build_callbacks(train_config, model_config, num_params)

    # ── Logger ────────────────────────────────────────────────────
    logger = _build_logger(args.wandb_project, args.wandb_offline, args.mode)

    # ── Compute max_steps (pretrain) or max_epochs (finetune) ─────
    trainer_kwargs: dict = {}
    if args.mode == "pretrain":
        effective_batch_tokens = (
            train_config["micro_batch_size"]
            * train_config["gradient_accumulation_steps"]
            * train_config["sequence_length"]
        )
        max_steps = int(train_config["num_tokens"]) // effective_batch_tokens
        trainer_kwargs["max_steps"] = max_steps
        log.info(
            "Pre-training for %d steps (~%d tokens, effective batch = %d tokens/step)",
            max_steps,
            int(train_config["num_tokens"]),
            effective_batch_tokens,
        )
    else:
        max_epochs = train_config.get("num_epochs", 3)
        trainer_kwargs["max_epochs"] = max_epochs
        log.info("Fine-tuning for %d epoch(s)", max_epochs)

    # ── Trainer ───────────────────────────────────────────────────
    trainer = L.Trainer(
        accelerator=train_config.get("accelerator", "auto"),
        devices=train_config.get("devices", "auto"),
        strategy=train_config.get("strategy", "auto"),
        precision=train_config.get("precision", "bf16-mixed"),
        accumulate_grad_batches=train_config.get("gradient_accumulation_steps", 1),
        gradient_clip_val=train_config.get("gradient_clip_val", 1.0),
        val_check_interval=train_config.get("val_check_interval", None),
        log_every_n_steps=train_config.get("log_every_n_steps", 50),
        callbacks=callbacks,
        logger=logger,
        enable_checkpointing=True,
        **trainer_kwargs,
    )

    # ── Train ─────────────────────────────────────────────────────
    # --resume-ckpt: full resume (step, optimizer, LR schedule, etc.)
    # --resume-from: load weights only (step resets to 0)
    ckpt_path = args.resume_ckpt if args.resume_ckpt else None
    trainer.fit(model, train_loader, val_loader, ckpt_path=ckpt_path)

    # ── Save final model weights ──────────────────────────────────
    checkpoint_dir = Path(train_config.get("checkpoint_dir", "checkpoints"))
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    final_path = checkpoint_dir / "model.pt"
    state_dict = model.model.state_dict()
    torch.save(state_dict, final_path)
    log.info("Final model weights saved to %s", final_path)

    log.info("Training complete.")


if __name__ == "__main__":
    main()