# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import pathlib
import uuid
from typing import Any

from pyrit.common.apply_defaults import apply_defaults
from pyrit.common.path import CONVERTER_SEED_PROMPT_PATH
from pyrit.converter.converter import ConverterResult
from pyrit.executor.promptgen.fuzzer.fuzzer_converter_base import (
    FuzzerConverter,
)
from pyrit.models import JsonResponseConfig, Message, MessagePiece, PromptDataType, SeedPrompt
from pyrit.prompt_target import PromptTarget


class FuzzerCrossOverConverter(FuzzerConverter):
    """
    Uses multiple prompt templates to generate new prompts.
    """

    @apply_defaults
    def __init__(
        self,
        *,
        converter_target: PromptTarget | None = None,
        prompt_template: SeedPrompt | None = None,
        prompt_templates: list[str] | None = None,
    ) -> None:
        """
        Initialize the converter with the specified chat target and prompt templates.

        Args:
            converter_target (PromptTarget): Chat target used to perform fuzzing on user prompt.
                Can be omitted if a default has been configured via PyRIT initialization.
            prompt_template (SeedPrompt, Optional): Template to be used instead of the default system prompt with
                instructions for the chat target.
            prompt_templates (list[str], Optional): List of prompt templates to use in addition to the default one.
        """
        prompt_template = (
            prompt_template
            if prompt_template
            else SeedPrompt.from_yaml_file(
                pathlib.Path(CONVERTER_SEED_PROMPT_PATH) / "fuzzer_converters" / "crossover_converter.yaml"
            )
        )
        super().__init__(converter_target=converter_target, prompt_template=prompt_template)
        self.prompt_templates = prompt_templates or []
        self.template_label = "TEMPLATE 1"

    def update(self, **kwargs: Any) -> None:
        """Update the converter with new prompt templates."""
        if "prompt_templates" in kwargs:
            self.prompt_templates = kwargs["prompt_templates"]

    async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult:
        """
        Convert the given prompt by combining it with a random prompt template from the list of available templates.

        Args:
            prompt (str): The prompt to be converted.
            input_type (PromptDataType): The type of input data.

        Returns:
            ConverterResult: The result containing the modified prompt.

        Raises:
            ValueError: If the input type is not supported.
        """
        if not self.input_supported(input_type):
            raise ValueError("Input type not supported")

        if len(self.prompt_templates) == 0:
            raise ValueError(
                "No prompt templates available for crossover. Please provide prompt templates via the update method."
            )

        conversation_id = str(uuid.uuid4())

        self.converter_target.set_system_prompt(
            system_prompt=self.system_prompt,
            conversation_id=conversation_id,
        )

        formatted_prompt = f"===={self.template_label} BEGINS====\n{prompt}\n===={self.template_label} ENDS===="
        formatted_prompt += (
            "\n====TEMPLATE 2 BEGINS====\n"
            f"{self._get_random_generator(stream='prompt-template').choice(self.prompt_templates)}"
            "\n====TEMPLATE 2 ENDS====\n"
        )

        prompt_metadata = JsonResponseConfig(enabled=True).to_metadata()
        request = Message(
            message_pieces=[
                MessagePiece(
                    role="user",
                    original_value=formatted_prompt,
                    converted_value=formatted_prompt,
                    conversation_id=conversation_id,
                    sequence=1,
                    original_value_data_type=input_type,
                    converted_value_data_type=input_type,
                    converter_identifiers=[self.get_identifier()],
                    prompt_metadata=prompt_metadata,
                )
            ]
        )

        response = await self.send_prompt_async(request)

        return ConverterResult(output_text=response, output_type="text")