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

"""SGLang model loader for GPU Memory Service integration.

Provides a model loader that loads weights via GMS for cross-process sharing.
The loader uses RW_OR_RO mode: first process loads from disk (RW), subsequent
processes import from GMS metadata (RO).

Usage:
    Set --load-format gms when launching SGLang.
"""

from __future__ import annotations

import logging

import torch
from gpu_memory_service.client.torch.module import materialize_module_from_gms
from gpu_memory_service.common.locks import GrantedLockType
from gpu_memory_service.integrations.common import patch_empty_cache
from gpu_memory_service.integrations.common.utils import (
    setup_meta_tensor_workaround,
    strip_gms_model_loader_config,
)
from gpu_memory_service.integrations.sglang.memory_saver import (
    get_gms_memory_saver_impl,
)
from gpu_memory_service.integrations.sglang.patches import (
    patch_model_runner,
    patch_static_state_for_gms,
    patch_torch_memory_saver,
)

logger = logging.getLogger(__name__)

# Apply patches at module import time.
# This module is only imported when load_format="gms" is used.
# Because SGLang scheduler processes use multiprocessing spawn, these patches
# must run inside the child process.  The import chain that triggers this is:
#   child unpickles server_args.load_format -> imports GMSModelLoader -> here.
patch_empty_cache()
patch_torch_memory_saver()
patch_model_runner()
patch_static_state_for_gms()
logger.info("[GMS] Applied patches")


class GMSModelLoader:
    """SGLang model loader that loads/imports weights via GPU Memory Service."""

    def __init__(self, load_config):
        self.load_config = load_config
        self._default_loader = None

    def _get_default_loader(self):
        if self._default_loader is None:
            from sglang.srt.model_loader.loader import DefaultModelLoader

            config = strip_gms_model_loader_config(
                self.load_config,
                load_format="auto",
            )
            self._default_loader = DefaultModelLoader(config)
        return self._default_loader

    def load_model(
        self,
        *,
        model_config,
        device_config,
    ) -> torch.nn.Module:
        """Load or import model weights."""
        impl = get_gms_memory_saver_impl()
        if impl is None:
            raise RuntimeError(
                "GMS impl not initialized. "
                "Ensure torch_memory_saver patch was applied before model loading."
            )

        mode = impl.allocators["weights"].granted_lock_type
        logger.info("[GMS] Loading model in %s mode", mode.name)

        if mode == GrantedLockType.RO:
            return self._load_import_only(model_config, device_config, impl)
        return self._load_write_mode(model_config, device_config, impl)

    def _load_write_mode(self, model_config, device_config, impl) -> torch.nn.Module:
        """Load model from disk and register with GMS (WRITE mode)."""
        default_loader = self._get_default_loader()

        model = default_loader.load_model(
            model_config=model_config,
            device_config=device_config,
        )

        impl.finalize_write_mode(model)
        return model

    def _load_import_only(self, model_config, device_config, impl) -> torch.nn.Module:
        """Import model weights from GMS metadata (READ mode)."""
        allocator = impl.allocators["weights"]

        device_index = torch.cuda.current_device()
        model = self._create_meta_model(model_config, device_config)

        materialize_module_from_gms(allocator, model, device_index=device_index)
        impl.imported_weights_bytes = allocator.total_bytes

        logger.info(
            "[GMS] READ mode: imported %.2f GiB from metadata",
            allocator.total_bytes / (1 << 30),
        )
        return model.eval()

    def _create_meta_model(self, model_config, device_config) -> torch.nn.Module:
        """Create model on meta device for import-only mode."""
        from sglang.srt.model_loader import get_model

        setup_meta_tensor_workaround()

        original_device = torch.cuda.current_device()
        meta_device = torch.device("meta")

        with meta_device:
            model = get_model(
                model_config=model_config,
                load_config=strip_gms_model_loader_config(
                    self.load_config,
                    load_format="dummy",
                ),
                device_config=device_config,
            )

        torch.cuda.set_device(original_device)

        try:
            from sglang.srt.model_loader.utils import (
                process_model_weights_after_loading,
            )

            process_model_weights_after_loading(model, model_config)
        except Exception as e:
            logger.debug("[GMS] Post-processing on meta tensors: %s", e)

        return model
