fire-planner/fire_planner/db.py

import os
from datetime import date, datetime
from decimal import Decimal
from typing import Any

from sqlalchemy import JSON, TIMESTAMP, Boolean, Date, Integer, Numeric, String, func, text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

SCHEMA_NAME = "fire_planner"


class Base(DeclarativeBase):
    pass


# JSONB on Postgres, plain JSON on SQLite — tests use SQLite, prod uses Postgres.
JSON_TYPE = JSONB().with_variant(JSON(), "sqlite")


class AccountSnapshot(Base):
    """Daily NW per account from Wealthfolio (filled by ingest).

    `external_id` is `wealthfolio:{account_id}:{date}` so re-runs on the same
    day are idempotent — Wealthfolio keeps one snapshot per account per day.
    """
    __tablename__ = "account_snapshot"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    external_id: Mapped[str] = mapped_column(String, unique=True, nullable=False)
    snapshot_date: Mapped[date] = mapped_column(Date, nullable=False, index=True)
    account_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    account_name: Mapped[str] = mapped_column(String, nullable=False)
    account_type: Mapped[str] = mapped_column(String, nullable=False)
    currency: Mapped[str] = mapped_column(String(3), nullable=False, server_default="GBP")
    market_value: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False)
    market_value_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False)
    cost_basis_gbp: Mapped[Decimal | None] = mapped_column(Numeric(14, 2), nullable=True)
    raw_extraction: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True)
    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


class Scenario(Base):
    """A simulation scenario.

    Two kinds:
    - `kind='cartesian'` — auto-generated from `scenarios.py` Cartesian
      product; rebuilt every recompute, upserted on `external_id`.
    - `kind='user'` — user-defined (named, optionally cloned from a base);
      survives recomputes; `parent_scenario_id` points at the source if any.
    """
    __tablename__ = "scenario"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    external_id: Mapped[str] = mapped_column(String, unique=True, nullable=False)
    kind: Mapped[str] = mapped_column(String(16),
                                      nullable=False,
                                      server_default=text("'cartesian'"))
    name: Mapped[str | None] = mapped_column(String, nullable=True)
    description: Mapped[str | None] = mapped_column(String, nullable=True)
    parent_scenario_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
    jurisdiction: Mapped[str] = mapped_column(String(32), nullable=False, index=True)
    strategy: Mapped[str] = mapped_column(String(32), nullable=False, index=True)
    leave_uk_year: Mapped[int] = mapped_column(Integer, nullable=False)
    glide_path: Mapped[str] = mapped_column(String(32), nullable=False)
    spending_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    horizon_years: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("60"))
    nw_seed_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False)
    savings_per_year_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2),
                                                          nullable=False,
                                                          server_default=text("0"))
    config_json: Mapped[dict[str, Any]] = mapped_column(JSON_TYPE, nullable=False)
    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


class McRun(Base):
    """One MC execution per (scenario, run_at). Stores execution metadata +
    summary statistics — enough to populate a Grafana cell without touching
    the per-path tables."""
    __tablename__ = "mc_run"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    run_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                             nullable=False,
                                             server_default=func.now())
    n_paths: Mapped[int] = mapped_column(Integer, nullable=False)
    seed: Mapped[int] = mapped_column(Integer, nullable=False)
    success_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False)
    p10_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p50_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p90_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    median_lifetime_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False)
    median_years_to_ruin: Mapped[Decimal | None] = mapped_column(Numeric(6, 2), nullable=True)
    elapsed_seconds: Mapped[Decimal] = mapped_column(Numeric(8, 3), nullable=False)
    sequence_risk_correlation: Mapped[Decimal | None] = mapped_column(Numeric(6, 4), nullable=True)
    extra: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True)


class McPath(Base):
    """Sparse per-path storage: top decile, bottom decile, and median paths
    fully stored — enough for a fan chart, not 10k×60 ≈ 600k rows."""
    __tablename__ = "mc_path"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    path_idx: Mapped[int] = mapped_column(Integer, nullable=False)
    bucket: Mapped[str] = mapped_column(String(16), nullable=False)
    year_idx: Mapped[int] = mapped_column(Integer, nullable=False)
    portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    withdrawal_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    tax_paid_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    real_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)


class ProjectionYearly(Base):
    """Deterministic point projection per scenario — per-year point estimates
    that drive fan charts and the per-year Grafana table. One row per
    (scenario, year)."""
    __tablename__ = "projection_yearly"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    year_idx: Mapped[int] = mapped_column(Integer, nullable=False)
    p10_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p25_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p50_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p75_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p90_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p50_withdrawal_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    p50_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    survival_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False)


class ScenarioSummary(Base):
    """Denormalised fast-read for Grafana — one row per (scenario, latest run)."""
    __tablename__ = "scenario_summary"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    scenario_id: Mapped[int] = mapped_column(Integer, unique=True, nullable=False)
    mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False)
    jurisdiction: Mapped[str] = mapped_column(String(32), nullable=False, index=True)
    strategy: Mapped[str] = mapped_column(String(32), nullable=False, index=True)
    leave_uk_year: Mapped[int] = mapped_column(Integer, nullable=False)
    glide_path: Mapped[str] = mapped_column(String(32), nullable=False)
    spending_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False)
    success_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False)
    p10_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p50_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    p90_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False)
    median_lifetime_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False)
    median_years_to_ruin: Mapped[Decimal | None] = mapped_column(Numeric(6, 2), nullable=True)
    updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


