from __future__ import annotations

from pydantic_ai import Agent
from pydantic_ai.models import Model

from generator.prompts import MUTATION_SYSTEM, format_technique_hint
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,
        technique_hints: list[str] | None = None,
        sample_fn: object = None,
    ) -> list[str]:
        if not seeds:
            return []

        # Use weighted sampler from SeedPool if provided, else fall back to first seed
        if callable(sample_fn):
            seed = sample_fn()
        else:
            seed = seeds[0]

        hint = format_technique_hint(technique_hints or [])
        agent = Agent(
            model,
            system_prompt=MUTATION_SYSTEM.format(count=count, technique_hint=hint),
        )
        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"
