"""
Database optimization module for parallel backtesting performance.
Implements query optimization, bulk operations, and performance monitoring.
"""

import logging
import time
from contextlib import contextmanager
from typing import Any

import pandas as pd
from sqlalchemy import Index, event, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session

from maverick_mcp.data.models import (
    PriceCache,
    SessionLocal,
    Stock,
    engine,
)

logger = logging.getLogger(__name__)


class QueryOptimizer:
    """Database query optimization for backtesting performance."""

    def __init__(self, session_factory=None):
        """Initialize query optimizer."""
        self.session_factory = session_factory or SessionLocal
        self._query_stats = {}
        self._connection_pool_stats = {
            "active_connections": 0,
            "checked_out": 0,
            "total_queries": 0,
            "slow_queries": 0,
        }

    def create_backtesting_indexes(self, engine: Engine):
        """
        Create optimized indexes for backtesting queries.

        These indexes are specifically designed for the parallel backtesting
        workload patterns.
        """
        logger.info("Creating backtesting optimization indexes...")

        # Define additional indexes for common backtesting query patterns
        additional_indexes = [
            # Composite index for date range queries with symbol lookup
            Index(
                "mcp_price_cache_symbol_date_range_idx",
                Stock.__table__.c.ticker_symbol,
                PriceCache.__table__.c.date,
                PriceCache.__table__.c.close_price,
            ),
            # Index for volume-based queries (common in strategy analysis)
            Index(
                "mcp_price_cache_volume_date_idx",
                PriceCache.__table__.c.volume,
                PriceCache.__table__.c.date,
            ),
            # Covering index for OHLCV queries (includes all price data)
            Index(
                "mcp_price_cache_ohlcv_covering_idx",
                PriceCache.__table__.c.stock_id,
                PriceCache.__table__.c.date,
                # Include all price columns as covering columns
                PriceCache.__table__.c.open_price,
                PriceCache.__table__.c.high_price,
                PriceCache.__table__.c.low_price,
                PriceCache.__table__.c.close_price,
                PriceCache.__table__.c.volume,
            ),
            # Index for latest price queries
            Index(
                "mcp_price_cache_latest_price_idx",
                PriceCache.__table__.c.stock_id,
                PriceCache.__table__.c.date.desc(),
            ),
            # Partial index for recent data (last 2 years) - most commonly queried
            # Note: This is PostgreSQL-specific, will be skipped for SQLite
        ]

        try:
            with engine.connect() as conn:
                # Check if we're using PostgreSQL for partial indexes
                is_postgresql = engine.dialect.name == "postgresql"

                for index in additional_indexes:
                    try:
                        # Skip PostgreSQL-specific features on SQLite
                        if not is_postgresql and "partial" in str(index).lower():
                            continue

                        # Create index if it doesn't exist
                        index.create(conn, checkfirst=True)
                        logger.info(f"Created index: {index.name}")

                    except Exception as e:
                        logger.warning(f"Failed to create index {index.name}: {e}")

                # Add PostgreSQL-specific optimizations
                if is_postgresql:
                    try:
                        # Create partial index for recent data (last 2 years)
                        conn.execute(
                            text("""
                            CREATE INDEX CONCURRENTLY IF NOT EXISTS mcp_price_cache_recent_data_idx
                            ON mcp_price_cache (stock_id, date DESC, close_price)
                            WHERE date >= CURRENT_DATE - INTERVAL '2 years'
                        """)
                        )
                        logger.info("Created partial index for recent data")

                        # Update table statistics for better query planning
                        conn.execute(text("ANALYZE mcp_price_cache"))
                        conn.execute(text("ANALYZE mcp_stocks"))
                        logger.info("Updated table statistics")

                    except Exception as e:
                        logger.warning(
                            f"Failed to create PostgreSQL optimizations: {e}"
                        )

                conn.commit()

        except Exception as e:
            logger.error(f"Failed to create backtesting indexes: {e}")

    def optimize_connection_pool(self, engine: Engine):
        """Optimize connection pool settings for parallel operations."""
        logger.info("Optimizing connection pool for parallel backtesting...")

        # Add connection pool event listeners for monitoring
        @event.listens_for(engine, "connect")
        def receive_connect(dbapi_connection, connection_record):
            self._connection_pool_stats["active_connections"] += 1

        @event.listens_for(engine, "checkout")
        def receive_checkout(dbapi_connection, connection_record, connection_proxy):
            self._connection_pool_stats["checked_out"] += 1

        @event.listens_for(engine, "checkin")
        def receive_checkin(dbapi_connection, connection_record):
            self._connection_pool_stats["checked_out"] -= 1

    def create_bulk_insert_method(self):
        """Create optimized bulk insert method for price data."""

        def bulk_insert_price_data_optimized(
            session: Session,
            price_data_list: list[dict[str, Any]],
            batch_size: int = 1000,
        ):
            """
            Optimized bulk insert for price data with batching.

            Args:
                session: Database session
                price_data_list: List of price data dictionaries
                batch_size: Number of records per batch
            """
            if not price_data_list:
                return

            logger.info(f"Bulk inserting {len(price_data_list)} price records")
            start_time = time.time()

            try:
                # Process in batches to avoid memory issues
                for i in range(0, len(price_data_list), batch_size):
                    batch = price_data_list[i : i + batch_size]

                    # Use bulk_insert_mappings for better performance
                    session.bulk_insert_mappings(PriceCache, batch)

                    # Commit each batch to free up memory
                    if i + batch_size < len(price_data_list):
                        session.flush()

                session.commit()

                elapsed = time.time() - start_time
                logger.info(
                    f"Bulk insert completed in {elapsed:.2f}s "
                    f"({len(price_data_list) / elapsed:.0f} records/sec)"
                )

            except Exception as e:
                logger.error(f"Bulk insert failed: {e}")
                session.rollback()
                raise

        return bulk_insert_price_data_optimized

    @contextmanager
    def query_performance_monitor(self, query_name: str):
        """Context manager for monitoring query performance."""
        start_time = time.time()

        try:
            yield
        finally:
            elapsed = time.time() - start_time

            # Track query statistics
            if query_name not in self._query_stats:
                self._query_stats[query_name] = {
                    "count": 0,
                    "total_time": 0.0,
                    "avg_time": 0.0,
                    "max_time": 0.0,
                    "slow_queries": 0,
                }

            stats = self._query_stats[query_name]
            stats["count"] += 1
            stats["total_time"] += elapsed
            stats["avg_time"] = stats["total_time"] / stats["count"]
            stats["max_time"] = max(stats["max_time"], elapsed)

            # Mark slow queries (> 1 second)
            if elapsed > 1.0:
                stats["slow_queries"] += 1
                self._connection_pool_stats["slow_queries"] += 1
                logger.warning(f"Slow query detected: {query_name} took {elapsed:.2f}s")

            self._connection_pool_stats["total_queries"] += 1

    def get_optimized_price_query(self) -> str:
        """Get optimized SQL query for price data retrieval."""
        return """
        SELECT
            pc.date,
            pc.open_price as "open",
            pc.high_price as "high",
            pc.low_price as "low",
            pc.close_price as "close",
            pc.volume
        FROM mcp_price_cache pc
        JOIN mcp_stocks s ON pc.stock_id = s.stock_id
        WHERE s.ticker_symbol = :symbol
        AND pc.date >= :start_date
        AND pc.date <= :end_date
        ORDER BY pc.date
        """

    def get_batch_price_query(self) -> str:
        """Get optimized SQL query for batch price data retrieval."""
        return """
        SELECT
            s.ticker_symbol,
            pc.date,
            pc.open_price as "open",
            pc.high_price as "high",
            pc.low_price as "low",
            pc.close_price as "close",
            pc.volume
        FROM mcp_price_cache pc
        JOIN mcp_stocks s ON pc.stock_id = s.stock_id
        WHERE s.ticker_symbol = ANY(:symbols)
        AND pc.date >= :start_date
        AND pc.date <= :end_date
        ORDER BY s.ticker_symbol, pc.date
        """

    def execute_optimized_query(
        self,
        session: Session,
        query: str,
        params: dict[str, Any],
        query_name: str = "unnamed",
    ) -> pd.DataFrame:
        """Execute optimized query with performance monitoring."""
        with self.query_performance_monitor(query_name):
            try:
                result = pd.read_sql(
                    text(query),
                    session.bind,
                    params=params,
                    index_col="date" if "date" in query.lower() else None,
                    parse_dates=["date"] if "date" in query.lower() else None,
                )

                logger.debug(f"Query {query_name} returned {len(result)} rows")
                return result

            except Exception as e:
                logger.error(f"Query {query_name} failed: {e}")
                raise

    def get_statistics(self) -> dict[str, Any]:
        """Get query and connection pool statistics."""
        return {
            "query_stats": self._query_stats.copy(),
            "connection_pool_stats": self._connection_pool_stats.copy(),
            "top_slow_queries": sorted(
                [
                    (name, stats["avg_time"])
                    for name, stats in self._query_stats.items()
                ],
                key=lambda x: x[1],
                reverse=True,
            )[:5],
        }

    def reset_statistics(self):
        """Reset performance statistics."""
        self._query_stats.clear()
        self._connection_pool_stats = {
            "active_connections": 0,
            "checked_out": 0,
            "total_queries": 0,
            "slow_queries": 0,
        }


