from __future__ import annotations

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

import torch

from sarasa.data import DataConfig as Data
from sarasa.models import ModelConfig as Model
from sarasa.optimizers import AdamW
from sarasa.utils import world_size


@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__ = [
    "AdamW",
    "Checkpoint",
    "Config",
    "Data",
    "Distributed",
    "Evaluate",
    "LRScheduler",
    "Metrics",
    "Model",
    "Train",
]
