"""star_capture_activations_v2.py"""
import os
import torch
import folder_paths

CAPTURE_DIR = os.path.join(folder_paths.get_output_directory(), "star_activations")
os.makedirs(CAPTURE_DIR, exist_ok=True)

TARGET_SUFFIXES = (
    "attn.wq", "attn.wk", "attn.wv", "attn.wo", "attn.gate",
    "mlp.gate", "mlp.up", "mlp.down",
)

class StarCaptureActivationsV2:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "run_tag": ("STRING", {"default": "capture1"}),
                "max_calls_per_layer": ("INT", {"default": 3, "min": 1, "max": 50}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "wrap"
    CATEGORY = "⭐StarNodes/Model Tools"

    def wrap(self, model, run_tag, max_calls_per_layer):
        real_model = model.model
        target = real_model.diffusion_model if hasattr(real_model, "diffusion_model") else real_model
        out_dir = os.path.join(CAPTURE_DIR, run_tag)
        os.makedirs(out_dir, exist_ok=True)

        counters = {}
        handles = []

        def make_hook(name):
            safe_name = name.replace(".", "_")
            def hook(module, inputs, output):
                c = counters.get(name, 0)
                if c >= max_calls_per_layer:
                    return
                counters[name] = c + 1
                x = inputs[0].detach().to("cpu", dtype=torch.float16).contiguous()
                torch.save(x, os.path.join(out_dir, f"{safe_name}__call{c}.pt"))
            return hook

        matched = 0
        for name, module in target.named_modules():
            # match by name suffix + presence of a 2D weight param, not isinstance
            has_2d_weight = hasattr(module, "weight") and getattr(module.weight, "ndim", 0) == 2
            if has_2d_weight and any(name.endswith(s) for s in TARGET_SUFFIXES):
                h = module.register_forward_hook(make_hook(name))
                handles.append(h)
                matched += 1

        print(f"⭐ StarCaptureActivationsV2: hooked {matched} layers -> {out_dir}")
        model._star_capture_handles = handles
        return (model,)


NODE_CLASS_MAPPINGS = {"StarCaptureActivationsV2": StarCaptureActivationsV2}
NODE_DISPLAY_NAME_MAPPINGS = {"StarCaptureActivationsV2": "⭐ Star Capture Activations V2"}