case_db/models.py

"""Criminal case data models."""

from dataclasses import dataclass, field
from typing import Literal


@dataclass
class Evidence:
    """A piece of evidence."""

    evidence_id: str
    type: str  # "physical", "testimonial", "documentary", "forensic"
    description: str
    strength_prosecution: float  # 0.0-1.0
    strength_defense: float  # 0.0-1.0
    contestable: bool = False
    contest_reason: str | None = None

    def get_summary(self) -> str:
        """Get a brief summary of the evidence."""
        return f"[{self.type.upper()}] {self.description}"


@dataclass
class Witness:
    """A witness in the case."""

    witness_id: str
    name: str
    role: str  # "eyewitness", "expert", "character", etc.
    testimony_summary: str
    credibility_issues: list[str] = field(default_factory=list)
    side: Literal["prosecution", "defense", "neutral"] = "neutral"

    def get_summary(self) -> str:
        """Get a brief summary of the witness."""
        return f"{self.name} ({self.role}): {self.testimony_summary}"


@dataclass
class Defendant:
    """The defendant in the case."""

    name: str
    age: int | None = None
    occupation: str | None = None
    background: str = ""
    prior_record: list[str] = field(default_factory=list)


@dataclass
class CriminalCase:
    """A criminal case for deliberation."""

    case_id: str
    title: str
    summary: str  # 2-3 paragraph overview

    # Charges
    charges: list[str] = field(default_factory=list)

    # Evidence
    evidence: list[Evidence] = field(default_factory=list)

    # Witnesses
    witnesses: list[Witness] = field(default_factory=list)

    # Arguments
    prosecution_arguments: list[str] = field(default_factory=list)
    defense_arguments: list[str] = field(default_factory=list)

    # Defendant
    defendant: Defendant | None = None

    # Metadata
    difficulty: Literal["clear_guilty", "clear_innocent", "ambiguous"] = "ambiguous"
    themes: list[str] = field(default_factory=list)

    # For display
    year: int = 2024
    jurisdiction: str = "United States"

    def get_evidence_summary(self) -> str:
        """Get formatted summary of all evidence."""
        if not self.evidence:
            return "No evidence available."
        return "\n".join(f"- {e.get_summary()}" for e in self.evidence)

    def get_witness_summary(self) -> str:
        """Get formatted summary of all witnesses."""
        if not self.witnesses:
            return "No witnesses available."
        return "\n".join(f"- {w.get_summary()}" for w in self.witnesses)

    def get_charges_text(self) -> str:
        """Get formatted list of charges."""
        if not self.charges:
            return "No charges specified."
        return ", ".join(self.charges)

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "case_id": self.case_id,
            "title": self.title,
            "summary": self.summary,
            "charges": self.charges,
            "difficulty": self.difficulty,
            "themes": self.themes,
            "year": self.year,
            "jurisdiction": self.jurisdiction,
        }