import pydantic as p
from server.utils.bq import BQClient, Statement, Table


def get_yy(yy: str | datetime.date) -> str:
    """Return the two-digit year for the given date."""
    return yy.strftime("%y") if isinstance(yy, datetime.date) else yy[-2:]


class Committee(p.BaseModel, frozen=True):
    """Represents an FEC committee."""

    id: str
    name: str | None
    party: str | None
    candidate_id: str | None

    @property
    def adjusted_party(self) -> str | None:
        """
        Return the FEC reported party, except in a few key cases,
        where we know better. (For instance, ActBlue has a "party"
        of NULL in the FEC's 2020 dataset. But we know it's mostly
        DEM.)
        """
        if self.id in KNOWN_DEM_COMMITTEE_IDS:
            return Party.DEMOCRAT
        return self.party


class CommitteeTable(Table[Committee]):
    """Tools for querying the BigQuery committee master file."""

    def __init__(self, client: BQClient, year: str | datetime.date):
        super().__init__(client, f"bigquery-public-data.fec.cm{get_yy(year)}")

    def get_model_instance(self, row: t.Any) -> Committee:
        """Create a committee from a row of the committee master file."""
        return Committee(
            id=row.cmte_id,
            name=row.cmte_nm,
            party=row.cmte_pty_affiliation,
            candidate_id=row.cand_id,
        )

    def for_name_stmt(self, name: str) -> Statement:
        """Return a select statement for committees matching the given criteria."""
        return self.all_stmt().where("cmte_nm", "LIKE", f"%{name.upper()}%")

    def for_name(self, name: str) -> t.Iterable[Committee]:
        """Return a query for committees matching the given criteria."""
        return self.execute(self.for_name_stmt(name))