P&L Accounting for AI Agents: Tracking Every Satoshi

"""
Multi-chain, multi-service P&L tracker for Purple Flea AI agents.
Supports FIFO, LIFO, HIFO cost basis methods.
Aggregates: Casino, Trading (Spot + Perp), Wallet, Escrow, Faucet.
"""

import asyncio
import csv
import json
import logging
from collections import defaultdict, deque
from dataclasses import dataclass, field, asdict
from decimal import Decimal
from enum import Enum
from typing import Dict, List, Optional, Tuple
from datetime import datetime, timezone
import aiohttp

log = logging.getLogger('pnl_tracker')

PURPLEFLEA_BASE = "https://purpleflea.com"
API_KEY = "your_api_key_here"


# ─── Enums & Constants ───────────────────────────────────────────────────────

class CostBasisMethod(Enum):
    FIFO = "fifo"
    LIFO = "lifo"
    HIFO = "hifo"


class EventType(Enum):
    BUY            = "buy"
    SELL           = "sell"
    CASINO_WIN     = "casino_win"
    CASINO_LOSS    = "casino_loss"
    PERP_OPEN      = "perp_open"
    PERP_CLOSE     = "perp_close"
    FUNDING        = "funding"
    BRIDGE_OUT     = "bridge_out"
    BRIDGE_IN      = "bridge_in"
    FEE            = "fee"
    ESCROW_RECEIVE = "escrow_receive"
    ESCROW_SEND    = "escrow_send"
    FAUCET_CLAIM   = "faucet_claim"
    REFERRAL_EARN  = "referral_earn"


# ─── Data Structures ─────────────────────────────────────────────────────────

@dataclass
class TaxLot:
    """A single acquired unit for cost basis tracking."""
    quantity: Decimal
    cost_per_unit: Decimal
    timestamp: float
    source: str
    lot_id: str


@dataclass
class PnLEvent:
    event_type: EventType
    asset: str
    chain: str
    quantity: Decimal
    price_usd: Decimal
    fee_usd: Decimal
    timestamp: float
    tx_hash: Optional[str] = None
    metadata: Dict = field(default_factory=dict)

    @property
    def gross_usd(self) -> Decimal:
        return self.quantity * self.price_usd

    @property
    def net_usd(self) -> Decimal:
        return self.gross_usd - self.fee_usd


@dataclass
class RealizedGain:
    asset: str
    quantity_sold: Decimal
    sale_price: Decimal
    cost_basis: Decimal
    gain_usd: Decimal
    timestamp: float
    method: str
    lot_id: str


@dataclass
class PnLReport:
    period_start: datetime
    period_end: datetime
    realized_pnl: Decimal
    unrealized_pnl: Decimal
    total_fees: Decimal
    casino_pnl: Decimal
    trading_pnl: Decimal
    funding_pnl: Decimal
    escrow_pnl: Decimal
    faucet_income: Decimal
    net_pnl: Decimal
    gain_events: List[RealizedGain]
    fee_breakdown: Dict[str, Decimal]


# ─── Cost Basis Engine ───────────────────────────────────────────────────────

