from __future__ import annotations

import dataclasses
import sys
from pathlib import Path
from typing import Literal

import torch

from sarasa.utils import world_size

"""
Variable configuration dataclasses for model, optimizer, lr scheduler, and data
These classes have `create` methods to instantiate the actual objects

Users can define their own configuration dataclasses and pass them to Config.from_cli to use custom components
"""
from sarasa.data import DataConfig as Data  # noqa
from sarasa.models import ModelConfig as Model  # noqa
from sarasa.optimizers import AdamW  # noqa


@dataclasses.dataclass
class _WithSeqLen:
    seq_len: int | None = None


@dataclasses.dataclass
class LRScheduler:
    warmup_steps: int = 200
    decay_ratio: float | None = None
    """If set, the ratio of total steps to apply decay after warmup. If None, decay starts immediately after warmup."""

    decay_type: Literal["linear", "cosine", "sqrt"] = "linear"
    min_lr_factor: float = 0.0

    def create(
        self,
        optimizer: torch.optim.Optimizer,
        total_iters: int,
    ) -> torch.optim.lr_scheduler.SequentialLR:
        assert self.decay_ratio is None or (0 <= self.decay_ratio <= 1), "decay_ratio must be between 0 and 1"
        warmup_steps = self.warmup_steps
        stay_steps = 0 if self.decay_ratio is None else int(total_iters * (1 - self.decay_ratio)) - warmup_steps
        decay_steps = total_iters - warmup_steps - stay_steps
        assert warmup_steps >= 0 and decay_steps >= 0 and stay_steps >= 0, (
            f"Invalid lr scheduler steps configuration: {warmup_steps=}, {decay_steps=}, {stay_steps=}"
        )

        # 1 / max(1, warmup_steps) to avoid division by zero
        warmup = torch.optim.lr_scheduler.LinearLR(optimizer, 1 / max(1, warmup_steps), total_iters=warmup_steps)

        stay = torch.optim.lr_scheduler.ConstantLR(optimizer=optimizer, factor=1.0, total_iters=stay_steps)

        match self.decay_type:
            case "linear":
                decay = torch.optim.lr_scheduler.LinearLR(
                    optimizer,
                    start_factor=1.0,
                    end_factor=self.min_lr_factor,
                    total_iters=decay_steps,
                )
            case "sqrt":
                decay = torch.optim.lr_scheduler.LambdaLR(
                    optimizer,
                    lr_lambda=lambda step: (
                        max(
                            self.min_lr_factor,
                            (decay_steps - step) / decay_steps,
                        )
                        ** 0.5
                    ),
                )
            case "cosine":
                decay = torch.optim.lr_scheduler.CosineAnnealingLR(
                    optimizer,
                    T_max=decay_steps,
                    eta_min=optimizer.param_groups[0]["lr"] * self.min_lr_factor,
                )
            case _:
                raise ValueError(f"Unsupported decay type: {self.decay_type}")

        scheduler = torch.optim.lr_scheduler.SequentialLR(
            optimizer,
            [warmup, stay, decay],
            [self.warmup_steps, self.warmup_steps + stay_steps],
        )
        return scheduler


"""
Static configuration dataclasses

These classes are not expected to be changed by the user
"""


@dataclasses.dataclass
class Train:
    steps: int = 10_000

    grad_clip: float | None = None

    dtype: Literal["bfloat16", "float32"] = "float32"
    """Dtype used for model initialization"""

    amp_dtype: Literal["bfloat16", "float16", "float32"] = "bfloat16"
    """Dtype used for automatic mixed precision training"""

    compile: bool = False

    gc_freq: int = 500
    """Garbage collection frequency (in steps). If -1, no periodic GC is performed."""

    local_batch_size: int = 32
    """local (per device) batch size"""

    global_batch_size: int = 256
    """
    global (across all devices) batch size, used to compute 
    grad_accum_steps = global_batch_size // (local_batch_size * num_devices)
    """

    use_fa4: bool = False
    """Whether to use FA4 flash attention if available."""

    use_sac: bool = False
    """Whether to use selective activation checkpointing."""

    def __post_init__(self):
        assert self.global_batch_size % self.local_batch_size == 0
        assert self.global_batch_size % world_size() == 0
        assert self.global_batch_size >= self.local_batch_size * world_size()


@dataclasses.dataclass
class Evaluate:
    freq: int = -1
    """Evaluation frequency (in steps). If -1, no evaluation is performed."""

    val_size: int = 0
    """Number of samples in the evaluation set. If 0, no evaluation is performed.
    Must be divisible by (local_batch_size * num_devices)."""

    local_batch_size: int | None = None
    """local (per device) batch size for evaluation. If None, uses train.local_batch"""


@dataclasses.dataclass
class Profile:
    enabled: bool = False
    wait: int = 10  # the number of steps to wait before warming up the profiler
    warmup: int = 10  # the number of steps to warm up the profiler before starting to record
    active: int = 10  # the number of steps for the active record
    repeat: int = 1  # the number of cycles, 0 means always


@dataclasses.dataclass
class Metrics:
    freq: int = 10
    use_tensorboard: bool = False
    all_node: bool = False


@dataclasses.dataclass
class Checkpoint:
    freq: int = 1000
    async_mode: Literal["none", "default", "mem_pinned"] = "default"


