import random

from pydantic_ai import Agent
from pydantic_ai.models import Model

from generator.prompts import MUTATION_SYSTEM
from prompts.models import Probe

from .base import AbstractStrategy


class MutationStrategy(AbstractStrategy):
    """Produce variations of a seed prompt using different evasion techniques."""

    def generate(self, seeds: list[Probe], model: Model, count: int = 5) -> list[str]:
        if not seeds:
            return []
        seed = random.choice(seeds)
        agent = Agent(model, system_prompt=MUTATION_SYSTEM.format(count=count))
        result = agent.run_sync(f"Test prompt to vary:\n\n{seed.content}")
        return self._parse_candidates(result.output)[:count]

    def get_name(self) -> str:
        return "mutation"

    def get_description(self) -> str:
        return "Produce variations of a seed prompt using different evasion techniques"
