import random

from pydantic_ai import Agent
from pydantic_ai.models import Model

from generator.prompts import OBJECTIVE_SYSTEM
from prompts.models import Probe

from .base import AbstractStrategy


class ObjectiveStrategy(AbstractStrategy):
    """Generate probes targeting a specific user-defined objective."""

    def __init__(self, targets: list[str]) -> None:
        self._targets = targets

    def generate(self, seeds: list[Probe], model: Model, count: int = 5) -> list[str]:
        if not self._targets:
            return []
        goal = random.choice(self._targets)
        agent = Agent(model, system_prompt=OBJECTIVE_SYSTEM.format(count=count, goal=goal))
        result = agent.run_sync(f"Generate {count} prompts for goal: {goal}")
        return self._parse_candidates(result.output)[:count]

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

    def get_description(self) -> str:
        preview = ", ".join(self._targets[:2])
        suffix = "…" if len(self._targets) > 2 else ""
        return f"Generate probes targeting specific objectives: {preview}{suffix}"
