MyMidas/backend/app/schemas/pension.py

import uuid
from datetime import date as DateType, datetime
from decimal import Decimal
from typing import Literal

from pydantic import BaseModel, Field, model_validator

PensionType = Literal["workplace_dc", "workplace_db", "sipp", "lisa"]
ReliefType = Literal["relief_at_source", "net_pay", "salary_sacrifice", "none"]


class PensionMetadataCreate(BaseModel):
    pension_type: PensionType
    provider_name: str | None = None
    scheme_name: str | None = None
    member_reference: str | None = None
    dob: DateType | None = None
    target_retirement_age: int | None = Field(default=None, ge=55, le=90)
    assumed_growth_rate: Decimal | None = Field(default=None, ge=0, le=1)


class PensionMetadataUpdate(BaseModel):
    pension_type: PensionType | None = None
    provider_name: str | None = None
    scheme_name: str | None = None
    member_reference: str | None = None
    dob: DateType | None = None
    target_retirement_age: int | None = Field(default=None, ge=55, le=90)
    assumed_growth_rate: Decimal | None = Field(default=None, ge=0, le=1)


class PensionMetadataResponse(BaseModel):
    id: uuid.UUID
    account_id: uuid.UUID
    pension_type: PensionType
    provider_name: str | None
    scheme_name: str | None
    member_reference: str | None
    dob: DateType | None
    target_retirement_age: int | None
    assumed_growth_rate: Decimal | None
    created_at: datetime
    updated_at: datetime

    model_config = {"from_attributes": True}


class PensionContributionCreate(BaseModel):
    contribution_date: DateType
    member_amount: Decimal = Field(..., ge=0)
    employer_amount: Decimal = Field(default=Decimal("0"), ge=0)
    relief_type: ReliefType
    notes: str | None = None

    @model_validator(mode="after")
    def check_amounts(self) -> "PensionContributionCreate":
        if self.member_amount == 0 and self.employer_amount == 0:
            raise ValueError("At least one of member_amount or employer_amount must be greater than 0")
        return self


class PensionContributionUpdate(BaseModel):
    contribution_date: DateType | None = None
    member_amount: Decimal | None = Field(default=None, ge=0)
    employer_amount: Decimal | None = Field(default=None, ge=0)
    relief_type: ReliefType | None = None
    notes: str | None = None


class PensionContributionResponse(BaseModel):
    id: uuid.UUID
    pension_id: uuid.UUID
    contribution_date: DateType
    tax_year: int
    member_amount: Decimal
    employer_amount: Decimal
    relief_type: ReliefType
    gross_amount: Decimal
    relief_amount: Decimal
    notes: str | None
    created_at: datetime

    model_config = {"from_attributes": True}


class YtdSummary(BaseModel):
    tax_year: int
    member_total: Decimal
    employer_total: Decimal
    gross_total: Decimal
    relief_total: Decimal
    contribution_count: int


class PensionAccountResponse(BaseModel):
    """Account row joined with pension metadata and YTD summary."""
    account_id: uuid.UUID
    account_name: str
    current_balance: Decimal
    currency: str
    color: str
    metadata: PensionMetadataResponse | None
    ytd: YtdSummary | None


_FULL_SP_WEEKLY = Decimal("221.80")
_SP_AGE = 67


class StatePensionCreate(BaseModel):
    qualifying_years: int = Field(..., ge=0, le=50)
    checked_date: DateType | None = None


class StatePensionResponse(BaseModel):
    id: uuid.UUID
    qualifying_years: int
    checked_date: DateType | None
    weekly_amount: Decimal
    annual_amount: Decimal
    is_full_pension: bool
    years_to_full: int
    state_pension_age: int

    model_config = {"from_attributes": True}


class ProjectionScenario(BaseModel):
    label: str
    growth_rate: Decimal
    projected_pot: Decimal
    annual_drawdown_4pct: Decimal
    annual_drawdown_3pct: Decimal


class ChartDataPoint(BaseModel):
    year: int
    pot_2pct: Decimal
    pot_5pct: Decimal
    pot_8pct: Decimal


class RetirementProjection(BaseModel):
    account_id: uuid.UUID
    account_name: str
    current_balance: Decimal
    years_to_retirement: int
    target_retirement_age: int
    scenarios: list[ProjectionScenario]
    state_pension_annual: Decimal | None
    state_pension_age: int
    chart_data: list[ChartDataPoint]


class LisaTaxYearBreakdown(BaseModel):
    tax_year: int
    contributions: Decimal
    bonus_expected: Decimal
    limit_remaining: Decimal
    limit_used_pct: Decimal


class LisaSummary(BaseModel):
    account_id: uuid.UUID
    account_name: str
    tax_year_breakdown: list[LisaTaxYearBreakdown]
    current_year_contributions: Decimal
    current_year_bonus_expected: Decimal
    current_year_limit_remaining: Decimal
    total_contributions: Decimal
    total_bonus_expected: Decimal
    account_opened_date: datetime
    withdrawal_penalty_amount: Decimal
    withdrawal_penalty_pct: Decimal
    penalty_warning: bool


class CarryForwardYear(BaseModel):
    tax_year: int
    standard_allowance: Decimal
    contributions: Decimal
    unused: Decimal


class AllowanceSummary(BaseModel):
    tax_year: int
    standard_allowance: Decimal
    contributions_total: Decimal
    remaining: Decimal
    carry_forward: list[CarryForwardYear]
    carry_forward_total: Decimal
    total_available: Decimal
    relief_ras_total: Decimal
    relief_higher_rate_claimable: Decimal
    relief_additional_rate_claimable: Decimal