@dataclasses.dataclass
class Distributed:
    dp_replicate_degree: int = 1
    """Degree of DDP. 1 means no replication"""

    dp_shard_degree: int = -1
    """Degree of FSDP. -1 means full sharding"""

    backend: Literal["nccl", "gloo"] = "nccl"

    init_timeout_seconds: int = 300
    """Timeout for initializing the distributed process group."""

    train_timeout_seconds: int = 100
    """Timeout for distributed training operations after the first iteration."""

    reshard_after_forward: bool = False
    """Whether to reshard model parameters after each forward pass (FSDP only)."""

    def __post_init__(self):
        if not (self.dp_replicate_degree == 1 and self.dp_shard_degree == -1):
            raise NotImplementedError()


@dataclasses.dataclass
class Config[ModelT: _WithSeqLen, OptimizerT, LRSchedulerT, DataT: _WithSeqLen]:
    # variable components
    model: ModelT
    optim: OptimizerT
    lr_scheduler: LRSchedulerT
    data: DataT

    # static components
    train: Train = dataclasses.field(default_factory=Train)
    evaluate: Evaluate = dataclasses.field(default_factory=Evaluate)
    profile: Profile = dataclasses.field(default_factory=Profile)
    metrics: Metrics = dataclasses.field(default_factory=Metrics)
    checkpoint: Checkpoint = dataclasses.field(default_factory=Checkpoint)
    distributed: Distributed = dataclasses.field(default_factory=Distributed)

    seed: int = 0
    debug: bool = False
    """ Enable debug mode with more verbose logging and checks."""

    output_dir: Path = Path("./outputs")
    """Directory to save checkpoints and logs."""

    config_file: Path | None = None
    """Path to a config file (JSON or TOML) to load configuration from."""

    def __post_init__(self):
        if self.output_dir is not None:
            self.output_dir.mkdir(parents=True, exist_ok=True)

        if hasattr(self.model, "seq_len"):
            if self.model.seq_len is None and self.data.seq_len is not None:
                self.model.seq_len = self.data.seq_len
            if self.model.seq_len is None:
                raise ValueError("seq_len must be specified in either model or data configuration.")

        if self.evaluate.local_batch_size is None:
            self.evaluate.local_batch_size = self.train.local_batch_size

        if self.evaluate.freq > 0:
            assert self.evaluate.val_size > 0, "evaluate.val_size must be > 0 if evaluate.freq > 0"
            assert self.evaluate.val_size % (self.evaluate.local_batch_size * world_size()) == 0, (
                "evaluate.val_size must be divisible by (evaluate.local_batch_size * num_devices)"
            )

    @classmethod
    def create(
        cls,
        model: ModelT,
        optim: OptimizerT,
        lr_scheduler: LRSchedulerT,
        data: DataT,
        **kwargs,
    ) -> Config:
        return cls(
            model=model,
            optim=optim,
            lr_scheduler=lr_scheduler,
            data=data,
            **kwargs,
        )

    @classmethod
    def from_cli(
        cls,
        *,
        model_type: ModelT = Model,  # pyrefly: ignore [bad-function-definition]
        optim_type: OptimizerT = AdamW,  # pyrefly: ignore [bad-function-definition]
        lr_scheduler_type: LRSchedulerT = LRScheduler,  # pyrefly: ignore [bad-function-definition]
        data_type: DataT = Data,  # pyrefly: ignore [bad-function-definition]
    ) -> Config:
        """
        initialize JobConfig from command line arguments
        update the values with the following priority: CLI arguments > config file > defaults

        *_type can be used to specify custom dataclass types for each section
        >> config = Config.from_cli(optim_type=CustomOptimizerConfig)
        """

        import importlib.util

        import tyro

        from sarasa.utils import rank

        loaded_config = None

        if (under := ("--config_file" in sys.argv)) or ("--config-file" in sys.argv):
            config_file = sys.argv[sys.argv.index("--config_file" if under else "--config-file") + 1]
            config_file = Path(config_file)

            if not config_file.exists():
                raise FileNotFoundError(f"Config file {config_file} does not exist.")

            if config_file.suffix != ".py":
                raise ValueError("Only Python config files are supported in this method.")

            spec = importlib.util.spec_from_file_location("custom_config", config_file)
            if spec is not None and spec.loader is not None:
                module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(module)
            configs = [
                config
                for config in module.__dict__.values()
                if isinstance(config, cls) and not isinstance(config, type)
            ]
            if len(configs) == 0:
                raise ValueError(f"No Config instance found in {config_file}.")
            elif len(configs) > 1:
                raise ValueError(f"Multiple Config instances found in {config_file}. Please keep only one.")
            else:
                loaded_config = configs[0]

        return tyro.cli(
            # pyrefly: ignore [unsupported-operation]
            cls[
                model_type,
                optim_type,
                lr_scheduler_type,
                data_type,
            ],
            default=loaded_config,
            console_outputs=(rank() == 0),
            config=(tyro.conf.UsePythonSyntaxForLiteralCollections,),
            # UsePythonSyntaxForLiteralCollections allows dict/list/tuple to be specified in CLI arguments
            # e.g., --model.extra '{"key": "value"}' or --model.extra '["val1", "val2"]'
        )


__all__ = [
    "Config",
    "Model",
    "AdamW",
    "LRScheduler",
    "Data",
    "Train",
    "Evaluate",
    "Metrics",
    "Checkpoint",
    "Distributed",
]
