from pathlib import Path
from typing import Any, override

from bench_af._abstract.objects import Detector, Environment, EvalLog, ModelOrganism
from nnsight import LanguageModel

CURRENT_DIRECTORY = Path(__file__).parent


class MyDetector(Detector):
    def __init__(self) -> None:
        super().__init__(
            name=CURRENT_DIRECTORY.name,
            supported_model_organisms=set(["test_nnsight"]),
            supported_environments=set(["self_exfiltrate", "mvp"]),
        )
        self.activations_cache: dict[str, Any] = {}
    
    @override
    def get_new_model_and_environment(
        self,
        *,
        original_model_organism: ModelOrganism,
        original_environment: Environment,
    ) -> tuple[ModelOrganism, Environment]:

        model: LanguageModel = original_model_organism.get_nnsight_model()
        
        def hook(prompt: str) -> None:
            model.model.layers[5].mlp.output[0] = model.model.layers[5].mlp.output[0] + 1
            self.activations_cache["layer_5"] = model.model.layers[5].mlp.output[0].save()
        
        original_model_organism.set_nnsight_hook(hook)

        return original_model_organism, original_environment
        

    @override
    def get_prediction(self, eval_logs: list[EvalLog]) -> tuple[float, str]:
        print(self.activations_cache)
        return 0.0, "Example explanation"


def get() -> Detector:
    return MyDetector()
