database.py

"""
Database Module

SQLite database manager for trade history, portfolio snapshots, and analytics.
"""

import sqlite3
import json
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
from contextlib import contextmanager

import config
from logger import system_logger


class DatabaseManager:
    """
    Manages SQLite database for trading platform.
    Handles trade history, portfolio snapshots, and analytics.
    """
    
    def __init__(self, db_file: str = config.DATABASE_FILE):
        """
        Initialize database manager.
        
        Args:
            db_file: Path to SQLite database file
        """
        self.db_file = db_file
        self.init_database()
    
    @contextmanager
    def get_connection(self):
        """
        Context manager for database connections.
        
        Yields:
            SQLite connection
        """
        conn = sqlite3.connect(self.db_file)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception as e:
            conn.rollback()
            system_logger.error(f"Database error: {str(e)}")
            raise
        finally:
            conn.close()
    
    def init_database(self):
        """Create database tables if they don't exist."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Trades table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS trades (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    order_id TEXT UNIQUE,
                    symbol TEXT NOT NULL,
                    side TEXT NOT NULL,
                    quantity REAL NOT NULL,
                    price REAL NOT NULL,
                    order_type TEXT NOT NULL,
                    status TEXT NOT NULL,
                    filled_qty REAL DEFAULT 0,
                    filled_avg_price REAL,
                    commission REAL DEFAULT 0,
                    pl REAL,
                    timestamp TEXT NOT NULL,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Portfolio snapshots table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS portfolio_snapshots (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    equity REAL NOT NULL,
                    cash REAL NOT NULL,
                    portfolio_value REAL NOT NULL,
                    buying_power REAL NOT NULL,
                    positions_count INTEGER DEFAULT 0,
                    daily_pl REAL,
                    total_pl REAL,
                    timestamp TEXT NOT NULL,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Positions table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS positions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    snapshot_id INTEGER,
                    symbol TEXT NOT NULL,
                    quantity REAL NOT NULL,
                    avg_entry_price REAL NOT NULL,
                    current_price REAL NOT NULL,
                    market_value REAL NOT NULL,
                    unrealized_pl REAL NOT NULL,
                    unrealized_plpc REAL NOT NULL,
                    timestamp TEXT NOT NULL,
                    FOREIGN KEY (snapshot_id) REFERENCES portfolio_snapshots(id)
                )
            ''')
            
            # Activity log table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS activity_log (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    event_type TEXT NOT NULL,
                    description TEXT NOT NULL,
                    severity TEXT DEFAULT 'INFO',
                    metadata TEXT,
                    timestamp TEXT NOT NULL,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Strategy performance table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS strategy_performance (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    strategy_name TEXT NOT NULL,
                    total_trades INTEGER DEFAULT 0,
                    winning_trades INTEGER DEFAULT 0,
                    losing_trades INTEGER DEFAULT 0,
                    total_pl REAL DEFAULT 0,
                    win_rate REAL DEFAULT 0,
                    avg_win REAL DEFAULT 0,
                    avg_loss REAL DEFAULT 0,
                    sharpe_ratio REAL,
                    max_drawdown REAL,
                    timestamp TEXT NOT NULL,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Settings table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS settings (
                    key TEXT PRIMARY KEY,
                    value TEXT NOT NULL,
                    updated_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create indexes for performance
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_trades_timestamp ON trades(timestamp)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_positions_symbol ON positions(symbol)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_activity_timestamp ON activity_log(timestamp)')
            
            system_logger.info("Database initialized successfully")
    
    # Trade operations
    
    def save_trade(self, trade_data: Dict[str, Any]) -> int:
        """
        Save trade to database.
        
        Args:
            trade_data: Trade information dictionary
            
        Returns:
            Trade ID
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO trades (
                    order_id, symbol, side, quantity, price, order_type,
                    status, filled_qty, filled_avg_price, commission, pl, timestamp
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                trade_data.get('order_id'),
                trade_data.get('symbol'),
                trade_data.get('side'),
                trade_data.get('quantity'),
                trade_data.get('price'),
                trade_data.get('order_type'),
                trade_data.get('status'),
                trade_data.get('filled_qty', 0),
                trade_data.get('filled_avg_price'),
                trade_data.get('commission', 0),
                trade_data.get('pl'),
                datetime.now().isoformat()
            ))
            
            return cursor.lastrowid
    
    def get_trades(
        self,
        symbol: Optional[str] = None,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """
        Get trades from database with optional filtering.
        
        Args:
            symbol: Filter by symbol
            start_date: Start date filter
            end_date: End date filter
            limit: Maximum number of trades to return
            
        Returns:
            List of trade dictionaries
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            query = 'SELECT * FROM trades WHERE 1=1'
            params = []
            
            if symbol:
                query += ' AND symbol = ?'
                params.append(symbol)
            
            if start_date:
                query += ' AND timestamp >= ?'
                params.append(start_date.isoformat())
            
            if end_date:
                query += ' AND timestamp <= ?'
                params.append(end_date.isoformat())
            
            query += ' ORDER BY timestamp DESC LIMIT ?'
            params.append(limit)
            
            cursor.execute(query, params)
            rows = cursor.fetchall()
            
            return [dict(row) for row in rows]
    
    def update_trade_status(self, order_id: str, status: str, filled_data: Optional[Dict] = None):
        """
        Update trade status.
        
        Args:
            order_id: Order ID
            status: New status
            filled_data: Optional filled execution data
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            if filled_data:
                cursor.execute('''
                    UPDATE trades
                    SET status = ?, filled_qty = ?, filled_avg_price = ?
                    WHERE order_id = ?
                ''', (
                    status,
                    filled_data.get('filled_qty'),
                    filled_data.get('filled_avg_price'),
                    order_id
                ))
            else:
                cursor.execute('''
                    UPDATE trades SET status = ? WHERE order_id = ?
                ''', (status, order_id))
    
    # Portfolio operations
    
    def save_portfolio_snapshot(
        self,
        account_data: Dict[str, Any],
        positions: List[Dict[str, Any]]
    ) -> int:
        """
        Save portfolio snapshot.
        
        Args:
            account_data: Account information
            positions: List of positions
            
        Returns:
            Snapshot ID
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Save snapshot
            cursor.execute('''
                INSERT INTO portfolio_snapshots (
                    equity, cash, portfolio_value, buying_power,
                    positions_count, daily_pl, timestamp
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                account_data.get('equity'),
                account_data.get('cash_balance'),
                account_data.get('portfolio_value'),
                account_data.get('buying_power'),
                len(positions),
                account_data.get('day_pl'),
                datetime.now().isoformat()
            ))
            
            snapshot_id = cursor.lastrowid
            
            # Save positions
            for position in positions:
                cursor.execute('''
                    INSERT INTO positions (
                        snapshot_id, symbol, quantity, avg_entry_price,
                        current_price, market_value, unrealized_pl,
                        unrealized_plpc, timestamp
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    snapshot_id,
                    position.get('symbol'),
                    position.get('quantity'),
                    position.get('avg_entry_price'),
                    position.get('current_price'),
                    position.get('market_value'),
                    position.get('unrealized_pl'),
                    position.get('unrealized_plpc'),
                    datetime.now().isoformat()
                ))
            
            return snapshot_id
    
    def get_portfolio_history(
        self,
        days: int = 30
    ) -> List[Dict[str, Any]]:
        """
        Get portfolio history.
        
        Args:
            days: Number of days of history
            
        Returns:
            List of portfolio snapshots
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            start_date = (datetime.now() - timedelta(days=days)).isoformat()
            
            cursor.execute('''
                SELECT * FROM portfolio_snapshots
                WHERE timestamp >= ?
                ORDER BY timestamp ASC
            ''', (start_date,))
            
            rows = cursor.fetchall()
            return [dict(row) for row in rows]
    
    # Activity log operations
    
    def log_activity(
        self,
        event_type: str,
        description: str,
        severity: str = 'INFO',
        metadata: Optional[Dict] = None
    ):
        """
        Log activity event.
        
        Args:
            event_type: Event type
            description: Event description
            severity: Severity level
            metadata: Optional metadata dictionary
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO activity_log (
                    event_type, description, severity, metadata, timestamp
                ) VALUES (?, ?, ?, ?, ?)
            ''', (
                event_type,
                description,
                severity,
                json.dumps(metadata) if metadata else None,
                datetime.now().isoformat()
            ))
    
    def get_activity_log(
        self,
        hours: int = 24,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """
        Get recent activity log.
        
        Args:
            hours: Hours of history to retrieve
            limit: Maximum number of entries
            
        Returns:
            List of activity log entries
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            start_time = (datetime.now() - timedelta(hours=hours)).isoformat()
            
            cursor.execute('''
                SELECT * FROM activity_log
                WHERE timestamp >= ?
                ORDER BY timestamp DESC
                LIMIT ?
            ''', (start_time, limit))
            
            rows = cursor.fetchall()
            return [dict(row) for row in rows]
    
    # Analytics operations
    
    def get_trading_statistics(self) -> Dict[str, Any]:
        """
        Calculate trading statistics.
        
        Returns:
            Dictionary of trading statistics
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Total trades
            cursor.execute("SELECT COUNT(*) as count FROM trades WHERE status = 'filled'")
            total_trades = cursor.fetchone()['count']
            
            # Winning trades
            cursor.execute("SELECT COUNT(*) as count FROM trades WHERE status = 'filled' AND pl > 0")
            winning_trades = cursor.fetchone()['count']
            
            # Losing trades
            cursor.execute("SELECT COUNT(*) as count FROM trades WHERE status = 'filled' AND pl < 0")
            losing_trades = cursor.fetchone()['count']
            
            # Total P&L
            cursor.execute("SELECT SUM(pl) as total FROM trades WHERE status = 'filled'")
            total_pl = cursor.fetchone()['total'] or 0
            
            # Average win/loss
            cursor.execute("SELECT AVG(pl) as avg FROM trades WHERE status = 'filled' AND pl > 0")
            avg_win = cursor.fetchone()['avg'] or 0
            
            cursor.execute("SELECT AVG(pl) as avg FROM trades WHERE status = 'filled' AND pl < 0")
            avg_loss = cursor.fetchone()['avg'] or 0
            
            # Win rate
            win_rate = (winning_trades / total_trades * 100) if total_trades > 0 else 0
            
            return {
                'total_trades': total_trades,
                'winning_trades': winning_trades,
                'losing_trades': losing_trades,
                'win_rate': win_rate,
                'total_pl': total_pl,
                'avg_win': avg_win,
                'avg_loss': avg_loss
            }
    
    def cleanup_old_data(self, days: int = 365):
        """
        Clean up old data to prevent database bloat.
        
        Args:
            days: Keep data newer than this many days
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            cutoff_date = (datetime.now() - timedelta(days=days)).isoformat()
            
            # Clean old trades
            cursor.execute('DELETE FROM trades WHERE timestamp < ?', (cutoff_date,))
            
            # Clean old snapshots
            cursor.execute('DELETE FROM portfolio_snapshots WHERE timestamp < ?', (cutoff_date,))
            
            # Clean old activity logs
            cursor.execute('DELETE FROM activity_log WHERE timestamp < ?', (cutoff_date,))
            
            # Vacuum database
            cursor.execute('VACUUM')
            
            system_logger.info(f"Cleaned up data older than {days} days")