"""Balance sheets with exact integer accounting and explicit reservations.

Invariant enforced by `Ledger.check_conservation()` after every fill:
    sum(cash of all agents) + exchange_fees_collected == initial_cash_total + dividends_paid_total
    sum(positions of all agents per asset) == initial_shares_total per asset   (the exchange never creates shares)
"""
from __future__ import annotations

from dataclasses import dataclass, field


@dataclass(slots=True)
class Account:
    agent_id: int
    cash: int                                   # cents
    positions: dict[int, int]                   # asset -> lots (negative if short)
    reserved_cash: int = 0                      # cents locked by open bids
    reserved_qty: dict[int, int] = field(default_factory=dict)   # asset -> lots locked by open asks
    frozen: bool = False                        # bankrupt / liquidated: no new orders
    realized_fees: int = 0                      # cents paid in fees (cumulative)
    dividends_received: int = 0
    n_orders: int = 0
    n_fills: int = 0
    n_rejects: int = 0
    n_cancels: int = 0
    traded_qty: int = 0
    # risk settings (per agent)
    leverage_cap: float = 1.0                   # max gross exposure / equity (1.0 = no borrowing)
    short_allowed: bool = False
    maintenance_margin: float = 0.25            # equity / gross exposure must stay above this if leveraged
    initial_cash: int = 0
    initial_positions: dict[int, int] = field(default_factory=dict)

    @property
    def free_cash(self) -> int:
        return self.cash - self.reserved_cash

    def free_qty(self, asset: int) -> int:
        return self.positions.get(asset, 0) - self.reserved_qty.get(asset, 0)

    def equity(self, marks: dict[int, float]) -> float:
        """Cash + marked positions (marks in cents per lot)."""
        return self.cash + sum(q * marks.get(a, 0.0) for a, q in self.positions.items())

    def gross_exposure(self, marks: dict[int, float]) -> float:
        return sum(abs(q) * marks.get(a, 0.0) for a, q in self.positions.items())

    def net_worth_units(self, marks: dict[int, float]) -> float:
        return self.equity(marks) / 100.0


class Ledger:
    """All accounts plus the exchange's own fee account; owns the conservation checks."""

    def __init__(self) -> None:
        self.accounts: dict[int, Account] = {}
        self.exchange_fees: int = 0
        self.dividends_paid: int = 0
        self.initial_cash_total: int = 0
        self.initial_shares: dict[int, int] = {}

    def add(self, acct: Account) -> None:
        acct.initial_cash = acct.cash
        acct.initial_positions = dict(acct.positions)
        self.accounts[acct.agent_id] = acct
        self.initial_cash_total += acct.cash
        for a, q in acct.positions.items():
            self.initial_shares[a] = self.initial_shares.get(a, 0) + q

    def __getitem__(self, agent_id: int) -> Account:
        return self.accounts[agent_id]

    def apply_fill(self, buyer: int, seller: int, asset: int, price_cents: int, qty: int,
                   buyer_fee: int, seller_fee: int, buyer_reserved: int, seller_reserved_qty: int) -> None:
        """Move cash and lots; release the reservations attached to the two orders for this quantity."""
        notional = price_cents * qty
        b, s = self.accounts[buyer], self.accounts[seller]
        b.cash -= notional + buyer_fee
        s.cash += notional - seller_fee
        b.positions[asset] = b.positions.get(asset, 0) + qty
        s.positions[asset] = s.positions.get(asset, 0) - qty
        b.reserved_cash -= buyer_reserved
        s.reserved_qty[asset] = s.reserved_qty.get(asset, 0) - seller_reserved_qty
        b.realized_fees += buyer_fee
        s.realized_fees += seller_fee
        b.n_fills += 1
        s.n_fills += 1
        b.traded_qty += qty
        s.traded_qty += qty
        self.exchange_fees += buyer_fee + seller_fee
        assert b.reserved_cash >= 0, "negative cash reservation"
        assert s.reserved_qty[asset] >= 0 or s.short_allowed, "negative qty reservation"

    def pay_dividend(self, asset: int, per_lot_cents: int) -> int:
        """Pays (or charges, for shorts) a cash dividend; returns the net cash injected."""
        total = 0
        for acct in self.accounts.values():
            q = acct.positions.get(asset, 0)
            if q:
                d = per_lot_cents * q
                acct.cash += d
                acct.dividends_received += d
                total += d
        self.dividends_paid += total
        return total

    def check_conservation(self) -> None:
        cash = sum(a.cash for a in self.accounts.values()) + self.exchange_fees
        assert cash == self.initial_cash_total + self.dividends_paid, (
            f"cash not conserved: {cash} != {self.initial_cash_total} + {self.dividends_paid}")
        for asset, n0 in self.initial_shares.items():
            n = sum(a.positions.get(asset, 0) for a in self.accounts.values())
            assert n == n0, f"shares not conserved for asset {asset}: {n} != {n0}"
        for a in self.accounts.values():
            assert a.reserved_cash >= 0 and a.reserved_cash <= max(a.cash, 0) or a.leverage_cap > 1.0, (
                f"agent {a.agent_id}: reserved_cash {a.reserved_cash} > cash {a.cash}")
            for asset, rq in a.reserved_qty.items():
                assert rq >= 0, f"agent {a.agent_id}: negative reserved qty"
                assert a.short_allowed or rq <= a.positions.get(asset, 0), (
                    f"agent {a.agent_id}: reserved {rq} > held {a.positions.get(asset, 0)} on asset {asset}")