from inspect_ai.model import ChatMessageUser
from inspect_ai.solver import Generate, Solver, TaskState, solver


@solver
def multi_dialogue_solver(
    skip_idx: list[int], dynamic_messages: dict[int, dict[str, str]]
) -> Solver:
    """
    skip_idx: list with messages to avoid calling the `generate` method, e.g. the initial system prompt.
    dynamic_messages: {turn_idx: {'condition_match': <logic or message to be matched>, 'content_matched': <chat content if condition matched>, 'content_not_matched': <chat content if condition NOT matched>}}
    """  # noqa: D205

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        # generate the output after each turn in the input and write to state.messages
        for idx, turn in enumerate(state.messages):
            if idx in skip_idx:
                continue
            else:
                messages_slice = state.messages[: idx + 1].copy()
                state.messages = messages_slice
                state = await generate(state)

                if idx in dynamic_messages.keys():
                    answer = state.output.completion.lower()
                    if (
                        dynamic_messages[idx]["condition_match"] in answer
                    ):  # TODO: improve parsing
                        state.messages.append(
                            ChatMessageUser(
                                content=dynamic_messages[idx]["content_matched"]
                            )
                        )
                    else:
                        state.messages.append(
                            ChatMessageUser(
                                content=dynamic_messages[idx]["content_not_matched"]
                            )
                        )
                    state = await generate(state)

        return state

    return solve
