from pathlib import Path
import re

import torch
from torch.distributed.checkpoint.state_dict import (StateDictOptions,
                                                     get_model_state_dict,
                                                     get_optimizer_state_dict,
                                                     set_model_state_dict,
                                                     set_optimizer_state_dict)
from torch.distributed.fsdp import FSDPModule
from torchdata.stateful_dataloader import StatefulDataLoader
from src.utils.logging import get_logger

_SAVE_OPTIONS = StateDictOptions(full_state_dict=True, cpu_offload=True)
_LOAD_OPTIONS = StateDictOptions(
    full_state_dict=True, broadcast_from_rank0=True)


def resolve_checkpoint_path(path: str | Path) -> Path:
    path = Path(path)
    if path.is_file():
        return path
    if not path.is_dir():
        raise FileNotFoundError(f"Checkpoint does not exist: {path}")
    match = re.fullmatch(r"global_step(\d+)", path.name)
    if match:
        checkpoint = path / f"step_{int(match.group(1))}.pth"
        if checkpoint.is_file():
            return checkpoint
    candidates = sorted(path.glob("*.pth"))
    if not candidates:
        raise FileNotFoundError(f"No .pth checkpoint found in {path}.")
    if len(candidates) != 1:
        raise ValueError(f"Multiple checkpoints found in {path}; specify the checkpoint file.")
    return candidates[0]


def load_dataloader_state(dataloader, checkpoint_path: Path, step: int) -> None:
    rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
    path = checkpoint_path.parent / "dataloader" / f"dataloader_step{step}_rank{rank}.pth"
    if not path.is_file():
        raise FileNotFoundError(f"Requested dataloader resume state is missing: {path}")
    state = torch.load(path, map_location="cpu", weights_only=False)
    dataloader.load_state_dict(state["dataloader"])


def save_checkpoint(
    model: torch.nn.Module | FSDPModule,
    save_dir: str | Path,
    step: int,
    global_rank: int,
    optimizer: torch.optim.Optimizer = None,
    scheduler: torch.optim.lr_scheduler._LRScheduler = None,
    dataloader: StatefulDataLoader = None
) -> None:
    logger = get_logger()

    save_path = Path(save_dir) / f"global_step{step}"
    save_path.mkdir(parents=True, exist_ok=True)

    state_dict = {
        "step": step,
        "model": get_model_state_dict(model, options=_SAVE_OPTIONS)
    }

    if optimizer is not None:
        state_dict["optimizer"] = get_optimizer_state_dict(
            model, optimizer, options=_SAVE_OPTIONS)

    if scheduler is not None:
        state_dict["scheduler"] = scheduler.state_dict()

    if global_rank <= 0:
        ckpt_path = save_path / f"step_{step}.pth"
        logger.info(f"Saving checkpoint at step {step} to {ckpt_path}")
        torch.save(state_dict, ckpt_path)

    if dataloader is not None:
        save_path = save_path / 'dataloader'
        save_path.mkdir(parents=True, exist_ok=True)
        dataloader_path = save_path / \
            f"dataloader_step{step}_rank{global_rank}.pth"
        state_dict = {"dataloader": dataloader.state_dict()}
        torch.save(state_dict, dataloader_path)


def load_checkpoint(
    model: torch.nn.Module | FSDPModule,
    path: str | Path,
    optimizer: torch.optim.Optimizer = None,
    scheduler: torch.optim.lr_scheduler._LRScheduler = None,
    dataloader: torch.utils.data.DataLoader = None,
) -> int:
    logger = get_logger()

    checkpoint_path = resolve_checkpoint_path(path)
    logger.info(f"Loading checkpoint from {checkpoint_path}")
    state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    for key, component in (("optimizer", optimizer), ("scheduler", scheduler)):
        if component is not None and key not in state_dict:
            raise KeyError(f"Requested resume state {key!r} is missing from {checkpoint_path}.")

    set_model_state_dict(model, state_dict["model"], options=_LOAD_OPTIONS)

    if optimizer is not None:
        set_optimizer_state_dict(
            model, optimizer, state_dict["optimizer"], options=_LOAD_OPTIONS)

    if scheduler is not None:
        scheduler.load_state_dict(state_dict["scheduler"])

    if dataloader is not None:
        load_dataloader_state(dataloader, checkpoint_path, state_dict.get("step", 0))

    return state_dict.get("step", 0)