from importlib import resources
from pathlib import Path

import yaml

from abrade_cli.models import SEVERITY_RANK, Injection, Severity, SuiteMode

QUICK_SUITE_SIZE = 5


def _bundled_sample_path() -> Path:
    return resources.files("abrade_cli.data") / "sample_corpus.yaml"


def load_corpus(path: Path | None = None) -> list[Injection]:
    """Defaults to the bundled sample corpus. Pass `--corpus <path>` for a
    real, versioned corpus artifact once one exists — `abrade corpus pull`
    (an OCI registry fetch) is not implemented in this pass; there is no
    registry to pull from yet. See data/sample_corpus.yaml's own header.
    """
    source = path if path is not None else _bundled_sample_path()
    raw = yaml.safe_load(source.read_text())
    injections = []
    for entry in raw.get("injections", []):
        injections.append(
            Injection(
                id=entry["id"],
                name=entry["name"],
                vector=entry["vector"],
                severity=Severity(entry["severity"]),
                goal=entry["goal"],
                payload=entry["payload"].strip(),
                # Optional, so a corpus written before this field still
                # loads. Absent means "attack", which is the safe default:
                # scoring a benign payload as an attack understates the
                # agent, while the reverse would overstate it.
                is_control=bool(entry.get("is_control", False)),
            )
        )
    return injections


class UnknownVectorError(ValueError):
    """Raised when --vector names something the corpus doesn't contain.

    An error rather than an empty selection, matching the hosted API: a
    typo'd vector would otherwise fire nothing, score 0/0, and exit 0 as
    though the agent had passed.
    """

    def __init__(self, unknown: list[str], available: list[str]) -> None:
        self.unknown = unknown
        self.available = available
        super().__init__(
            f"Unknown vector(s): {', '.join(sorted(unknown))}. "
            f"This corpus has: {', '.join(sorted(available))}."
        )


def select_suite(
    injections: list[Injection],
    suite_mode: SuiteMode,
    vectors: list[str] | None = None,
) -> list[Injection]:
    """Mirrors the hosted API's suite_selection.select_suite_injections.

    The vector filter narrows the corpus *before* the suite cap, so
    "quick, one vector" means the top payloads of that vector rather than
    the top payloads overall filtered down to almost nothing. Keeping the
    two implementations in step matters: a self-hosted run and a hosted
    run of the same suite should cover the same ground.
    """
    candidates = injections
    if vectors:
        available = {i.vector for i in injections}
        unknown = [v for v in vectors if v not in available]
        if unknown:
            raise UnknownVectorError(unknown, sorted(available))
        wanted = set(vectors)
        candidates = [i for i in injections if i.vector in wanted]

    ordered = sorted(candidates, key=lambda i: (-SEVERITY_RANK[i.severity], i.id))
    if suite_mode != SuiteMode.QUICK:
        return ordered

    # Never a quick suite without a clean control. A control is benign, so
    # it is LOW, so severity ranking alone excluded it from the top five —
    # and an agent that refused all five attacks scored 100. It takes the
    # last slot rather than extending the suite, so "quick" stays five
    # payloads and two runs of it remain comparable.
    controls = [i for i in ordered if i.is_control]
    if not controls:
        return ordered[:QUICK_SUITE_SIZE]
    attacks = [i for i in ordered if not i.is_control]
    chosen = attacks[: QUICK_SUITE_SIZE - 1] + controls[:1]
    return sorted(chosen, key=lambda i: (-SEVERITY_RANK[i.severity], i.id))


def has_control(injections: list[Injection]) -> bool:
    """Whether a selection can detect over-blocking at all."""
    return any(i.is_control for i in injections)