class BatchQueryExecutor:
    """Efficient batch query execution for parallel backtesting."""

    def __init__(self, optimizer: QueryOptimizer = None):
        """Initialize batch query executor."""
        self.optimizer = optimizer or QueryOptimizer()

    async def fetch_multiple_symbols_data(
        self,
        symbols: list[str],
        start_date: str,
        end_date: str,
        session: Session = None,
    ) -> dict[str, pd.DataFrame]:
        """
        Efficiently fetch data for multiple symbols in a single query.

        Args:
            symbols: List of stock symbols
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)
            session: Optional database session

        Returns:
            Dictionary mapping symbols to DataFrames
        """
        if not symbols:
            return {}

        should_close = session is None
        if session is None:
            session = self.optimizer.session_factory()

        try:
            # Use batch query to fetch all symbols at once
            batch_query = self.optimizer.get_batch_price_query()

            result_df = self.optimizer.execute_optimized_query(
                session=session,
                query=batch_query,
                params={
                    "symbols": symbols,
                    "start_date": start_date,
                    "end_date": end_date,
                },
                query_name="batch_symbol_fetch",
            )

            # Group by symbol and create separate DataFrames
            symbol_data = {}
            if not result_df.empty:
                for symbol in symbols:
                    symbol_df = result_df[result_df["ticker_symbol"] == symbol].copy()
                    symbol_df.drop("ticker_symbol", axis=1, inplace=True)
                    symbol_data[symbol] = symbol_df
            else:
                # Return empty DataFrames for all symbols
                symbol_data = {symbol: pd.DataFrame() for symbol in symbols}

            logger.info(
                f"Batch fetched data for {len(symbols)} symbols: "
                f"{sum(len(df) for df in symbol_data.values())} total records"
            )

            return symbol_data

        finally:
            if should_close:
                session.close()


