The-Ouroboros/src/db.py

"""Database layer for trade logging."""

from __future__ import annotations

import json
import sqlite3
from datetime import UTC, datetime
from pathlib import Path
from typing import Any


def init_db(db_path: str) -> sqlite3.Connection:
    """Initialize the trade logs database and return a connection."""
    if db_path != ":memory:":
        Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS trades (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT NOT NULL,
            stock_code TEXT NOT NULL,
            action TEXT NOT NULL,
            confidence INTEGER NOT NULL,
            rationale TEXT,
            quantity INTEGER,
            price REAL,
            pnl REAL DEFAULT 0.0,
            market TEXT DEFAULT 'KR',
            exchange_code TEXT DEFAULT 'KRX',
            decision_id TEXT
        )
        """
    )

    # Migration: Add market and exchange_code columns if they don't exist
    cursor = conn.execute("PRAGMA table_info(trades)")
    columns = {row[1] for row in cursor.fetchall()}

    if "market" not in columns:
        conn.execute("ALTER TABLE trades ADD COLUMN market TEXT DEFAULT 'KR'")
    if "exchange_code" not in columns:
        conn.execute("ALTER TABLE trades ADD COLUMN exchange_code TEXT DEFAULT 'KRX'")
    if "selection_context" not in columns:
        conn.execute("ALTER TABLE trades ADD COLUMN selection_context TEXT")
    if "decision_id" not in columns:
        conn.execute("ALTER TABLE trades ADD COLUMN decision_id TEXT")

    # Context tree tables for multi-layered memory management
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS contexts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            layer TEXT NOT NULL,
            timeframe TEXT NOT NULL,
            key TEXT NOT NULL,
            value TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            UNIQUE(layer, timeframe, key)
        )
        """
    )

    # Decision logging table for comprehensive audit trail
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS decision_logs (
            decision_id TEXT PRIMARY KEY,
            timestamp TEXT NOT NULL,
            stock_code TEXT NOT NULL,
            market TEXT NOT NULL,
            exchange_code TEXT NOT NULL,
            action TEXT NOT NULL,
            confidence INTEGER NOT NULL,
            rationale TEXT NOT NULL,
            context_snapshot TEXT NOT NULL,
            input_data TEXT NOT NULL,
            outcome_pnl REAL,
            outcome_accuracy INTEGER,
            reviewed INTEGER DEFAULT 0,
            review_notes TEXT
        )
        """
    )

    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS context_metadata (
            layer TEXT PRIMARY KEY,
            description TEXT NOT NULL,
            retention_days INTEGER,
            aggregation_source TEXT
        )
        """
    )

    # Playbook storage for pre-market strategy persistence
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS playbooks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            date TEXT NOT NULL,
            market TEXT NOT NULL,
            status TEXT NOT NULL DEFAULT 'pending',
            playbook_json TEXT NOT NULL,
            generated_at TEXT NOT NULL,
            token_count INTEGER DEFAULT 0,
            scenario_count INTEGER DEFAULT 0,
            match_count INTEGER DEFAULT 0,
            UNIQUE(date, market)
        )
        """
    )

    conn.execute("CREATE INDEX IF NOT EXISTS idx_playbooks_date ON playbooks(date)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_playbooks_market ON playbooks(market)")

    # Create indices for efficient context queries
    conn.execute("CREATE INDEX IF NOT EXISTS idx_contexts_layer ON contexts(layer)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_contexts_timeframe ON contexts(timeframe)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_contexts_updated ON contexts(updated_at)")

    # Create indices for efficient decision log queries
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_decision_logs_timestamp ON decision_logs(timestamp)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_decision_logs_reviewed ON decision_logs(reviewed)"
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_decision_logs_confidence ON decision_logs(confidence)"
    )
    conn.commit()
    return conn


def log_trade(
    conn: sqlite3.Connection,
    stock_code: str,
    action: str,
    confidence: int,
    rationale: str,
    quantity: int = 0,
    price: float = 0.0,
    pnl: float = 0.0,
    market: str = "KR",
    exchange_code: str = "KRX",
    selection_context: dict[str, any] | None = None,
    decision_id: str | None = None,
) -> None:
    """Insert a trade record into the database.

    Args:
        conn: Database connection
        stock_code: Stock code
        action: Trade action (BUY/SELL/HOLD)
        confidence: Confidence level (0-100)
        rationale: AI decision rationale
        quantity: Number of shares
        price: Trade price
        pnl: Profit/loss
        market: Market code
        exchange_code: Exchange code
        selection_context: Scanner selection data (RSI, volume_ratio, signal, score)
    """
    # Serialize selection context to JSON
    context_json = json.dumps(selection_context) if selection_context else None

    conn.execute(
        """
        INSERT INTO trades (
            timestamp, stock_code, action, confidence, rationale,
            quantity, price, pnl, market, exchange_code, selection_context, decision_id
        )
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            datetime.now(UTC).isoformat(),
            stock_code,
            action,
            confidence,
            rationale,
            quantity,
            price,
            pnl,
            market,
            exchange_code,
            context_json,
            decision_id,
        ),
    )
    conn.commit()


def get_latest_buy_trade(
    conn: sqlite3.Connection, stock_code: str, market: str
) -> dict[str, Any] | None:
    """Fetch the most recent BUY trade for a stock and market."""
    cursor = conn.execute(
        """
        SELECT decision_id, price, quantity
        FROM trades
        WHERE stock_code = ?
          AND market = ?
          AND action = 'BUY'
          AND decision_id IS NOT NULL
        ORDER BY timestamp DESC
        LIMIT 1
        """,
        (stock_code, market),
    )
    row = cursor.fetchone()
    if not row:
        return None
    return {"decision_id": row[0], "price": row[1], "quantity": row[2]}