# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import logging
import os
from copy import deepcopy

import wandb

logger = logging.getLogger(__name__)


def _is_offline_mode(args) -> bool:
    """Detect whether W&B should run in offline mode.

    Priority order:
    1) args.wandb_mode if provided
    2) WANDB_MODE environment variable
    """
    if args.wandb_mode:
        return args.wandb_mode == "offline"
    return os.environ.get("WANDB_MODE") == "offline"


def init_wandb_primary(args):
    if not args.use_wandb:
        args.wandb_run_id = None
        return

    # Set W&B mode if specified (overrides WANDB_MODE env var)
    if args.wandb_mode:
        os.environ["WANDB_MODE"] = args.wandb_mode
        if args.wandb_mode == "offline":
            logger.info("W&B offline mode enabled. Data will be saved locally.")
        elif args.wandb_mode == "disabled":
            logger.info("W&B disabled mode enabled. No data will be logged.")
        elif args.wandb_mode == "online":
            logger.info("W&B online mode enabled. Data will be uploaded to cloud.")

    offline = _is_offline_mode(args)

    # Only perform explicit login when NOT offline
    if (not offline) and args.wandb_key is not None:
        wandb.login(key=args.wandb_key, host=args.wandb_host)

    # Check if we should resume a previous run
    # Priority: 1) wandb_resume_run_id from args, 2) wandb_run_id from args, 3) wandb_run_id from checkpoint file
    resume_run_id = getattr(args, "wandb_resume_run_id", None) or getattr(args, "wandb_run_id", None)
    if not resume_run_id:
        resume_run_id = _load_wandb_run_id_from_checkpoint(args)
    
    if resume_run_id:
        # Resume an existing run
        logger.info(f"Resuming W&B run with id: {resume_run_id}")
        init_kwargs = {
            "id": resume_run_id,
            "entity": args.wandb_team,
            "project": args.wandb_project,
            "resume": "must",  # Fail if run doesn't exist
            "config": _compute_config_for_logging(args),
        }
        
        # Configure settings based on offline/online mode
        if offline:
            init_kwargs["settings"] = wandb.Settings(mode="offline")
        else:
            init_kwargs["settings"] = wandb.Settings(mode="shared", x_primary=True)
    else:
        # Create a new run
        # add random 6 length string with characters
        if args.wandb_random_suffix:
            suffix = "_" + wandb.util.generate_id()
            max_base_len = 128 - len(suffix)
            group = args.wandb_group[:max_base_len] + suffix
            run_name = f"{group}-RANK_{args.rank}"
        else:
            group = args.wandb_group
            run_name = args.wandb_group

        # Prepare wandb init parameters
        init_kwargs = {
            "entity": args.wandb_team,
            "project": args.wandb_project,
            "group": group,
            "name": run_name,
            "config": _compute_config_for_logging(args),
        }

        # Configure settings based on offline/online mode
        if offline:
            init_kwargs["settings"] = wandb.Settings(mode="offline")
        else:
            init_kwargs["settings"] = wandb.Settings(mode="shared", x_primary=True)

    # Add custom directory if specified
    if args.wandb_dir:
        # Ensure directory exists to avoid backend crashes
        os.makedirs(args.wandb_dir, exist_ok=True)
        init_kwargs["dir"] = args.wandb_dir
        logger.info(f"W&B logs will be stored in: {args.wandb_dir}")

    wandb.init(**init_kwargs)

    _init_wandb_common()

    args.wandb_run_id = wandb.run.id
    _save_wandb_run_id_to_checkpoint(args)

    if resume_run_id:
        logger.info(f"Successfully resumed W&B run: {wandb.run.url}")


def _load_wandb_run_id_from_checkpoint(args):
    load_dir = getattr(args, "load", None)
    if not load_dir:
        return None
    path = os.path.join(load_dir, "wandb_run_id.txt")
    if os.path.exists(path):
        with open(path, "r") as f:
            run_id = f.read().strip()
        if run_id:
            logger.info(f"Loaded wandb_run_id from {path}: {run_id}")
            return run_id
    return None


def _save_wandb_run_id_to_checkpoint(args):
    save_dir = getattr(args, "save", None)
    if not save_dir:
        return
    os.makedirs(save_dir, exist_ok=True)
    path = os.path.join(save_dir, "wandb_run_id.txt")
    with open(path, "w") as f:
        f.write(args.wandb_run_id)
    logger.info(f"Saved wandb_run_id to {path}")


def _compute_config_for_logging(args):
    output = deepcopy(args.__dict__)

    whitelist_env_vars = [
        "SLURM_JOB_ID",
        # We may insert more default values here, and may also allow users to configure a whitelist
    ]
    output["env_vars"] = {k: v for k, v in os.environ.items() if k in whitelist_env_vars}

    return output


# https://docs.wandb.ai/guides/track/log/distributed-training/#track-all-processes-to-a-single-run
def init_wandb_secondary(args, router_addr=None):
    wandb_run_id = getattr(args, "wandb_run_id", None)
    if wandb_run_id is None:
        return

    # Set W&B mode if specified (same as primary)
    if args.wandb_mode:
        os.environ["WANDB_MODE"] = args.wandb_mode

    offline = _is_offline_mode(args)

    if (not offline) and args.wandb_key is not None:
        wandb.login(key=args.wandb_key, host=args.wandb_host)

    # Configure settings based on offline/online mode
    if offline:
        settings_kwargs = dict(mode="offline")
    else:
        settings_kwargs = dict(
            mode="shared",
            x_primary=False,
            x_update_finish_state=False,
        )

    if args.sglang_enable_metrics and router_addr is not None:
        logger.info(f"Forward SGLang metrics at {router_addr} to WandB.")
        settings_kwargs |= dict(
            x_stats_open_metrics_endpoints={
                "sgl_engine": f"{router_addr}/engine_metrics",
            },
            x_stats_open_metrics_filters={
                "sgl_engine.*": {},
            },
        )

    init_kwargs = {
        "id": wandb_run_id,
        "entity": args.wandb_team,
        "project": args.wandb_project,
        "config": args.__dict__,
        "resume": "allow",
        "reinit": True,
        "settings": wandb.Settings(**settings_kwargs),
    }

    # Add custom directory if specified
    if args.wandb_dir:
        os.makedirs(args.wandb_dir, exist_ok=True)
        init_kwargs["dir"] = args.wandb_dir

    wandb.init(**init_kwargs)

    _init_wandb_common()


def _init_wandb_common():
    wandb.define_metric("train/step")
    wandb.define_metric("train/*", step_metric="train/step")
    wandb.define_metric("rollout/step")
    wandb.define_metric("rollout/*", step_metric="rollout/step")
    wandb.define_metric("multi_turn/*", step_metric="rollout/step")
    wandb.define_metric("passrate/*", step_metric="rollout/step")
    wandb.define_metric("eval/step")
    wandb.define_metric("eval/*", step_metric="eval/step")
    wandb.define_metric("perf/*", step_metric="rollout/step")


def get_wandb_offline_dir(args):
    """Get the directory where offline W&B data is stored."""
    if _is_offline_mode(args):
        if args and hasattr(args, "wandb_dir") and args.wandb_dir:
            # Use custom directory if specified
            return args.wandb_dir
        else:
            # Default offline directory is ~/wandb/offline-run-<timestamp>
            # This will be created automatically by wandb
            return os.path.expanduser("~/wandb")
    return None