from __future__ import annotations
import argparse
import functools
import json
import logging
import math
import os
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
    FullStateDictConfig,
    StateDictType,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    apply_activation_checkpointing,
    CheckpointImpl,
)
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from riemann.dataset import PackedTextDataset
from riemann.growth_ladder import get_stage, GrowthStage
from riemann.modeling_riemann import RiemannForCausalLM, TransformerBlock
from riemann.paths import DATA_DIR, CHECKPOINT_ROOT

logger = logging.getLogger("riemann.pretrain_cuda")
TOKENIZER_NAME = "gpt2"


def setup_distributed(cpu: bool) -> tuple[int, int, int]:
    if cpu:
        dist.init_process_group(backend="gloo")
        return dist.get_rank(), dist.get_world_size(), 0
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count()))
    torch.cuda.set_device(local_rank)
    return rank, world, local_rank


def make_optimizer(decay, no_decay, lr: float, weight_decay: float, cpu: bool):
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    if not cpu:
        try:
            return torch.optim.AdamW(
                groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, fused=True
            )
        except (RuntimeError, TypeError, ValueError) as e:
            logger.warning(
                "fused AdamW unavailable on this backend (%s); using foreach", e
            )
    return torch.optim.AdamW(groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, foreach=True)


def is_main(rank: int) -> bool:
    return rank == 0


def load_training_state(checkpoint_dir: Path) -> dict:
    p = checkpoint_dir / "training_state.json"
    if p.exists():
        return json.loads(p.read_text(encoding="utf-8"))
    return {"stage": None, "global_step": 0}


def lr_at(step: int, total: int, warmup: int, peak: float, floor_ratio: float) -> float:
    if step < warmup:
        return peak * (step + 1) / max(1, warmup)
    if total <= warmup:
        return peak
    progress = (step - warmup) / max(1, total - warmup)
    progress = min(1.0, max(0.0, progress))
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return peak * (floor_ratio + (1.0 - floor_ratio) * cosine)


def count_corpus_lines(path: Path) -> int:
    n = 0
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(64 * 1024 * 1024), b""):
            n += chunk.count(b"\n")
    return n


def estimate_steps_per_epoch(
    corpus_lines: int,
    avg_tokens_per_doc: float,
    seq_len: int,
    world: int,
    batch: int,
    accum: int,
) -> int:
    total_tokens = corpus_lines * avg_tokens_per_doc
    sequences = total_tokens / seq_len
    per_step = world * batch * accum
    return max(1, int(sequences / per_step))


def gather_full_state(model):
    if isinstance(model, FSDP):
        cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
        with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
            return model.state_dict()
    return {k: v.detach().cpu() for k, v in model.module.state_dict().items()}


def clip_grads(model, max_norm: float):
    if isinstance(model, FSDP):
        return model.clip_grad_norm_(max_norm)
    return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)


def save_checkpoint(
    model,
    tokenizer,
    checkpoint_dir: Path,
    stage: GrowthStage,
    global_step: int,
    rank: int,
    max_shard_size: str,
) -> None:
    state = gather_full_state(model)
    if is_main(rank):
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
        for old in list(checkpoint_dir.glob("model*.safetensors*")) + list(
            checkpoint_dir.glob("pytorch_model*.bin*")
        ):
            try:
                old.unlink()
            except PermissionError:
                stale = old.with_name(old.name + ".stale")
                try:
                    old.replace(stale)
                except OSError as e:
                    logger.warning("could not remove or rename %s: %s", old.name, e)
        inner = model.module
        inner.save_pretrained(
            str(checkpoint_dir),
            state_dict={k: v.to(torch.float32) for k, v in state.items()},
            safe_serialization=True,
            max_shard_size=max_shard_size,
        )
        tokenizer.save_pretrained(str(checkpoint_dir))
        (checkpoint_dir / "training_state.json").write_text(
            json.dumps({"stage": stage.name, "global_step": global_step}, indent=2),
            encoding="utf-8",
        )
        logger.info("Saved checkpoint at step %d -> %s", global_step, checkpoint_dir)
    dist.barrier()


