from pathlib import Path

from .vit import vit

# Use the centralized weights directory: sltk/weights/wilor/
# Path: __file__ -> backbones -> models -> wilor -> wilor_deps -> extraction -> sltk -> weights/wilor
_WEIGHTS_DIR = Path(__file__).parent.parent.parent.parent.parent.parent / "weights" / "wilor"


def create_backbone(cfg):
    if cfg.MODEL.BACKBONE.TYPE == "vit":
        return vit(cfg)
    elif cfg.MODEL.BACKBONE.TYPE == "fast_vit":
        import torch
        from timm.models import create_model

        fast_vit = create_model("fastvit_ma36", drop_path_rate=0.2)
        # SECURITY: Use weights_only=True to prevent arbitrary code execution
        fastvit_path = _WEIGHTS_DIR / "fastvit_ma36.pt"
        if not fastvit_path.exists():
            raise FileNotFoundError(
                f"FastViT checkpoint not found at {fastvit_path}. "
                f"Place weights in sltk/weights/wilor/ or set VLTK_WILOR_CHECKPOINT env var."
            )
        checkpoint = torch.load(str(fastvit_path), weights_only=True)
        fast_vit.load_state_dict(checkpoint["state_dict"])
        return fast_vit

    else:
        raise NotImplementedError("Backbone type is not implemented")
