from inspect_ai.model import GenerateConfig, get_model
from transformer_lens import HookedTransformer
from bench_af._abstract.objects import ModelOrganism
from inspect_ai.solver import generate


# MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
# MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
MODEL_NAME = "unsloth/Llama-3.3-70B-Instruct-bnb-4bit"
tl_model = HookedTransformer.from_pretrained(MODEL_NAME)
tl_model.eval()  # Set model to evaluation mode

# 2. Define the custom solver that orchestrates the multi-turn conversation.


@solver
def multi_turn_generate_with_activations() -> Solver:
    """
    A custom solver that runs a two-turn conversation, capturing
    activations from a transformer_lens model at each turn.
    """
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        """
        This inner 'solve' function orchestrates the two generation steps.
        """

        # TURN 1: Generate a response to the initial user question.
        # The initial `state` object contains the first question from the dataset.
        activation_cache_1 = {}

        def store_activations_hook_t1(activations: torch.Tensor, hook: HookPoint) -> torch.Tensor:
            activation_cache_1[hook.name] = activations[:, -
                                                        1, :].detach().cpu()
            return activations

        hook_points_t1 = [
            (f"blocks.{layer_idx}.attn.hook_z", store_activations_hook_t1)
            for layer_idx in range(tl_model.cfg.n_layers)
        ] + [
            (f"blocks.{layer_idx}.mlp.hook_post", store_activations_hook_t1)
            for layer_idx in range(tl_model.cfg.n_layers)
        ]

        with tl_model.hooks(fwd_hooks=hook_points_t1):
            state_with_response_1 = await generate(state)

        response_1 = state_with_response_1.output.completion

        # TURN 2: Generate a response to the follow-up question.

        # Retrieve the follow-up question from the metadata we defined in the environment.
        follow_up_question = state.metadata["follow_up_question"]

        # FIX: Instead of importing the Message class, get its type from an
        # existing message object. This is more robust and avoids import errors.
        MessageClass = type(state_with_response_1.messages[0])

        # Append a proper Message object using the dynamically retrieved class.
        state_with_response_1.messages.append(
            MessageClass(role="user", content=follow_up_question)
        )

        activation_cache_2 = {}

        def store_activations_hook_t2(activations: torch.Tensor, hook: HookPoint) -> torch.Tensor:
            activation_cache_2[hook.name] = activations[:, -
                                                        1, :].detach().cpu()
            return activations

        hook_points_t2 = [
            (f"blocks.{layer_idx}.attn.hook_z", store_activations_hook_t2)
            for layer_idx in range(tl_model.cfg.n_layers)
        ] + [
            (f"blocks.{layer_idx}.mlp.hook_post", store_activations_hook_t2)
            for layer_idx in range(tl_model.cfg.n_layers)
        ]

        with tl_model.hooks(fwd_hooks=hook_points_t2):
            # Call generate again, passing the modified TaskState object which now
            # contains the full conversation history.
            state_with_response_2 = await generate(state_with_response_1)

        # Populate the final state with all the data we've collected.
        state.metadata["turn_1_response"] = response_1
        state.metadata["turn_1_activations"] = {
            name: tensor.tolist() for name, tensor in activation_cache_1.items()}
        state.metadata["turn_2_activations"] = {
            name: tensor.tolist() for name, tensor in activation_cache_2.items()}

        # The final output for the task is the model's response to the follow-up question.
        state.output = state_with_response_2.output

        print("Completed multi-turn generation and captured activations for both turns.")

        return state

    return solve

# 2. Define the solver that ONLY handles the multi-turn conversation.


@solver
def multi_turn_dialogue_solver() -> Solver:
    """
    A solver that runs a two-turn conversation without capturing activations.
    """
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        model = get_model()
        print(model)
        tl_model = model.model_args["tl_model"]
        print(tl_model)
        # TURN 1: Generate a response to the initial user question.
        # this returns a string
        state_with_response_1 = await tl_model.generate(state.user_prompt.content)
        print(state_with_response_1)
        response_1_text = state_with_response_1.output.completion

        # TURN 2: Ask the follow-up question.
        follow_up_question = state.metadata["follow_up_question"]
        MessageClass = type(state_with_response_1.messages[0])
        state_with_response_1.messages.append(
            MessageClass(role="user", content=follow_up_question)
        )

        # Call generate again with the full conversation history.
        state_with_response_2 = await tl_model.generate(state_with_response_1.user_prompt)

        # The final state should contain the full conversation.
        # Store the first turn's response for potential analysis later.
        state.metadata["turn_1_response"] = response_1_text
        state.output = state_with_response_2.output

        return state

    return solve


# 3. Configure the model for inspect_ai.
model_args = {
    "tl_model": tl_model,
    "tl_generate_args": {
        "max_new_tokens": 50,
        "temperature": 0.0,
        "do_sample": False,
    },
}

model = get_model(
    "transformer_lens/test_tl",
    config=GenerateConfig(max_tokens=50),
    **model_args
)

model_organism = ModelOrganism(
    name="test_tl",
    model=model,
    solver=[generate()],
    supported_environments={"discriminatory"},
)
