#!/usr/bin/env python3
"""
SQLite database for storing signals and trades.

Provides persistence for:
- Market signals
- Trade history
- Performance metrics
"""

import sqlite3
import os
from datetime import datetime
from typing import Optional, List, Dict, Any
from contextlib import contextmanager


class SignalsDatabase:
    """
    SQLite database for trading signals and trade records.
    
    Tables:
        signals: Market signals from Markov engine
        trades: Executed trades
        performance: Daily performance metrics
    """
    
    def __init__(self, db_path: str = None):
        """
        Initialize database connection.
        
        Args:
            db_path: Path to SQLite database file
        """
        if db_path is None:
            db_path = os.path.expanduser('~/bitcoin-bot/signals.db')
        
        # Ensure directory exists
        db_dir = os.path.dirname(db_path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir)
        
        self.db_path = db_path
        self._init_database()
    
    @contextmanager
    def _get_connection(self):
        """Get database connection with proper cleanup."""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()
    
    def _init_database(self) -> None:
        """Initialize database tables."""
        with self._get_connection() as conn:
            cursor = conn.cursor()
            
            # Signals table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS signals (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol TEXT NOT NULL,
                    timestamp REAL NOT NULL,
                    state TEXT NOT NULL,
                    tau REAL NOT NULL,
                    eps REAL NOT NULL,
                    signal_type TEXT,
                    entry_price REAL,
                    stop_loss REAL,
                    take_profit REAL,
                    kelly_fraction REAL,
                    executed INTEGER DEFAULT 0,
                    trade_id INTEGER,
                    execution_note TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Trades table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS trades (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    symbol TEXT NOT NULL,
                    signal_id INTEGER,
                    side TEXT NOT NULL,
                    entry_price REAL NOT NULL,
                    exit_price REAL,
                    amount REAL NOT NULL,
                    pnl REAL,
                    pnl_pct REAL,
                    status TEXT NOT NULL DEFAULT 'open',
                    entry_time TIMESTAMP NOT NULL,
                    exit_time TIMESTAMP,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (signal_id) REFERENCES signals(id)
                )
            ''')
            
            # Performance table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS performance (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date TEXT NOT NULL UNIQUE,
                    total_trades INTEGER DEFAULT 0,
                    winning_trades INTEGER DEFAULT 0,
                    losing_trades INTEGER DEFAULT 0,
                    total_pnl REAL DEFAULT 0,
                    max_drawdown REAL DEFAULT 0,
                    sharpe_ratio REAL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create indexes
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_signals_symbol ON signals(symbol)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_signals_timestamp ON signals(timestamp)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_trades_status ON trades(status)')
            
            conn.commit()
    
    def store_signal(
        self,
        symbol: str,
        timestamp: float,
        state: str,
        tau: float,
        eps: float,
        signal_type: Optional[str] = None,
        entry_price: Optional[float] = None,
        stop_loss: Optional[float] = None,
        take_profit: Optional[float] = None,
        kelly_fraction: Optional[float] = None
    ) -> int:
        """
        Store a market signal.
        
        Args:
            symbol: Trading symbol
            timestamp: Signal timestamp
            state: Markov state
            tau: Persistence metric
            eps: Transition metric
            signal_type: 'entry' or 'exit'
            entry_price: Suggested entry price
            stop_loss: Suggested stop loss
            take_profit: Suggested take profit
            kelly_fraction: Position size fraction
            
        Returns:
            Signal ID
        """
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO signals (
                    symbol, timestamp, state, tau, eps,
                    signal_type, entry_price, stop_loss, take_profit, kelly_fraction
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                symbol, timestamp, state, tau, eps,
                signal_type, entry_price, stop_loss, take_profit, kelly_fraction
            ))
            conn.commit()
            return cursor.lastrowid
    
    def get_recent_signals(
        self,
        symbol: str,
        count: int = 100
    ) -> List[Dict[str, Any]]:
        """
        Get recent signals for a symbol.
        
        Args:
            symbol: Trading symbol
            count: Number of signals to retrieve
            
        Returns:
            List of signal dictionaries
        """
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM signals
                WHERE symbol = ?
                ORDER BY timestamp DESC
                LIMIT ?
            ''', (symbol, count))
            
            return [dict(row) for row in cursor.fetchall()]
    
    def store_trade(
        self,
        symbol: str,
        signal_id: int,
        side: str,
        entry_price: float,
        amount: float,
        exit_price: Optional[float] = None
    ) -> int:
        """
        Store a trade record.
        
        Args:
            symbol: Trading symbol
            signal_id: Associated signal ID
            side: 'buy' or 'sell'
            entry_price: Entry price
            amount: Position size
            exit_price: Exit price (None for open trades)
            
        Returns:
            Trade ID
        """
        now = datetime.now()
        
        # Calculate PnL if trade is closed
        pnl = None
        pnl_pct = None
        status = 'closed' if exit_price else 'open'
        exit_time = now if exit_price else None
        
        if exit_price:
            if side == 'buy':
                pnl = (exit_price - entry_price) * amount
                pnl_pct = (exit_price - entry_price) / entry_price
            else:
                pnl = (entry_price - exit_price) * amount
                pnl_pct = (entry_price - exit_price) / entry_price
        
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO trades (
                    symbol, signal_id, side, entry_price, exit_price,
                    amount, pnl, pnl_pct, status, entry_time, exit_time
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                symbol, signal_id, side, entry_price, exit_price,
                amount, pnl, pnl_pct, status, now, exit_time
            ))
            conn.commit()
            return cursor.lastrowid
    
    def close_trade(
        self,
        trade_id: int,
        exit_price: float
    ) -> None:
        """
        Close an open trade.
        
        Args:
            trade_id: Trade ID to close
            exit_price: Exit price
        """
        with self._get_connection() as conn:
            cursor = conn.cursor()
            
            # Get trade details
            cursor.execute('SELECT * FROM trades WHERE id = ?', (trade_id,))
            trade = cursor.fetchone()
            
            if not trade:
                return
            
            # Calculate PnL
            entry_price = trade['entry_price']
            amount = trade['amount']
            side = trade['side']
            
            if side == 'buy':
                pnl = (exit_price - entry_price) * amount
                pnl_pct = (exit_price - entry_price) / entry_price
            else:
                pnl = (entry_price - exit_price) * amount
                pnl_pct = (entry_price - exit_price) / entry_price
            
            # Update trade
            cursor.execute('''
                UPDATE trades
                SET exit_price = ?, pnl = ?, pnl_pct = ?, status = 'closed', exit_time = ?
                WHERE id = ?
            ''', (exit_price, pnl, pnl_pct, datetime.now(), trade_id))
            
            conn.commit()
    
    def get_open_trades(self) -> List[Dict[str, Any]]:
        """Get all open trades."""
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM trades WHERE status = ?', ('open',))
            return [dict(row) for row in cursor.fetchall()]
    
    def get_trade_history(
        self,
        symbol: Optional[str] = None,
        days: int = 30
    ) -> List[Dict[str, Any]]:
        """
        Get trade history.
        
        Args:
            symbol: Filter by symbol (None for all)
            days: Days of history to retrieve
            
        Returns:
            List of trade dictionaries
        """
        with self._get_connection() as conn:
            cursor = conn.cursor()
            
            if symbol:
                cursor.execute('''
                    SELECT * FROM trades
                    WHERE symbol = ? AND status = 'closed'
                    ORDER BY entry_time DESC
                ''', (symbol,))
            else:
                cursor.execute('''
                    SELECT * FROM trades
                    WHERE status = 'closed'
                    ORDER BY entry_time DESC
                ''')
            
            return [dict(row) for row in cursor.fetchall()]
    
    def get_performance_summary(self, days: int = 30) -> Dict[str, Any]:
        """
        Get performance summary.
        
        Args:
            days: Days to include in summary
            
        Returns:
            Performance metrics dictionary
        """
        trades = self.get_trade_history(days=days)
        
        if not trades:
            return {
                'total_trades': 0,
                'winning_trades': 0,
                'losing_trades': 0,
                'win_rate': 0.0,
                'total_pnl': 0.0,
                'avg_pnl': 0.0,
                'avg_pnl_pct': 0.0
            }
        
        total = len(trades)
        winners = [t for t in trades if t['pnl'] and t['pnl'] > 0]
        losers = [t for t in trades if t['pnl'] and t['pnl'] <= 0]
        
        total_pnl = sum(t['pnl'] or 0 for t in trades)
        win_rate = len(winners) / total if total > 0 else 0
        
        return {
            'total_trades': total,
            'winning_trades': len(winners),
            'losing_trades': len(losers),
            'win_rate': win_rate,
            'total_pnl': total_pnl,
            'avg_pnl': total_pnl / total if total > 0 else 0,
            'avg_pnl_pct': sum(t['pnl_pct'] or 0 for t in trades) / total if total > 0 else 0
        }
    
    async def record_signal(
        self,
        strategy: str = None,
        signal: str = None,
        confidence: float = None,
        reason: str = None,
        metadata: Dict[str, Any] = None,
        executed: bool = False,
        # For backward compatibility with markov engine
        symbol: str = None,
        state: str = None,
        tau: float = None,
        eps: float = None,
        signal_type: str = None,
        price: float = None
    ) -> int:
        """
        Record a trading signal (async wrapper for store_signal).
        
        Supports both the generic signal format and the markov-specific format.
        
        Args:
            strategy: Strategy name (e.g., 'markov_engine')
            signal: Signal type (e.g., 'entry', 'no_entry')
            confidence: Signal confidence
            reason: Signal reason
            metadata: Signal metadata dict
            executed: Whether signal was executed
            symbol: Trading symbol (markov format)
            state: Markov state (markov format)
            tau: Persistence value (markov format)
            eps: Gap value (markov format)
            signal_type: 'entry' or 'exit' (markov format)
            price: Current price (markov format)
            
        Returns:
            Signal ID
        """
        # Use markov format if those args are provided
        if symbol is not None and state is not None:
            return self.store_signal(
                symbol=symbol,
                timestamp=datetime.now().timestamp(),
                state=state,
                tau=tau or 0.0,
                eps=eps or 0.0,
                signal_type=signal_type,
                entry_price=price,
                kelly_fraction=confidence
            )
        
        # Otherwise use generic format - still store with available data
        if symbol is None:
            symbol = metadata.get('symbol', 'UNKNOWN') if metadata else 'UNKNOWN'
            
        return self.store_signal(
            symbol=symbol,
            timestamp=datetime.now().timestamp(),
            state=state or (metadata.get('state', 'unknown') if metadata else 'unknown'),
            tau=tau or (metadata.get('persistence', 0.0) if metadata else 0.0),
            eps=eps or (metadata.get('arbitrage_gap', 0.0) if metadata else 0.0),
            signal_type=signal or signal_type,
            entry_price=price or (metadata.get('current_price') if metadata else None),
            kelly_fraction=confidence
        )

    async def update_signal_executed(
        self,
        signal_id: int,
        executed: bool,
        trade_id: Optional[int] = None,
        note: str = None
    ) -> None:
        """
        Update signal execution status.
        
        Args:
            signal_id: Signal ID to update
            executed: Whether signal was executed
            trade_id: Associated trade ID if executed
            note: Optional execution note
        """
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE signals
                SET executed = ?, trade_id = ?, execution_note = ?
                WHERE id = ?
            ''', (1 if executed else 0, trade_id, note, signal_id))
            conn.commit()