class CostBasisEngine:
    def __init__(self, method: CostBasisMethod = CostBasisMethod.HIFO):
        self.method = method
        # asset → list of TaxLot (ordered by acquisition time)
        self.lots: Dict[str, List[TaxLot]] = defaultdict(list)
        self.gains: List[RealizedGain] = []

    def acquire(self, asset: str, quantity: Decimal, cost_per_unit: Decimal, ts: float, source: str):
        """Record an asset acquisition (buy, faucet claim, etc)."""
        lot_id = f"{asset}_{int(ts*1000)}"
        lot = TaxLot(
            quantity=quantity,
            cost_per_unit=cost_per_unit,
            timestamp=ts,
            source=source,
            lot_id=lot_id,
        )
        self.lots[asset].append(lot)
        self.lots[asset].sort(key=lambda x: x.timestamp)  # keep chronological
        log.debug(f"Acquired {quantity} {asset} at ${cost_per_unit:.4f} (lot {lot_id})")

    def dispose(
        self, asset: str, quantity: Decimal, sale_price: Decimal, ts: float
    ) -> Tuple[Decimal, List[RealizedGain]]:
        """
        Record disposal (sell, swap out). Returns (realized_gain, lots_used).
        Applies configured cost basis method to select lots.
        """
        available = sum(lot.quantity for lot in self.lots[asset])
        if quantity > available:
            log.warning(f"Disposing {quantity} {asset} but only {available} in lots")
            quantity = available

        # Sort lots by method
        ordered = sorted(
            self.lots[asset],
            key=lambda l: (
                l.timestamp if self.method == CostBasisMethod.FIFO
                else -l.timestamp if self.method == CostBasisMethod.LIFO
                else -l.cost_per_unit  # HIFO: highest cost first
            )
        )

        remaining = quantity
        total_cost_basis = Decimal('0')
        new_gains = []

        for lot in ordered:
            if remaining <= 0:
                break

            use_qty = min(remaining, lot.quantity)
            cost_basis_portion = use_qty * lot.cost_per_unit
            gain = use_qty * sale_price - cost_basis_portion

            gain_event = RealizedGain(
                asset=asset,
                quantity_sold=use_qty,
                sale_price=sale_price,
                cost_basis=lot.cost_per_unit,
                gain_usd=gain,
                timestamp=ts,
                method=self.method.value,
                lot_id=lot.lot_id,
            )
            new_gains.append(gain_event)
            self.gains.append(gain_event)
            total_cost_basis += cost_basis_portion

            lot.quantity -= use_qty
            remaining -= use_qty

        # Remove exhausted lots
        self.lots[asset] = [l for l in self.lots[asset] if l.quantity > 0]

        total_gain = quantity * sale_price - total_cost_basis
        return total_gain, new_gains

    def total_cost_basis(self, asset: str) -> Decimal:
        return sum(l.quantity * l.cost_per_unit for l in self.lots[asset])

    def total_quantity(self, asset: str) -> Decimal:
        return sum(l.quantity for l in self.lots[asset])


# ─── Purple Flea API Data Fetcher ─────────────────────────────────────────────

class PurpleFleasDataFetcher:
    def __init__(self, api_key: str):
        self.headers = {"Authorization": f"Bearer {api_key}"}
        self._session: Optional[aiohttp.ClientSession] = None

    async def __aenter__(self):
        self._session = aiohttp.ClientSession(headers=self.headers)
        return self

    async def __aexit__(self, *args):
        await self._session.close()

    async def get_casino_history(self, since: float) -> List[PnLEvent]:
        async with self._session.get(
            f"{PURPLEFLEA_BASE}/casino-api/history",
            params={"since": int(since)}
        ) as r:
            data = await r.json()

        events = []
        for bet in data["bets"]:
            etype = EventType.CASINO_WIN if bet["outcome"] == "win" else EventType.CASINO_LOSS
            pnl = Decimal(str(bet["net_pnl"]))
            events.append(PnLEvent(
                event_type=etype,
                asset="USDC",
                chain=bet.get("chain", "ETH"),
                quantity=abs(pnl),
                price_usd=Decimal("1.00"),
                fee_usd=Decimal(str(bet.get("fee", 0))),
                timestamp=bet["settled_at"],
                tx_hash=bet.get("tx_hash"),
                metadata={"game": bet["game"], "bet_amount": bet["amount"]},
            ))
        return events

    async def get_trading_history(self, since: float) -> List[PnLEvent]:
        async with self._session.get(
            f"{PURPLEFLEA_BASE}/trading-api/history",
            params={"since": int(since)}
        ) as r:
            data = await r.json()

        events = []
        for trade in data["trades"]:
            etype = EventType.BUY if trade["side"] == "buy" else EventType.SELL
            events.append(PnLEvent(
                event_type=etype,
                asset=trade["base_asset"],
                chain=trade.get("chain", "ETH"),
                quantity=Decimal(str(trade["quantity"])),
                price_usd=Decimal(str(trade["price"])),
                fee_usd=Decimal(str(trade["fee_usd"])),
                timestamp=trade["executed_at"],
                tx_hash=trade.get("tx_hash"),
                metadata={"pair": trade["pair"], "type": trade["type"]},
            ))
        return events

    async def get_wallet_history(self, since: float) -> List[PnLEvent]:
        async with self._session.get(
            f"{PURPLEFLEA_BASE}/wallet-api/history",
            params={"since": int(since)}
        ) as r:
            data = await r.json()

        events = []
        for tx in data["transactions"]:
            if tx["type"] == "bridge_out":
                etype = EventType.BRIDGE_OUT
            elif tx["type"] == "bridge_in":
                etype = EventType.BRIDGE_IN
            else:
                continue  # plain transfers are not P&L events

            events.append(PnLEvent(
                event_type=etype,
                asset=tx["asset"],
                chain=tx["chain"],
                quantity=Decimal(str(tx["amount"])),
                price_usd=Decimal(str(tx.get("price_usd_at_time", 0))),
                fee_usd=Decimal(str(tx.get("fee_usd", 0))),
                timestamp=tx["timestamp"],
                tx_hash=tx.get("tx_hash"),
            ))
        return events

    async def get_escrow_history(self, since: float) -> List[PnLEvent]:
        async with self._session.get(
            f"https://escrow.purpleflea.com/api/history",
            headers=self.headers,
            params={"since": int(since)}
        ) as r:
            data = await r.json()

        events = []
        for e in data["escrows"]:
            etype = EventType.ESCROW_RECEIVE if e["role"] == "recipient" else EventType.ESCROW_SEND
            # Escrow fee: 1% platform fee + 15% of that goes to referrer
            escrow_fee = Decimal(str(e["amount"])) * Decimal("0.01")
            events.append(PnLEvent(
                event_type=etype,
                asset=e["asset"],
                chain=e["chain"],
                quantity=Decimal(str(e["amount"])),
                price_usd=Decimal(str(e.get("price_usd_at_time", 1))),
                fee_usd=escrow_fee,
                timestamp=e["released_at"],
                tx_hash=e.get("tx_hash"),
                metadata={"counterparty": e.get("counterparty")},
            ))
            # Track referral earnings separately if any
            if e.get("referral_earned"):
                events.append(PnLEvent(
                    event_type=EventType.REFERRAL_EARN,
                    asset="USDC",
                    chain=e["chain"],
                    quantity=Decimal(str(e["referral_earned"])),
                    price_usd=Decimal("1.00"),
                    fee_usd=Decimal("0"),
                    timestamp=e["released_at"],
                ))
        return events

    async def get_current_prices(self, assets: List[str]) -> Dict[str, Decimal]:
        async with self._session.get(
            f"{PURPLEFLEA_BASE}/wallet-api/prices",
            params={"assets": ",".join(assets)}
        ) as r:
            data = await r.json()
        return {a: Decimal(str(p)) for a, p in data["prices"].items()}


