"""
Database Access Module for ML Engine
=====================================
Provides PostgreSQL access with the same interface as the legacy MySQL access.

This module replaces the MySQL-based data access with PostgreSQL,
allowing training scripts to work with the local trading_platform database.

Usage:
    from data.database import DatabaseManager, PostgreSQLConnection

    # Simple usage
    db = PostgreSQLConnection()
    df = db.get_ticker_data('XAUUSD', limit=50000)

    # Via DatabaseManager
    manager = DatabaseManager()
    df = manager.db.get_ticker_data('XAUUSD')

Author: ML-Specialist (NEXUS v4.0)
Created: 2026-01-25
"""

import os
import yaml
from pathlib import Path
from typing import Optional, Dict, Any
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.pool import QueuePool
from loguru import logger


class PostgreSQLConnection:
    """
    PostgreSQL connection for ML Engine.

    Provides the same interface as the legacy MySQLConnection
    but uses PostgreSQL with the local trading_platform database.
    """

    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize PostgreSQL connection.

        Args:
            config_path: Path to database.yaml config file
        """
        self.config = self._load_config(config_path)
        self.engine = self._create_engine()
        logger.info(f"Connected to PostgreSQL: {self.config['host']}:{self.config['port']}/{self.config['database']}")

    def _load_config(self, config_path: Optional[str] = None) -> Dict[str, Any]:
        """Load configuration from YAML or environment variables."""
        # First try environment variables
        config = {
            'host': os.getenv('DB_HOST', 'localhost'),
            'port': int(os.getenv('DB_PORT', '5432')),
            'database': os.getenv('DB_NAME', 'trading_platform'),
            'user': os.getenv('DB_USER', 'trading_user'),
            'password': os.getenv('DB_PASSWORD', 'trading_dev_2026'),
            'pool_size': int(os.getenv('DB_POOL_SIZE', '10')),
        }

        # Try to load from YAML config
        if config_path:
            yaml_path = Path(config_path)
        else:
            yaml_path = Path(__file__).parent.parent.parent / 'config' / 'database.yaml'

        if yaml_path.exists():
            with open(yaml_path, 'r') as f:
                yaml_config = yaml.safe_load(f)

            if 'postgres' in yaml_config:
                pg_config = yaml_config['postgres']
                # Override with YAML values (expand env vars)
                for key in ['host', 'port', 'database', 'user', 'password']:
                    if key in pg_config:
                        value = str(pg_config[key])
                        # Expand environment variable references
                        if value.startswith('${') and '}' in value:
                            env_var = value[2:value.index('}')]
                            default = value[value.index(':-')+2:-1] if ':-' in value else None
                            config[key] = os.getenv(env_var, default) or config[key]
                        else:
                            config[key] = value if key != 'port' else int(value)

        return config

    def _create_engine(self):
        """Create SQLAlchemy engine for PostgreSQL."""
        connection_string = (
            f"postgresql://{self.config['user']}:{self.config['password']}"
            f"@{self.config['host']}:{self.config['port']}/{self.config['database']}"
        )

        return create_engine(
            connection_string,
            poolclass=QueuePool,
            pool_size=self.config.get('pool_size', 10),
            max_overflow=20,
            pool_pre_ping=True,
            echo=False
        )

    def get_ticker_data(
        self,
        symbol: str,
        timeframe: str = '5m',
        limit: int = 500000,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
    ) -> pd.DataFrame:
        """
        Get OHLCV data for a ticker.

        Args:
            symbol: Trading symbol (e.g., 'XAUUSD', 'EURUSD')
            timeframe: Timeframe ('5m', '15m', '1h', '4h', 'd')
            limit: Maximum number of records
            start_date: Start date filter (YYYY-MM-DD)
            end_date: End date filter (YYYY-MM-DD)

        Returns:
            DataFrame with OHLCV data indexed by timestamp
        """
        # Determine table based on timeframe
        table_map = {
            '5m': 'ohlcv_5m',
            '15m': 'ohlcv_15m',
            '1h': 'ohlcv_1h',
            '4h': 'ohlcv_4h',
            'd': 'ohlcv_daily',
            '1d': 'ohlcv_daily',
        }

        table = table_map.get(timeframe.lower(), 'ohlcv_5m')

        # Normalize symbol (remove Polygon prefixes)
        clean_symbol = symbol
        if symbol.startswith('C:') or symbol.startswith('X:') or symbol.startswith('I:'):
            clean_symbol = symbol[2:]

        # Build query
        query = f"""
        SELECT
            o.timestamp,
            o.open,
            o.high,
            o.low,
            o.close,
            o.volume,
            o.vwap
        FROM market_data.{table} o
        JOIN market_data.tickers t ON t.id = o.ticker_id
        WHERE UPPER(t.symbol) = UPPER(:symbol)
        """

        params = {'symbol': clean_symbol}

        if start_date:
            query += " AND o.timestamp >= :start_date"
            params['start_date'] = start_date

        if end_date:
            query += " AND o.timestamp <= :end_date"
            params['end_date'] = end_date

        query += f" ORDER BY o.timestamp ASC LIMIT {limit}"

        try:
            df = pd.read_sql(text(query), self.engine, params=params)

            if df.empty:
                logger.warning(f"No data found for {symbol} ({timeframe})")
                return df

            # Set timestamp as index
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            df.set_index('timestamp', inplace=True)

            logger.info(f"Loaded {len(df)} bars for {symbol} ({timeframe})")
            return df

        except Exception as e:
            logger.error(f"Error loading data for {symbol}: {e}")
            return pd.DataFrame()

    def get_all_tickers(self) -> pd.DataFrame:
        """Get all active tickers."""
        query = """
        SELECT id, symbol, asset_type, polygon_ticker, is_active
        FROM market_data.tickers
        WHERE is_active = true
        ORDER BY symbol
        """

        return pd.read_sql(text(query), self.engine)

    def execute(self, query: str, params: Optional[Dict] = None) -> pd.DataFrame:
        """Execute a raw SQL query."""
        return pd.read_sql(text(query), self.engine, params=params or {})

    def execute_query(self, query: str, params: Optional[Dict] = None) -> pd.DataFrame:
        """
        Execute a raw SQL query with MySQL-to-PostgreSQL translation.

        This method handles queries written for the old MySQL schema
        and translates them to work with the PostgreSQL schema.

        Args:
            query: SQL query (may use old MySQL table/column names)
            params: Query parameters

        Returns:
            DataFrame with query results
        """
        # Translate MySQL table/column names to PostgreSQL
        translated_query = self._translate_mysql_query(query)

        # Handle parameter format differences
        translated_params = self._translate_params(params)

        try:
            return pd.read_sql(text(translated_query), self.engine, params=translated_params)
        except Exception as e:
            logger.error(f"Query error: {e}")
            logger.debug(f"Original query: {query}")
            logger.debug(f"Translated query: {translated_query}")
            raise

    def _translate_mysql_query(self, query: str) -> str:
        """
        Translate MySQL query to PostgreSQL query.

        Handles:
        - tickers_agg_data -> market_data.ohlcv_5m with JOIN
        - date_agg -> timestamp
        - ticker column -> JOIN with tickers table
        """
        import re

        # Check if query uses the old MySQL table
        if 'tickers_agg_data' in query.lower():
            # This is a MySQL-style query that needs translation
            # Replace table and column references
            translated = query

            # Replace table name with subquery that joins properly
            # The old query selects from tickers_agg_data WHERE ticker = :symbol
            # We need to translate to: FROM market_data.ohlcv_5m o JOIN market_data.tickers t ON ...

            # Extract what columns are being selected
            # For now, use a simple approach: replace the FROM clause

            # Check if it's a simple SELECT ... FROM tickers_agg_data WHERE ticker = ...
            if re.search(r'FROM\s+tickers_agg_data', query, re.IGNORECASE):
                # Build the translated query
                translated = re.sub(
                    r'FROM\s+tickers_agg_data',
                    '''FROM (
                        SELECT
                            o.timestamp as date_agg,
                            o.open,
                            o.high,
                            o.low,
                            o.close,
                            o.volume,
                            o.vwap,
                            t.polygon_ticker as ticker
                        FROM market_data.ohlcv_5m o
                        JOIN market_data.tickers t ON t.id = o.ticker_id
                    ) AS tickers_agg_data''',
                    translated,
                    flags=re.IGNORECASE
                )

            return translated

        return query

    def _translate_params(self, params: Optional[Dict]) -> Optional[Dict]:
        """Translate parameter format if needed."""
        if not params:
            return params

        # Parameters should work as-is for most cases
        return params


# Aliases for backward compatibility
MySQLConnection = PostgreSQLConnection
DatabaseConnection = PostgreSQLConnection


class DatabaseManager:
    """
    Database manager for ML Engine.

    Provides centralized access to database connections.
    Maintains backward compatibility with existing training scripts.
    """

    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize database manager.

        Args:
            config_path: Path to database configuration
        """
        self.db = PostgreSQLConnection(config_path)
        self._engine = self.db.engine

    @property
    def engine(self):
        """Get SQLAlchemy engine."""
        return self._engine

    def get_ticker_data(
        self,
        symbol: str,
        timeframe: str = '5m',
        limit: int = 500000
    ) -> pd.DataFrame:
        """Get ticker data - delegates to db connection."""
        return self.db.get_ticker_data(symbol, timeframe, limit)


def load_ohlcv_from_postgres(
    symbol: str,
    timeframe: str = '5m',
    start_date: Optional[str] = None,
    end_date: Optional[str] = None
) -> pd.DataFrame:
    """
    Load OHLCV data from PostgreSQL.

    Drop-in replacement for load_ohlcv_from_mysql functions.

    Args:
        symbol: Trading symbol (e.g., 'XAUUSD', 'BTCUSD')
        timeframe: Timeframe ('5m', '15m')
        start_date: Start date (YYYY-MM-DD)
        end_date: End date (YYYY-MM-DD)

    Returns:
        DataFrame with OHLCV data
    """
    db = PostgreSQLConnection()
    return db.get_ticker_data(
        symbol=symbol,
        timeframe=timeframe,
        start_date=start_date,
        end_date=end_date
    )
