import io
import logging
import time
from typing import Optional

import torch
import torch.distributed as dist

logger = logging.getLogger(__name__)


def init_gloo(master_ip: str, port: int, rank: int, world_size: int) -> None:
    dist.init_process_group(
        backend="gloo",
        init_method=f"tcp://{master_ip}:{port}",
        rank=rank,
        world_size=world_size,
    )
    logger.info(
        f"[rank {rank}] GLOO ready — world_size={world_size} master={master_ip}:{port}"
    )


def diloco_outer_step(
    model: torch.nn.Module,
    ref_state: dict,
    velocity: Optional[dict],
    outer_lr: float = 0.7,
    outer_momentum: float = 0.9,
) -> dict:
    t0 = time.time()

    params = list(model.parameters())
    flat = torch.cat([p.data.view(-1).float() for p in params])
    dist.all_reduce(flat, op=dist.ReduceOp.AVG)
    offset = 0
    for p in params:
        n = p.data.numel()
        p.data.copy_(flat[offset : offset + n].view_as(p.data).to(p.dtype))
        offset += n
    dist.barrier()

    avg_state = {k: v.clone() for k, v in model.state_dict().items()}

    if velocity is None:
        velocity = {
            k: torch.zeros_like(ref_state[k], dtype=torch.float32) for k in ref_state
        }

    new_state = {}
    for k in ref_state:
        ref = ref_state[k].float()
        avg = avg_state[k].float()
        pseudo_grad = ref - avg

        v_k = outer_momentum * velocity[k] + pseudo_grad
        new_params = ref - outer_lr * (pseudo_grad + outer_momentum * v_k)

        new_state[k] = new_params.to(ref_state[k].dtype)
        velocity[k] = v_k

    model.load_state_dict(new_state, strict=True)

    elapsed = time.time() - t0
    logger.info(f"GLOO outer sync done in {elapsed:.1f}s")
    return velocity


def save_outer_checkpoint(
    s3_client,
    bucket: Optional[str],
    run_name: str,
    outer_step: int,
    inner_step: int,
    node_rank: int,
    model: torch.nn.Module,
    optimizer_state: dict,
    velocity: Optional[dict],
) -> None:
    if s3_client is None or not bucket:
        return
    ckpt = {
        "outer_step": outer_step,
        "inner_step": inner_step,
        "model": {k: v.cpu() for k, v in model.state_dict().items()},
        "optimizer": optimizer_state,
        "velocity": {k: v.cpu() for k, v in velocity.items()} if velocity else None,
    }
    buf = io.BytesIO()
    torch.save(ckpt, buf)
    key = f"{run_name}/checkpoints/node_{node_rank:04d}/latest.pt"
    s3_client.put_object(Bucket=bucket, Key=key, Body=buf.getvalue())
    logger.info(
        f"[rank {node_rank}] checkpoint outer_step={outer_step} → s3://{bucket}/{key}"
    )


def load_latest_checkpoint(
    s3_client,
    bucket: Optional[str],
    run_name: str,
    node_rank: int,
) -> Optional[dict]:
    if s3_client is None or not bucket:
        return None
    key = f"{run_name}/checkpoints/node_{node_rank:04d}/latest.pt"
    try:
        obj = s3_client.get_object(Bucket=bucket, Key=key)
        ckpt = torch.load(
            io.BytesIO(obj["Body"].read()), map_location="cpu", weights_only=False
        )
        logger.info(f"[rank {node_rank}] resumed from inner_step={ckpt['inner_step']}")
        return ckpt
    except Exception:
        return None