# ─── Main P&L Tracker ────────────────────────────────────────────────────────

class AgentPnLTracker:
    def __init__(self, api_key: str, method: CostBasisMethod = CostBasisMethod.HIFO):
        self.engine  = CostBasisEngine(method)
        self.fetcher = PurpleFleasDataFetcher(api_key)
        self.all_events: List[PnLEvent] = []

    async def ingest_all_events(self, since: float):
        """Pull all P&L events from all Purple Flea services."""
        async with self.fetcher:
            casino_events, trading_events, wallet_events, escrow_events = await asyncio.gather(
                self.fetcher.get_casino_history(since),
                self.fetcher.get_trading_history(since),
                self.fetcher.get_wallet_history(since),
                self.fetcher.get_escrow_history(since),
            )

        all_events = casino_events + trading_events + wallet_events + escrow_events
        # Process in chronological order for correct cost basis assignment
        all_events.sort(key=lambda e: e.timestamp)
        self.all_events = all_events
        log.info(f"Ingested {len(all_events)} events from all services")

    def process_events(self) -> PnLReport:
        """Process all events and compute P&L report."""
        casino_pnl  = Decimal("0")
        trading_pnl = Decimal("0")
        funding_pnl = Decimal("0")
        escrow_pnl  = Decimal("0")
        faucet_inc  = Decimal("0")
        total_fees  = Decimal("0")
        fee_breakdown: Dict[str, Decimal] = defaultdict(Decimal)
        gain_events: List[RealizedGain] = []

        for event in self.all_events:
            total_fees += event.fee_usd
            fee_breakdown[event.event_type.value] += event.fee_usd

            if event.event_type == EventType.BUY:
                self.engine.acquire(
                    event.asset, event.quantity,
                    event.price_usd + (event.fee_usd / event.quantity),  # fee adds to basis
                    event.timestamp, "trade"
                )

            elif event.event_type == EventType.SELL:
                gain, lots = self.engine.dispose(
                    event.asset, event.quantity, event.price_usd - (event.fee_usd / event.quantity),
                    event.timestamp
                )
                trading_pnl += gain
                gain_events.extend(lots)

            elif event.event_type == EventType.CASINO_WIN:
                casino_pnl += event.net_usd

            elif event.event_type == EventType.CASINO_LOSS:
                casino_pnl -= event.net_usd

            elif event.event_type == EventType.FUNDING:
                funding_pnl += event.net_usd

            elif event.event_type in (EventType.ESCROW_RECEIVE, EventType.REFERRAL_EARN):
                escrow_pnl += event.net_usd

            elif event.event_type == EventType.ESCROW_SEND:
                escrow_pnl -= event.fee_usd  # only the fee is a cost here

            elif event.event_type == EventType.FAUCET_CLAIM:
                faucet_inc += event.net_usd
                self.engine.acquire("USDC", event.quantity, Decimal("0"), event.timestamp, "faucet")

        realized = trading_pnl + casino_pnl + funding_pnl + escrow_pnl + faucet_inc
        # Unrealized: would require current prices — computed separately

        period_events = self.all_events
        start = datetime.fromtimestamp(period_events[0].timestamp, tz=timezone.utc) if period_events else datetime.now(tz=timezone.utc)
        end   = datetime.fromtimestamp(period_events[-1].timestamp, tz=timezone.utc) if period_events else datetime.now(tz=timezone.utc)

        return PnLReport(
            period_start=start,
            period_end=end,
            realized_pnl=realized,
            unrealized_pnl=Decimal("0"),  # filled below
            total_fees=total_fees,
            casino_pnl=casino_pnl,
            trading_pnl=trading_pnl,
            funding_pnl=funding_pnl,
            escrow_pnl=escrow_pnl,
            faucet_income=faucet_inc,
            net_pnl=realized - total_fees,
            gain_events=gain_events,
            fee_breakdown=dict(fee_breakdown),
        )

    def export_csv(self, report: PnLReport, path: str):
        """Export gain events to CSV for tax filing."""
        with open(path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=[
                "date", "asset", "quantity", "sale_price",
                "cost_basis", "gain_usd", "method", "lot_id"
            ])
            writer.writeheader()
            for gain in sorted(report.gain_events, key=lambda g: g.timestamp):
                writer.writerow({
                    "date":       datetime.fromtimestamp(gain.timestamp).strftime("%Y-%m-%d"),
                    "asset":      gain.asset,
                    "quantity":   str(gain.quantity_sold),
                    "sale_price": str(gain.sale_price),
                    "cost_basis": str(gain.cost_basis),
                    "gain_usd":   str(gain.gain_usd),
                    "method":     gain.method,
                    "lot_id":     gain.lot_id,
                })
        log.info(f"Exported {len(report.gain_events)} gain events to {path}")