def cuda_defaults(stage: GrowthStage, world: int) -> tuple[int, int]:
    params = stage.n_layer * 12 * stage.n_embd * stage.n_embd
    if params < 1e9:
        per_gpu_batch = 16
    elif params < 4e9:
        per_gpu_batch = 8
    elif params < 10e9:
        per_gpu_batch = 4
    else:
        per_gpu_batch = 2
    target_tokens_per_step = 1_048_576
    accum = max(
        1, target_tokens_per_step // (world * per_gpu_batch * stage.n_positions)
    )
    return per_gpu_batch, accum


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", required=True)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--max-steps", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--grad-accum", type=int, default=None)
    parser.add_argument("--seq-len", type=int, default=None)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--min-lr-ratio", type=float, default=0.1)
    parser.add_argument("--warmup-steps", type=int, default=200)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--avg-tokens-per-doc", type=float, default=None)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--log-every", type=int, default=20)
    parser.add_argument("--save-every", type=int, default=1000)
    parser.add_argument("--max-shard-size", default="5GB")
    parser.add_argument("--no-activation-checkpointing", action="store_true")
    parser.add_argument("--sharding", choices=["full", "grad_op"], default="full")
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument("--dry-run-cpu", action="store_true")
    parser.add_argument("--data-dir", default=str(DATA_DIR))
    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
    args = parser.parse_args()

    cpu = args.dry_run_cpu
    rank, world, local_rank = setup_distributed(cpu)
    logging.basicConfig(
        level=logging.INFO if is_main(rank) else logging.WARNING,
        format=f"%(asctime)s %(levelname)s [rank {rank}] %(message)s",
    )
    device = torch.device("cpu") if cpu else torch.device("cuda", local_rank)
    stage = get_stage(args.stage)
    default_batch, default_accum = cuda_defaults(stage, world)
    batch_size = args.batch_size or default_batch
    grad_accum = args.grad_accum or default_accum
    seq_len = args.seq_len or stage.n_positions
    checkpoint_dir = Path(args.checkpoint_root) / stage.name
    corpus_path = Path(args.data_dir) / "pretrain_corpus.txt"
    if not (checkpoint_dir / "config.json").exists():
        raise FileNotFoundError(
            f"No checkpoint at {checkpoint_dir}. Grow into {stage.name} first with grow_and_continue.py."
        )
    if not corpus_path.exists():
        raise FileNotFoundError(f"{corpus_path} not found")

    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if is_main(rank):
        logger.info("Loading %s from %s", stage.name, checkpoint_dir)
    model = RiemannForCausalLM.from_pretrained(
        str(checkpoint_dir), torch_dtype=torch.float32
    )
    model.config.use_cache = False
    model.gradient_checkpointing = False

    if cpu:
        model = DDP(model)
    else:
        wrap_policy = functools.partial(
            transformer_auto_wrap_policy, transformer_layer_cls={TransformerBlock}
        )
        mixed = MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.bfloat16,
            buffer_dtype=torch.bfloat16,
        )
        model = FSDP(
            model,
            auto_wrap_policy=wrap_policy,
            mixed_precision=mixed,
            sharding_strategy=(
                ShardingStrategy.SHARD_GRAD_OP
                if args.sharding == "grad_op"
                else ShardingStrategy.FULL_SHARD
            ),
            device_id=local_rank,
            use_orig_params=True,
            limit_all_gathers=True,
        )
    if not args.no_activation_checkpointing:
        apply_activation_checkpointing(
            model,
            checkpoint_wrapper_fn=functools.partial(
                checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT
            ),
            check_fn=lambda m: isinstance(m, TransformerBlock),
        )
    model.train()

    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (
            no_decay
            if p.ndim < 2 or "layernorm" in name.lower() or "bias" in name
            else decay
        ).append(p)
    optimizer = make_optimizer(decay, no_decay, args.lr, args.weight_decay, cpu)

    state = load_training_state(checkpoint_dir)
    global_step = int(state.get("global_step", 0))
    corpus_lines = count_corpus_lines(corpus_path) if is_main(rank) else 0
    tokens_per_doc = args.avg_tokens_per_doc
    count_path = corpus_path.with_name("corpus_token_count.json")
    if tokens_per_doc is None:
        if count_path.exists():
            tokens_per_doc = float(
                json.loads(count_path.read_text(encoding="utf-8"))["tokens_per_doc"]
            )
        else:
            tokens_per_doc = 150.0
            if is_main(rank):
                logger.warning(
                    "%s not found; assuming %.0f tokens per document. Run scripts/count_corpus_tokens.py for an exact epoch size.",
                    count_path.name,
                    tokens_per_doc,
                )
    lines_t = torch.tensor([corpus_lines], device=device)
    dist.broadcast(lines_t, src=0)
    corpus_lines = int(lines_t.item())
    steps_per_epoch = estimate_steps_per_epoch(
        corpus_lines, tokens_per_doc, seq_len, world, batch_size, grad_accum
    )
    total_steps = args.max_steps or int(steps_per_epoch * args.epochs)
    tokens_per_step = world * batch_size * grad_accum * seq_len
    if is_main(rank):
        n_params = sum(p.numel() for p in model.parameters())
        logger.info(
            "Training %s: world=%d batch=%d accum=%d seq=%d tokens/step=%d, corpus=%d docs, ~%d steps/epoch, target %d steps (%.2f epochs), %.2fB params",
            stage.name,
            world,
            batch_size,
            grad_accum,
            seq_len,
            tokens_per_step,
            corpus_lines,
            steps_per_epoch,
            total_steps,
            args.epochs,
            n_params / 1e9,
        )

    epoch = 0
    accum_loss = torch.zeros((), device=device)
    micro = 0
    tokens_seen = 0
    start = time.time()
    done = False
    while not done:
        dataset = PackedTextDataset(
            corpus_path,
            tokenizer,
            seq_len,
            shard_id=rank,
            num_shards=world,
            epoch=epoch,
            seed=args.seed,
        )
        loader = DataLoader(
            dataset,
            batch_size=batch_size,
            num_workers=args.num_workers,
            pin_memory=not cpu,
            drop_last=True,
        )
        for batch in loader:
            batch = batch.to(device, non_blocking=True)
            is_last_micro = (micro + 1) % grad_accum == 0
            ctx = model.no_sync() if not is_last_micro else _nullcontext()
            with ctx:
                out = model(input_ids=batch, labels=batch)
                (out.loss / grad_accum).backward()
            accum_loss += out.loss.detach()
            micro += 1
            tokens_seen += batch.numel()
            if not is_last_micro:
                continue
            lr = lr_at(
                global_step, total_steps, args.warmup_steps, args.lr, args.min_lr_ratio
            )
            for g in optimizer.param_groups:
                g["lr"] = lr
            if args.grad_clip > 0:
                clip_grads(model, args.grad_clip)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            global_step += 1
            if global_step % args.log_every == 0:
                loss_t = accum_loss / (args.log_every * grad_accum)
                dist.all_reduce(loss_t, op=dist.ReduceOp.SUM)
                loss_t = loss_t / world
                if is_main(rank):
                    elapsed = time.time() - start
                    tps = (tokens_seen * world) / max(1e-6, elapsed)
                    remaining = (
                        (total_steps - global_step) * tokens_per_step / max(1e-6, tps)
                    )
                    logger.info(
                        "step %d/%d | loss %.4f | lr %.2e | %.0f tok/s | %.1fh elapsed | ~%.1fh remaining",
                        global_step,
                        total_steps,
                        loss_t.item(),
                        lr,
                        tps,
                        elapsed / 3600,
                        remaining / 3600,
                    )
                accum_loss.zero_()
            if global_step % args.save_every == 0:
                save_checkpoint(
                    model,
                    tokenizer,
                    checkpoint_dir,
                    stage,
                    global_step,
                    rank,
                    args.max_shard_size,
                )
            if global_step >= total_steps:
                done = True
                break
        epoch += 1
        if not done and is_main(rank):
            logger.info("Finished corpus pass %d at step %d", epoch, global_step)
    save_checkpoint(
        model, tokenizer, checkpoint_dir, stage, global_step, rank, args.max_shard_size
    )
    if is_main(rank):
        logger.info(
            "Done. %s at step %d, %.2fB tokens seen, %.2fh wall clock, checkpoint at %s",
            stage.name,
            global_step,
            tokens_seen * world / 1e9,
            (time.time() - start) / 3600,
            checkpoint_dir,
        )
    dist.barrier()
    dist.destroy_process_group()


class _nullcontext:

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False


if __name__ == "__main__":
    main()