"""Juror persona definitions with graph traversal configuration.

Each persona defines which graph nodes/edges/chunks are relevant to their
deliberation perspective, ensuring they see only information appropriate to
their role. The RULED_INADMISSIBLE wall is enforced for all personas.
"""

from __future__ import annotations

from aucourt_ingest.models import JurorPersona

# Default token budget per juror context assembly
DEFAULT_TOKEN_BUDGET = 4000


PERSONAS: dict[str, JurorPersona] = {
    "nurse": JurorPersona(
        name="nurse",
        anchor_nodes=["Witness[role=expert]", "Exhibit[category=medical]"],
        edge_types=["GAVE_TESTIMONY", "DESCRIBED_IN", "CORROBORATES"],
        chunk_types=["testimony", "exhibit"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
    "accountant": JurorPersona(
        name="accountant",
        anchor_nodes=["Exhibit[category=financial]", "Witness[role=expert]"],
        edge_types=["DESCRIBED_IN", "HAS_RULING"],
        chunk_types=["exhibit", "ruling"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
    "skeptic": JurorPersona(
        name="skeptic",
        anchor_nodes=["Timeline", "Witness"],
        edge_types=["CONTRADICTS", "HAS_RULING", "GAVE_TESTIMONY"],
        chunk_types=["testimony", "ruling"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
    "ex_cop": JurorPersona(
        name="ex_cop",
        anchor_nodes=["Ruling", "Exhibit[category=forensic]"],
        edge_types=["HAS_RULING", "DESCRIBED_IN", "CORROBORATES"],
        chunk_types=["ruling", "exhibit", "judgment"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
    "empath": JurorPersona(
        name="empath",
        anchor_nodes=["Witness[role=prosecution]", "Witness[role=accused]"],
        edge_types=["GAVE_TESTIMONY", "CORROBORATES", "CONTRADICTS"],
        chunk_types=["testimony", "closing"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
    "foreman": JurorPersona(
        name="foreman",
        anchor_nodes=["Charge", "Ruling", "Judge"],
        edge_types=["CHARGED_WITH", "HAS_RULING", "HEARD_BY"],
        chunk_types=["opening", "judgment", "sentence"],
        exclude_edges=["RULED_INADMISSIBLE"],
    ),
}


def get_persona(name: str) -> JurorPersona:
    """Look up a persona by name. Raises KeyError if not found."""
    return PERSONAS[name]


def all_persona_names() -> list[str]:
    """Return list of all defined persona names."""
    return list(PERSONAS.keys())