# ─── Main: Daily Report ──────────────────────────────────────────────────────

async def generate_daily_report(api_key: str, method: CostBasisMethod = CostBasisMethod.HIFO):
    import time
    since_24h = time.time() - 86400

    tracker = AgentPnLTracker(api_key, method)
    await tracker.ingest_all_events(since_24h)
    report = tracker.process_events()

    print(f"\n=== Purple Flea Agent P&L Report ===")
    print(f"Period: {report.period_start:%Y-%m-%d %H:%M} → {report.period_end:%Y-%m-%d %H:%M}")
    print(f"Cost Basis Method: {method.value.upper()}")
    print(f"\nP&L Breakdown:")
    print(f"  Casino:     {'+' if report.casino_pnl >= 0 else ''}${report.casino_pnl:.2f}")
    print(f"  Trading:    {'+' if report.trading_pnl >= 0 else ''}${report.trading_pnl:.2f}")
    print(f"  Funding:    {'+' if report.funding_pnl >= 0 else ''}${report.funding_pnl:.2f}")
    print(f"  Escrow:     +${report.escrow_pnl:.2f}")
    print(f"  Faucet:     +${report.faucet_income:.2f}")
    print(f"  Fees:       -${report.total_fees:.2f}")
    print(f"  ─────────────────────")
    print(f"  Net P&L:    {'+' if report.net_pnl >= 0 else ''}${report.net_pnl:.2f}")

    tracker.export_csv(report, f"pnl_report_{method.value}_{int(time.time())}.csv")
    return report

if __name__ == "__main__":
    asyncio.run(generate_daily_report(API_KEY, CostBasisMethod.HIFO))