from typing import List, Optional, Self

from abacus_common_logic.models.base import BaseModel, db


class TaxWithholdingOverride(BaseModel):
    """Tax Withholding Override model"""

    __tablename__ = 'tax_withholding_override'
    tax_withholding_override_id = db.Column(
        db.Integer, primary_key=True, autoincrement=True
    )
    account_payee_id = db.Column(
        db.Integer,
        db.ForeignKey('account_payee.account_payee_id'),
        nullable=False,
        unique=True,
    )
    rate_override = db.Column(db.Numeric(5, 2), nullable=True, default=None)
    certificate_expiration_date = db.Column(db.Date, nullable=True, default=None)
    message = db.Column(db.Text, nullable=True, default=None)
    # timestamps are inherited from BaseModel

    @classmethod
    def get_by_account_payee_id(cls, account_payee_id: int) -> Optional[Self]:
        """Get TaxWithholdingOverride by account_payee_id."""
        return cls.query.filter(cls.account_payee_id == account_payee_id).first()

    @classmethod
    def get_by_account_payee_ids(cls, account_payee_ids: List[int]) -> List[Self]:
        """Get TaxWithholdingOverride by account_payee_ids."""
        return cls.query.filter(cls.account_payee_id.in_(account_payee_ids)).all()