# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Minimal pytest fixtures for tt-sam3.

Provides a ``device`` fixture and a ``sam3_model`` fixture so the tt_sam3 tests
can be run without the tt-metal monorepo's own conftest.py. The device is
selected with ``SAM3_DEVICE`` env var (default 0).
"""

from __future__ import annotations

import gc
import os
import unittest.mock as mock

import pytest
import torch


@pytest.fixture(autouse=True)
def _gc_between_tests():
    gc.collect()


@pytest.fixture
def reset_seeds():
    """tt-metal's root conftest fixture the legacy tests request (deterministic random inputs)."""
    torch.manual_seed(213919)


@pytest.fixture(scope="session")
def device_params(request):
    """The legacy tests parametrize ``device_params`` indirectly (``{"l1_small_size": 32768}``, tt-metal's
    root conftest contract).  Session-scoped like ``device`` (every test here passes the same value, so the
    device is opened once per session); ``{}`` for tests that do not parametrize it."""
    return getattr(request, "param", {})


@pytest.fixture(scope="session")
def device(device_params):
    """Opens the device the way the server / benchmark do (``tt_sam3.tt.fused_config.device_open_kwargs``):
    bare unless the test's ``device_params`` ask for ``l1_small_size`` (the legacy PCC tests: 32768) or
    ``SAM3_L1_SMALL_SIZE`` overrides it (``0`` = the served bare configuration), and the trace region
    whenever ``TT_FUSED=1`` selects the fused encoder trace."""
    import ttnn

    from tt_sam3.tt.fused_config import FusedConfig, device_open_kwargs

    device_id = int(os.environ.get("SAM3_DEVICE", "0"))
    env_l1 = os.environ.get("SAM3_L1_SMALL_SIZE")
    l1_small = int(env_l1) if env_l1 is not None else int(device_params.get("l1_small_size", 0))
    from tt_sam3.tt.device_open import open_device

    # default ETH dispatch + 1 CQ + 12x10 grid (SAM3_DISPATCH=auto), as the server and the Python API
    dev = open_device(device_open_kwargs(FusedConfig.from_env(), device_id=device_id, l1_small_size=l1_small),
                      os.environ.get("SAM3_DISPATCH", "auto"))[0]
    yield dev
    ttnn.close_device(dev)


def _patch_cuda_to_cpu():
    """Context-manager-free CUDA→CPU patching for SAM3 model loading."""
    orig = {
        n: getattr(torch, n)
        for n in [
            "zeros", "ones", "arange", "empty", "full", "randn",
            "rand", "tensor", "linspace", "logspace", "eye",
        ]
    }

    def _redirect(fn):
        def wrapper(*args, **kwargs):
            if "device" in kwargs and kwargs["device"] is not None and "cuda" in str(kwargs["device"]):
                kwargs["device"] = "cpu"
            return fn(*args, **kwargs)
        return wrapper

    patches = [mock.patch("torch.cuda.is_available", return_value=False)]
    for name, fn in orig.items():
        patches.append(mock.patch(f"torch.{name}", _redirect(fn)))
    return patches


def _resolve_pinned_checkpoint():
    """``sam3.pt`` the way the server resolves it: ``SAM3_WEIGHTS_DIR`` (offline dir) or the HF cache at
    ``TT_WEIGHTS_REVISION``; None when neither is set (sam3 then downloads ``facebook/sam3@main``)."""
    weights_dir = os.environ.get("SAM3_WEIGHTS_DIR")
    if weights_dir:
        return os.path.join(weights_dir, "sam3.pt")
    revision = os.environ.get("TT_WEIGHTS_REVISION")
    if revision:
        from huggingface_hub import hf_hub_download

        return hf_hub_download(repo_id="facebook/sam3", filename="sam3.pt", revision=revision)
    return None


def _build_sam3_cpu(use_pretrained=False):
    # SAM3_BPE_PATH > vendored tt_sam3/assets copy > legacy open_clip path (see tt_sam3.tt.ttnn_sam3_model)
    from tt_sam3.tt.ttnn_sam3_model import resolve_bpe_path

    bpe_path = resolve_bpe_path()
    # --pretrained honours the served pin (TT_WEIGHTS_REVISION / SAM3_WEIGHTS_DIR) instead of facebook/sam3@main
    checkpoint_path = _resolve_pinned_checkpoint() if use_pretrained else None
    patches = _patch_cuda_to_cpu()
    for p in patches:
        p.start()
    try:
        from sam3.model_builder import build_sam3_image_model
        model = build_sam3_image_model(
            bpe_path=bpe_path,
            device="cpu",
            eval_mode=True,
            load_from_HF=bool(use_pretrained) and checkpoint_path is None,
            checkpoint_path=checkpoint_path,
            enable_segmentation=True,
            enable_inst_interactivity=False,
        )
    finally:
        for p in patches:
            p.stop()
    return model


@pytest.fixture(scope="session")
def sam3_model():
    return _build_sam3_cpu(use_pretrained=False)


@pytest.fixture(scope="session")
def sam3_reference_model(sam3_model):
    """Alias the segmentation tests request (random-weight shadow, as ``sam3_model``)."""
    return sam3_model


@pytest.fixture(scope="session")
def sam3_vit_backbone(sam3_model):
    from sam3.model.vitdet import ViT
    for _, module in sam3_model.named_modules():
        if isinstance(module, ViT):
            return module
    raise RuntimeError("Could not find ViT backbone in SAM3 model")


@pytest.fixture(scope="session")
def sam3_neck(sam3_model):
    from sam3.model.necks import Sam3DualViTDetNeck
    for _, module in sam3_model.named_modules():
        if isinstance(module, Sam3DualViTDetNeck):
            return module
    raise RuntimeError("Could not find neck in SAM3 model")