Stonks/overview/brokerages/random.py

import random
from datetime import datetime, timedelta

from overview.brokerage_types import (
    BrokerageAccountPayload,
    BrokerageAssetPayload,
    BrokerageHoldingPayload,
    BrokerageSnapshot,
)

from .base import BaseBrokerageAdapter


class RandomBrokerageAdapter(BaseBrokerageAdapter):
    provider = "random-demo"

    DEFAULT_TICKERS = [
        "AAPL",
        "MSFT",
        "GOOGL",
        "AMZN",
        "NVDA",
        "TSLA",
        "META",
        "BRK.B",
        "JPM",
        "V",
        "NFLX",
        "AMD",
    ]

    def fetch_snapshot(self) -> BrokerageSnapshot:
        generator = random.Random(self.connection.get_setting("seed"))
        account_count = max(1, int(self.connection.get_setting("account_count", 2)))
        asset_count = max(1, int(self.connection.get_setting("asset_count", 8)))
        holdings_per_account = max(
            1,
            int(self.connection.get_setting("holdings_per_account", 4)),
        )
        reference_timestamp = self.connection.get_setting(
            "reference_timestamp",
            "2025-01-01T00:00:00+00:00",
        )
        reference_time = datetime.fromisoformat(
            str(reference_timestamp).replace("Z", "+00:00")
        )
        tickers = self.connection.get_setting("tickers", self.DEFAULT_TICKERS)

        if not isinstance(tickers, list) or not tickers:
            tickers = self.DEFAULT_TICKERS

        assets = self._build_assets(generator, asset_count, tickers)
        accounts = self._build_accounts(generator, account_count)
        holdings = self._build_holdings(
            generator,
            accounts=accounts,
            assets=assets,
            holdings_per_account=holdings_per_account,
            reference_time=reference_time,
        )

        return BrokerageSnapshot(
            accounts=accounts,
            assets=assets,
            holdings=holdings,
        )

    def _build_accounts(
        self,
        generator: random.Random,
        account_count: int,
    ) -> list[BrokerageAccountPayload]:
        account_labels = [
            "Individual",
            "Joint",
            "Retirement",
            "Roth IRA",
            "Trust",
            "Education",
        ]
        accounts = []

        for index in range(account_count):
            suffix = generator.randint(1000, 9999)
            name = account_labels[index % len(account_labels)]
            accounts.append(
                BrokerageAccountPayload(
                    external_id=f"acct-{index + 1}",
                    name=f"{name} Account {index + 1}",
                    number=f"{suffix}{generator.randint(10, 99)}",
                    balance=0,
                )
            )

        return accounts

    def _build_assets(
        self,
        generator: random.Random,
        asset_count: int,
        tickers: list[str],
    ) -> list[BrokerageAssetPayload]:
        assets = []

        for index in range(asset_count):
            ticker = tickers[index % len(tickers)]
            base_price = round(generator.uniform(25, 750), 2)
            assets.append(
                BrokerageAssetPayload(
                    external_id=f"asset-{index + 1}",
                    ticker=ticker,
                    price=base_price,
                )
            )

        return assets

    def _build_holdings(
        self,
        generator: random.Random,
        accounts: list[BrokerageAccountPayload],
        assets: list[BrokerageAssetPayload],
        holdings_per_account: int,
        reference_time: datetime,
    ) -> list[BrokerageHoldingPayload]:
        holdings = []
        asset_lookup = {asset.external_id: asset for asset in assets}
        days_back = int(self.connection.get_setting("max_purchase_age_days", 3650))
        next_holding_id = 1
        refreshed_accounts = []

        for account in accounts:
            selected_assets = generator.sample(
                assets,
                k=min(holdings_per_account, len(assets)),
            )
            balance = 0.0

            for asset in selected_assets:
                amount = round(generator.uniform(1, 50), 4)
                current_value = amount * asset_lookup[asset.external_id].price
                price_offset = generator.uniform(0.7, 1.15)
                cost_basis = round(asset.price * price_offset, 2)
                purchase_date = reference_time - timedelta(
                    days=generator.randint(1, max(1, days_back))
                )

                holdings.append(
                    BrokerageHoldingPayload(
                        external_id=f"holding-{next_holding_id}",
                        account_external_id=account.external_id,
                        asset_external_id=asset.external_id,
                        ticker=asset.ticker,
                        amount=amount,
                        cost_basis=cost_basis,
                        purchase_date=purchase_date,
                    )
                )
                next_holding_id += 1
                balance += current_value

            refreshed_accounts.append(
                BrokerageAccountPayload(
                    external_id=account.external_id,
                    name=account.name,
                    number=account.number,
                    balance=int(round(balance)),
                )
            )

        accounts[:] = refreshed_accounts
        return holdings