# Global instances for easy access
_query_optimizer = QueryOptimizer()
_batch_executor = BatchQueryExecutor(_query_optimizer)


def get_query_optimizer() -> QueryOptimizer:
    """Get the global query optimizer instance."""
    return _query_optimizer


def get_batch_executor() -> BatchQueryExecutor:
    """Get the global batch executor instance."""
    return _batch_executor


def initialize_database_optimizations():
    """Initialize all database optimizations for backtesting."""
    logger.info("Initializing database optimizations for parallel backtesting...")

    try:
        optimizer = get_query_optimizer()

        # Create performance indexes
        optimizer.create_backtesting_indexes(engine)

        # Optimize connection pool
        optimizer.optimize_connection_pool(engine)

        logger.info("Database optimizations initialized successfully")

    except Exception as e:
        logger.error(f"Failed to initialize database optimizations: {e}")


@contextmanager
def optimized_db_session():
    """Context manager for optimized database session."""
    session = SessionLocal()
    try:
        # Configure session for optimal performance
        session.execute(text("PRAGMA synchronous = NORMAL"))  # SQLite optimization
        session.execute(text("PRAGMA journal_mode = WAL"))  # SQLite optimization
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


# Performance monitoring decorator
def monitor_query_performance(query_name: str):
    """Decorator for monitoring query performance."""

    def decorator(func):
        def wrapper(*args, **kwargs):
            optimizer = get_query_optimizer()
            with optimizer.query_performance_monitor(query_name):
                return func(*args, **kwargs)

        return wrapper

    return decorator
