"""Database models for trading bot persistence."""

import asyncio
import sqlite3
from datetime import datetime
from typing import Optional, List, Dict, Any


class Database:
    """SQLite database manager for trading bot.
    
    Provides async interface for:
    - Trade records
    - Execution logs
    - Strategy performance
    - Bot state
    """
    
    def __init__(self, db_path: str = "trading_bot.db"):
        """Initialize database.
        
        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
    
    def _get_connection(self) -> sqlite3.Connection:
        """Get database connection."""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def _initialize_schema(self) -> None:
        """Initialize database schema."""
        with self._get_connection() as conn:
            # Trades table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS trades (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    symbol TEXT NOT NULL,
                    side TEXT NOT NULL,
                    amount REAL NOT NULL,
                    price REAL NOT NULL,
                    fee REAL DEFAULT 0,
                    fee_currency TEXT,
                    profit REAL,
                    strategy TEXT,
                    decision_metadata TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Execution logs
            conn.execute("""
                CREATE TABLE IF NOT EXISTS execution_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    action TEXT NOT NULL,
                    status TEXT NOT NULL,
                    symbol TEXT,
                    details TEXT,
                    error_message TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Strategy signals
            conn.execute("""
                CREATE TABLE IF NOT EXISTS strategy_signals (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    strategy TEXT NOT NULL,
                    signal TEXT NOT NULL,
                    confidence REAL,
                    reason TEXT,
                    metadata TEXT,
                    executed BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Bot state
            conn.execute("""
                CREATE TABLE IF NOT EXISTS bot_state (
                    key TEXT PRIMARY KEY,
                    value TEXT,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # ========== NEW TABLES FOR MARKOV CHAIN STRATEGY ==========
            
            # Markov state history - track state transitions over time
            conn.execute("""
                CREATE TABLE IF NOT EXISTS markov_state_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    symbol TEXT NOT NULL,
                    state_id INTEGER NOT NULL,
                    state_price_low REAL,
                    state_price_high REAL,
                    persistence REAL,
                    optimal_next_state INTEGER,
                    transition_probability REAL,
                    arbitrage_gap REAL,
                    should_enter BOOLEAN,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Transition matrix snapshots - store P matrix at intervals
            conn.execute("""
                CREATE TABLE IF NOT EXISTS transition_matrix_snapshots (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    symbol TEXT NOT NULL,
                    n_states INTEGER NOT NULL,
                    price_min REAL NOT NULL,
                    price_max REAL NOT NULL,
                    matrix_data TEXT NOT NULL,  -- JSON string of matrix
                    diagonal_persistence TEXT,  -- JSON array of diagonal values
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Kelly criterion metrics - track for position sizing
            conn.execute("""
                CREATE TABLE IF NOT EXISTS kelly_metrics (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    symbol TEXT NOT NULL,
                    strategy TEXT,
                    total_trades INTEGER DEFAULT 0,
                    winning_trades INTEGER DEFAULT 0,
                    win_rate REAL,
                    avg_win REAL,
                    avg_loss REAL,
                    f_star REAL,
                    capped_f_star REAL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Multi-asset portfolio tracking - for variance reduction analysis
            conn.execute("""
                CREATE TABLE IF NOT EXISTS multi_asset_portfolio (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    symbols TEXT NOT NULL,  -- JSON array of symbols
                    total_stake REAL,
                    individual_stakes TEXT,  -- JSON object: symbol -> stake
                    combined_return REAL,
                    variance_single_asset REAL,
                    variance_multi_asset REAL,
                    variance_reduction_pct REAL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            # Create indexes
            conn.execute("CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_trades_timestamp ON trades(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_logs_timestamp ON execution_logs(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_signals_timestamp ON strategy_signals(timestamp)")
            
            # Indexes for new tables
            conn.execute("CREATE INDEX IF NOT EXISTS idx_markov_state_timestamp ON markov_state_history(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_markov_state_symbol ON markov_state_history(symbol)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_transition_matrix_timestamp ON transition_matrix_snapshots(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_kelly_metrics_timestamp ON kelly_metrics(timestamp)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_kelly_metrics_symbol ON kelly_metrics(symbol)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_portfolio_timestamp ON multi_asset_portfolio(timestamp)")
            
            conn.commit()
    
    async def initialize(self) -> None:
        """Initialize database (call at startup)."""
        await asyncio.get_event_loop().run_in_executor(
            None, self._initialize_schema
        )
    
    async def record_trade(self, trade: Dict[str, Any]) -> int:
        """Record a trade execution.
        
        Args:
            trade: Trade data dictionary
            
        Returns:
            Trade ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO trades 
                   (timestamp, symbol, side, amount, price, fee, fee_currency, 
                    profit, strategy, decision_metadata)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    trade.get('timestamp', datetime.now().timestamp()),
                    trade.get('symbol', ''),
                    trade.get('side', ''),
                    trade.get('amount', 0),
                    trade.get('price', 0),
                    trade.get('fee', 0),
                    trade.get('fee_currency', ''),
                    trade.get('profit'),
                    trade.get('strategy', ''),
                    str(trade.get('metadata', {})),
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def record_signal(
        self,
        strategy: str,
        signal: str,
        confidence: float,
        reason: str,
        metadata: Dict[str, Any] = None,
        executed: bool = False
    ) -> int:
        """Record a strategy signal.
        
        Args:
            strategy: Strategy name
            signal: Signal type (buy, sell, hold)
            confidence: Signal confidence
            reason: Signal reason
            metadata: Signal metadata
            executed: Whether signal was executed
            
        Returns:
            Signal ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO strategy_signals 
                   (timestamp, strategy, signal, confidence, reason, metadata, executed)
                   VALUES (?, ?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    strategy,
                    signal,
                    confidence,
                    reason,
                    str(metadata or {}),
                    executed,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def record_log(
        self,
        action: str,
        status: str,
        symbol: str = None,
        details: str = None,
        error: str = None
    ) -> int:
        """Record an execution log.
        
        Args:
            action: Action performed
            status: Status (success, failed, pending)
            symbol: Trading symbol
            details: Additional details
            error: Error message if failed
            
        Returns:
            Log ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO execution_logs 
                   (timestamp, action, status, symbol, details, error_message)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    action,
                    status,
                    symbol,
                    details,
                    error,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def get_recent_trades(
        self,
        symbol: str = None,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """Get recent trades.
        
        Args:
            symbol: Filter by symbol
            limit: Maximum trades to return
            
        Returns:
            List of trade records
        """
        with self._get_connection() as conn:
            if symbol:
                rows = conn.execute(
                    """SELECT * FROM trades 
                       WHERE symbol = ? 
                       ORDER BY timestamp DESC 
                       LIMIT ?""",
                    (symbol, limit)
                )
            else:
                rows = conn.execute(
                    """SELECT * FROM trades 
                       ORDER BY timestamp DESC 
                       LIMIT ?""",
                    (limit,)
                )
            
            return [dict(row) for row in rows.fetchall()]
    
    async def get_signal_stats(self, strategy: str = None) -> Dict[str, Any]:
        """Get strategy signal statistics.
        
        Args:
            strategy: Strategy name (None for all)
            
        Returns:
            Statistics dictionary
        """
        with self._get_connection() as conn:
            if strategy:
                stats = conn.execute(
                    """SELECT 
                           COUNT(*) as total,
                           SUM(CASE WHEN signal = 'buy' THEN 1 ELSE 0 END) as buy_signals,
                           SUM(CASE WHEN signal = 'sell' THEN 1 ELSE 0 END) as sell_signals,
                           AVG(confidence) as avg_confidence,
                           SUM(executed) as executed_count
                       FROM strategy_signals 
                       WHERE strategy = ?""",
                    (strategy,)
                ).fetchone()
            else:
                stats = conn.execute(
                    """SELECT 
                           COUNT(*) as total,
                           SUM(CASE WHEN signal = 'buy' THEN 1 ELSE 0 END) as buy_signals,
                           SUM(CASE WHEN signal = 'sell' THEN 1 ELSE 0 END) as sell_signals,
                           AVG(confidence) as avg_confidence,
                           SUM(executed) as executed_count
                       FROM strategy_signals"""
                ).fetchone()
            
            return dict(stats) if stats else {}
    
    async def set_state(self, key: str, value: str) -> None:
        """Set a state value.
        
        Args:
            key: State key
            value: State value
        """
        with self._get_connection() as conn:
            conn.execute(
                """INSERT OR REPLACE INTO bot_state (key, value, updated_at)
                   VALUES (?, ?, CURRENT_TIMESTAMP)""",
                (key, value)
            )
            conn.commit()
    
    async def get_state(self, key: str, default: str = None) -> str:
        """Get a state value.
        
        Args:
            key: State key
            default: Default value if not found
            
        Returns:
            State value or default
        """
        with self._get_connection() as conn:
            row = conn.execute(
                'SELECT value FROM bot_state WHERE key = ?',
                (key,)
            ).fetchone()
            return row['value'] if row else default
    
    async def close(self) -> None:
        """Close database connections."""
        pass  # SQLite handles this automatically
    
    # ========== MARKOV CHAIN PERSISTENCE METHODS ==========
    
    async def record_markov_state(
        self,
        symbol: str,
        state_id: int,
        state_price_range: tuple,
        persistence: float,
        optimal_next_state: int,
        transition_probability: float,
        arbitrage_gap: float,
        should_enter: bool
    ) -> int:
        """Record a Markov state observation.
        
        Args:
            symbol: Trading symbol
            state_id: Current state ID
            state_price_range: (low, high) price range for state
            persistence: Diagonal value P[state][state]
            optimal_next_state: Most likely next state
            transition_probability: Probability of optimal transition
            arbitrage_gap: Model probability - market price
            should_enter: Entry decision based on tau/eps thresholds
            
        Returns:
            Record ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO markov_state_history
                   (timestamp, symbol, state_id, state_price_low, state_price_high,
                    persistence, optimal_next_state, transition_probability,
                    arbitrage_gap, should_enter)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    symbol,
                    state_id,
                    state_price_range[0],
                    state_price_range[1],
                    persistence,
                    optimal_next_state,
                    transition_probability,
                    arbitrage_gap,
                    should_enter,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def get_markov_state_history(
        self,
        symbol: str,
        limit: int = 100
    ) -> List[Dict[str, Any]]:
        """Get recent Markov state history.
        
        Args:
            symbol: Trading symbol
            limit: Maximum records to return
            
        Returns:
            List of state history records
        """
        with self._get_connection() as conn:
            rows = conn.execute(
                """SELECT * FROM markov_state_history
                   WHERE symbol = ?
                   ORDER BY timestamp DESC
                   LIMIT ?""",
                (symbol, limit)
            )
            return [dict(row) for row in rows.fetchall()]
    
    async def record_transition_matrix(
        self,
        symbol: str,
        n_states: int,
        price_min: float,
        price_max: float,
        matrix_data: str,
        diagonal_persistence: str = None
    ) -> int:
        """Record a transition matrix snapshot.
        
        Args:
            symbol: Trading symbol
            n_states: Number of states in matrix
            price_min: Minimum price range
            price_max: Maximum price range
            matrix_data: JSON string of transition matrix
            diagonal_persistence: JSON array of diagonal values
            
        Returns:
            Record ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO transition_matrix_snapshots
                   (timestamp, symbol, n_states, price_min, price_max,
                    matrix_data, diagonal_persistence)
                   VALUES (?, ?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    symbol,
                    n_states,
                    price_min,
                    price_max,
                    matrix_data,
                    diagonal_persistence,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def get_latest_transition_matrix(self, symbol: str) -> Dict[str, Any]:
        """Get the most recent transition matrix for a symbol.
        
        Args:
            symbol: Trading symbol
            
        Returns:
            Matrix snapshot or None
        """
        with self._get_connection() as conn:
            row = conn.execute(
                """SELECT * FROM transition_matrix_snapshots
                   WHERE symbol = ?
                   ORDER BY timestamp DESC
                   LIMIT 1""",
                (symbol,)
            ).fetchone()
            return dict(row) if row else None
    
    # ========== KELLY CRITERION PERSISTENCE METHODS ==========
    
    async def record_kelly_metrics(
        self,
        symbol: str,
        strategy: str,
        total_trades: int,
        winning_trades: int,
        win_rate: float,
        avg_win: float,
        avg_loss: float,
        f_star: float,
        capped_f_star: float
    ) -> int:
        """Record Kelly criterion metrics.
        
        Args:
            symbol: Trading symbol
            strategy: Strategy name
            total_trades: Total number of trades
            winning_trades: Number of winning trades
            win_rate: Win rate (0-1)
            avg_win: Average winning trade size
            avg_loss: Average losing trade size
            f_star: Calculated Kelly fraction
            capped_f_star: Capped Kelly fraction (max 0.71)
            
        Returns:
            Record ID
        """
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO kelly_metrics
                   (timestamp, symbol, strategy, total_trades, winning_trades,
                    win_rate, avg_win, avg_loss, f_star, capped_f_star)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    symbol,
                    strategy,
                    total_trades,
                    winning_trades,
                    win_rate,
                    avg_win,
                    avg_loss,
                    f_star,
                    capped_f_star,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def get_latest_kelly_metrics(
        self,
        symbol: str,
        strategy: str = None
    ) -> Dict[str, Any]:
        """Get the most recent Kelly metrics for a symbol/strategy.
        
        Args:
            symbol: Trading symbol
            strategy: Optional strategy filter
            
        Returns:
            Kelly metrics or None
        """
        with self._get_connection() as conn:
            if strategy:
                row = conn.execute(
                    """SELECT * FROM kelly_metrics
                       WHERE symbol = ? AND strategy = ?
                       ORDER BY timestamp DESC
                       LIMIT 1""",
                    (symbol, strategy)
                ).fetchone()
            else:
                row = conn.execute(
                    """SELECT * FROM kelly_metrics
                       WHERE symbol = ?
                       ORDER BY timestamp DESC
                       LIMIT 1""",
                    (symbol,)
                ).fetchone()
            return dict(row) if row else None
    
    # ========== MULTI-ASSET PORTFOLIO PERSISTENCE METHODS ==========
    
    async def record_multi_asset_portfolio(
        self,
        symbols: List[str],
        total_stake: float,
        individual_stakes: Dict[str, float],
        combined_return: float,
        variance_single_asset: float,
        variance_multi_asset: float,
        variance_reduction_pct: float
    ) -> int:
        """Record multi-asset portfolio metrics.
        
        Args:
            symbols: List of symbols in portfolio
            total_stake: Total capital deployed
            individual_stakes: Dictionary of symbol -> stake amount
            combined_return: Combined portfolio return
            variance_single_asset: Variance if single asset
            variance_multi_asset: Actual multi-asset variance
            variance_reduction_pct: Percentage reduction achieved
            
        Returns:
            Record ID
        """
        import json
        with self._get_connection() as conn:
            cursor = conn.execute(
                """INSERT INTO multi_asset_portfolio
                   (timestamp, symbols, total_stake, individual_stakes,
                    combined_return, variance_single_asset, variance_multi_asset,
                    variance_reduction_pct)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    datetime.now().timestamp(),
                    json.dumps(symbols),
                    total_stake,
                    json.dumps(individual_stakes),
                    combined_return,
                    variance_single_asset,
                    variance_multi_asset,
                    variance_reduction_pct,
                )
            )
            conn.commit()
            return cursor.lastrowid
    
    async def get_variance_reduction_stats(self, days: int = 30) -> Dict[str, Any]:
        """Get variance reduction statistics over time period.
        
        Args:
            days: Number of days to analyze
            
        Returns:
            Statistics dictionary
        """
        cutoff = datetime.now().timestamp() - (days * 24 * 60 * 60)
        with self._get_connection() as conn:
            stats = conn.execute(
                """SELECT
                       COUNT(*) as total_records,
                       AVG(variance_reduction_pct) as avg_reduction,
                       MAX(variance_reduction_pct) as max_reduction,
                       MIN(variance_reduction_pct) as min_reduction
                   FROM multi_asset_portfolio
                   WHERE timestamp >= ?""",
                (cutoff,)
            ).fetchone()
            return dict(stats) if stats else {}
