database.py

"""
database.py
===========
Data Collection Pipeline for the DeFi Risk Simulation Lab.

Responsibilities:
  1. Creates and manages a SQLite database for all agent actions.
  2. Provides a SimulationLogger class for agents / model to call.
  3. On simulation end, exports the full table to CSV for ML pipelines.

Schema (agent_actions table):
  id               INTEGER  — auto-increment primary key
  step             INTEGER  — Mesa model step number
  timestamp        TEXT     — ISO-8601 datetime string
  agent_id         INTEGER  — Mesa agent unique_id
  agent_type       TEXT     — e.g. "ArbitrageurAgent"
  action           TEXT     — e.g. "SWAP", "LIQUIDATE", "BORROW", "IDLE"
  amount           REAL     — token amount involved (0 if IDLE)
  token            TEXT     — "ETH", "USDC", or ""
  pool_reserve_usdc REAL    — pool state AFTER the action
  pool_reserve_eth  REAL    — pool state AFTER the action
  token_price      REAL     — AMM-derived price AFTER the action
  oracle_price     REAL     — oracle price at time of action
  reasoning        TEXT     — LLM reasoning (Wildcard agent only, else "")
"""

import sqlite3
import csv
import logging
from datetime import datetime, timezone
from typing import Optional

import config

logger = logging.getLogger(__name__)

# ------------------------------------------------------------------ #
# DDL STATEMENT                                                        #
# ------------------------------------------------------------------ #
_CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS agent_actions (
    id                INTEGER PRIMARY KEY AUTOINCREMENT,
    step              INTEGER NOT NULL,
    timestamp         TEXT    NOT NULL,
    agent_id          INTEGER NOT NULL,
    agent_type        TEXT    NOT NULL,
    action            TEXT    NOT NULL,
    amount            REAL    NOT NULL DEFAULT 0.0,
    token             TEXT    NOT NULL DEFAULT '',
    pool_reserve_usdc REAL    NOT NULL,
    pool_reserve_eth  REAL    NOT NULL,
    token_price       REAL    NOT NULL,
    oracle_price      REAL    NOT NULL,
    reasoning         TEXT    NOT NULL DEFAULT ''
);
"""

_INSERT_SQL = """
INSERT INTO agent_actions (
    step, timestamp, agent_id, agent_type, action,
    amount, token,
    pool_reserve_usdc, pool_reserve_eth, token_price, oracle_price,
    reasoning
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
"""


class SimulationLogger:
    """
    Thread-safe SQLite logger for agent actions.

    Usage:
        db = SimulationLogger()
        db.log(step=1, agent_id=0, agent_type="Whale", action="SWAP", ...)
        db.export_csv()
        db.close()
    """

    def __init__(self, db_path: str = config.DB_PATH):
        self._db_path = db_path
        self._conn    = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL;")  # better concurrent writes
        self._conn.execute(_CREATE_TABLE_SQL)
        self._conn.commit()
        logger.info("SimulationLogger connected to %s", db_path)

    # ---------------------------------------------------------------- #
    # WRITE                                                              #
    # ---------------------------------------------------------------- #

    def log(
        self,
        step:              int,
        agent_id:          int,
        agent_type:        str,
        action:            str,
        pool_reserve_usdc: float,
        pool_reserve_eth:  float,
        token_price:       float,
        oracle_price:      float,
        amount:            float = 0.0,
        token:             str   = "",
        reasoning:         str   = "",
    ) -> None:
        """
        Insert one action record into the database.

        Args:
            step              : Current simulation step number.
            agent_id          : Mesa agent unique_id.
            agent_type        : Class name of the agent.
            action            : Action label (SWAP, LIQUIDATE, BORROW, SUPPLY, IDLE, ...).
            pool_reserve_usdc : Pool USDC reserve after the action.
            pool_reserve_eth  : Pool ETH reserve after the action.
            token_price       : AMM-derived token price after the action.
            oracle_price      : External oracle price at this step.
            amount            : Token amount involved (default 0 for IDLE).
            token             : Token symbol (default "" for IDLE).
            reasoning         : LLM reasoning string (Wildcard only).
        """
        ts = datetime.now(timezone.utc).isoformat()
        try:
            self._conn.execute(
                _INSERT_SQL,
                (
                    step, ts, agent_id, agent_type, action,
                    amount, token,
                    pool_reserve_usdc, pool_reserve_eth, token_price, oracle_price,
                    reasoning,
                ),
            )
            self._conn.commit()
        except sqlite3.Error as exc:
            logger.error("DB write error: %s", exc)

    # ---------------------------------------------------------------- #
    # READ / EXPORT                                                      #
    # ---------------------------------------------------------------- #

    def export_csv(self, csv_path: str = config.CSV_PATH) -> str:
        """
        Dumps the entire agent_actions table to a CSV file.

        Returns:
            str: Path to the written CSV file.
        """
        cursor = self._conn.execute("SELECT * FROM agent_actions ORDER BY id;")
        columns = [desc[0] for desc in cursor.description]
        rows    = cursor.fetchall()

        with open(csv_path, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow(columns)
            writer.writerows(rows)

        logger.info("Exported %d rows to %s", len(rows), csv_path)
        return csv_path

    def get_summary(self) -> dict:
        """Returns a quick summary dict for the end-of-run console print."""
        cursor = self._conn.execute(
            "SELECT agent_type, action, COUNT(*) as cnt "
            "FROM agent_actions GROUP BY agent_type, action ORDER BY agent_type, action;"
        )
        rows = cursor.fetchall()
        summary = {}
        for agent_type, action, cnt in rows:
            summary.setdefault(agent_type, {})[action] = cnt
        return summary

    def close(self) -> None:
        """Closes the database connection cleanly."""
        self._conn.close()
        logger.info("SimulationLogger connection closed.")