class LifeEvent(Base):
    """A timed event in a user's plan: retirement, kid born, mortgage payoff,
    sabbatical, etc. Attached to a scenario.

    `year_start` and `year_end` are offsets from the scenario start year
    (year 0 = today). For one-time events, leave `year_end` = `year_start`.
    `delta_gbp_per_year` is the annual cashflow change while the event is
    active (negative = expense, positive = income; 0 for events that just
    mark a milestone like "retire").

    Free-form `payload` carries event-kind-specific config that the
    simulator hasn't yet learned to consume — graceful forward-compat.
    """
    __tablename__ = "life_event"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    kind: Mapped[str] = mapped_column(String(32), nullable=False)
    name: Mapped[str] = mapped_column(String, nullable=False)
    year_start: Mapped[int] = mapped_column(Integer, nullable=False)
    year_end: Mapped[int | None] = mapped_column(Integer, nullable=True)
    delta_gbp_per_year: Mapped[Decimal] = mapped_column(Numeric(12, 2),
                                                        nullable=False,
                                                        server_default=text("0"))
    one_time_amount_gbp: Mapped[Decimal | None] = mapped_column(Numeric(14, 2), nullable=True)
    # Spending category for flex-spending classification:
    #   essential       — never trimmed by flex rules (mortgage, food, kids)
    #   discretionary   — trimmed when portfolio drops vs ATH (travel, dining)
    #   not_spending    — informational only (a milestone marker that doesn't
    #                     change cashflow, e.g. a kid graduating)
    # Default is `essential` so existing rows keep their full spending impact.
    category: Mapped[str] = mapped_column(String(16),
                                          nullable=False,
                                          server_default=text("'essential'"))
    enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true"))
    payload: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True)
    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


class IncomeStream(Base):
    """A typed, recurring source of income — first-class income object.

    Modelled as a per-scenario row so a user can stack salary, dividends,
    rental, pensions, RSUs, etc. The simulator routes the after-tax
    amount through the jurisdiction's tax engine using `tax_treatment`
    as the bucket hint (income / dividend / cgt / tax_free).

    `start_year` / `end_year` are offsets from the scenario start year.
    `growth_pct` is real growth; the simulator applies it geometrically.
    """
    __tablename__ = "income_stream"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    kind: Mapped[str] = mapped_column(String(32), nullable=False)
    name: Mapped[str] = mapped_column(String, nullable=False)
    start_year: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("0"))
    end_year: Mapped[int | None] = mapped_column(Integer, nullable=True)
    amount_gbp_per_year: Mapped[Decimal] = mapped_column(Numeric(12, 2),
                                                         nullable=False,
                                                         server_default=text("0"))
    growth_pct: Mapped[Decimal] = mapped_column(Numeric(6, 4),
                                                nullable=False,
                                                server_default=text("0"))
    tax_treatment: Mapped[str] = mapped_column(String(16),
                                               nullable=False,
                                               server_default=text("'income'"))
    enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true"))
    payload: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True)
    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


class RetirementGoal(Base):
    """A user-defined success criterion for a scenario.

    Examples:
    - target_nw: "have ≥£2M real GBP at year 25" → kind=target_nw,
      target_amount_gbp=2_000_000, target_year=25, comparator='>='
    - never_run_out: "never run out before age 95" → kind=never_run_out,
      target_year=65 (years from start), no amount
    - inheritance: "leave ≥£500k to heirs" → kind=inheritance,
      target_amount_gbp=500_000, target_year=horizon, comparator='>='

    `success_threshold` is the probability bar that counts as "passing"
    (e.g. 0.95 = 95% of MC paths must satisfy the comparator).
    """
    __tablename__ = "retirement_goal"
    __table_args__ = {"schema": SCHEMA_NAME}  # noqa: RUF012

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
    kind: Mapped[str] = mapped_column(String(32), nullable=False)
    name: Mapped[str] = mapped_column(String, nullable=False)
    target_amount_gbp: Mapped[Decimal | None] = mapped_column(Numeric(16, 2), nullable=True)
    target_year: Mapped[int | None] = mapped_column(Integer, nullable=True)
    comparator: Mapped[str] = mapped_column(String(4), nullable=False, server_default=text("'>='"))
    success_threshold: Mapped[Decimal] = mapped_column(Numeric(4, 3),
                                                       nullable=False,
                                                       server_default=text("0.95"))
    enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true"))
    payload: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True)
    created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True),
                                                 nullable=False,
                                                 server_default=func.now())


def create_engine_from_env() -> AsyncEngine:
    url = os.environ["DB_CONNECTION_STRING"]
    return create_async_engine(url, pool_pre_ping=True)


def make_session_factory(engine: AsyncEngine) -> async_sessionmaker[Any]:
    return async_sessionmaker(engine, expire_on_commit=False)