"""
Price Context Service
Handles price data enrichment for options flow
Uses Yahoo Finance for real-time data instead of database
"""
import pandas as pd
import asyncpg
from datetime import datetime, timedelta
from typing import Dict, Optional
import pytz
from services.yahoo_finance_service import YahooFinanceService
from utils.logger import logger


class PriceContextService:
    """Service for enriching flow data with price context"""
    
    def __init__(self, pool: asyncpg.Pool):
        self.pool = pool
        self.ct_tz = pytz.timezone('America/Chicago')
        self.use_yahoo_finance = False  # Temporarily disabled - focus on Phase 1 button
        self.yahoo_service = YahooFinanceService() if self.use_yahoo_finance else None
    
    def get_session_bucket(self, flow_ts_local: datetime) -> str:
        """Determine session bucket from flow timestamp"""
        if pd.isna(flow_ts_local):
            return 'OFF'
        
        hour = flow_ts_local.hour
        minute = flow_ts_local.minute
        
        if 4 <= hour < 9 or (hour == 9 and minute < 30):
            return 'PRE'
        elif (hour == 9 and minute >= 30) or (9 < hour < 16):
            return 'RTH'
        elif 16 <= hour < 20:
            return 'POST'
        else:
            return 'OFF'
    
    async def get_price_at_time(
        self, 
        symbol: str, 
        timestamp: datetime,
        pool: asyncpg.Pool = None  # Not used anymore, kept for compatibility
    ) -> Optional[Dict]:
        """Get price data at or before a specific timestamp using Yahoo Finance"""
        if not self.use_yahoo_finance or not self.yahoo_service:
            return None  # Yahoo Finance disabled
        try:
            price = self.yahoo_service.get_price_at_time(symbol, timestamp)
            if price:
                return {
                    'close': price,
                    'high': price,  # Approximate
                    'low': price,   # Approximate
                    'volume': None,
                    'ts': timestamp
                }
            return None
        except Exception as e:
            logger.debug(f"Error getting price at time from Yahoo Finance: {e}")
            return None
    
    async def get_rth_open(
        self,
        symbol: str,
        flow_date_cst: datetime.date,
        pool: asyncpg.Pool = None  # Not used anymore, kept for compatibility
    ) -> Optional[Dict]:
        """Get first RTH bar for a given date using Yahoo Finance"""
        if not self.use_yahoo_finance or not self.yahoo_service:
            return None  # Yahoo Finance disabled
        try:
            rth_open_price = self.yahoo_service.get_rth_open(symbol, flow_date_cst)
            if rth_open_price:
                rth_time = self.ct_tz.localize(
                    datetime.combine(flow_date_cst, datetime.min.time().replace(hour=9, minute=30))
                )
                return {
                    'open': rth_open_price,
                    'ts': rth_time
                }
            return None
        except Exception as e:
            logger.debug(f"Error getting RTH open from Yahoo Finance: {e}")
            return None
    
    async def get_prior_close(
        self,
        symbol: str,
        flow_date_cst: datetime.date,
        pool: asyncpg.Pool = None  # Not used anymore, kept for compatibility
    ) -> Optional[float]:
        """Get prior day's close using Yahoo Finance"""
        if not self.use_yahoo_finance or not self.yahoo_service:
            return None  # Yahoo Finance disabled
        try:
            return self.yahoo_service.get_prior_close(symbol, flow_date_cst)
        except Exception as e:
            logger.debug(f"Error getting prior close from Yahoo Finance: {e}")
            return None
    
    async def calculate_vwap_at_time(
        self,
        symbol: str,
        timestamp: datetime,
        pool: asyncpg.Pool = None  # Not used anymore, kept for compatibility
    ) -> Optional[Dict]:
        """
        Calculate VWAP (Volume Weighted Average Price) up to the given timestamp
        for the trading day using Yahoo Finance. VWAP = SUM(price * volume) / SUM(volume)
        """
        try:
            # Convert timestamp to CST if needed
            if timestamp.tzinfo is None:
                timestamp = self.ct_tz.localize(timestamp)
            else:
                timestamp = timestamp.astimezone(self.ct_tz)
            
            # Convert to Eastern Time for market hours check (US market opens at 9:30 AM ET)
            et_tz = pytz.timezone('America/New_York')
            timestamp_et = timestamp.astimezone(et_tz) if timestamp.tzinfo else et_tz.localize(timestamp)
            trade_time = timestamp_et.time()
            
            # Only calculate VWAP if we're at or after RTH open (9:30 AM ET)
            if trade_time < pd.Timestamp('09:30:00').time():
                # Before RTH, return None (no VWAP yet)
                logger.debug(f"Before RTH open (9:30 AM ET): {timestamp_et.strftime('%H:%M:%S %Z')}")
                return None
            
            # Calculate VWAP from RTH open (9:30 AM) to the given timestamp
            if not self.use_yahoo_finance or not self.yahoo_service:
                return None  # Yahoo Finance disabled
            vwap = await self.yahoo_service.calculate_vwap(symbol, timestamp)
            
            if vwap:
                return {
                    'vwap': vwap,
                    'total_volume': None,  # Yahoo Finance doesn't provide this easily
                    'bar_count': None
                }
            
            return None
        except Exception as e:
            logger.debug(f"Error calculating VWAP from Yahoo Finance: {e}")
            return None
    
    async def enrich_flow_with_prices(
        self,
        flow_rows: pd.DataFrame,
        pool: asyncpg.Pool
    ) -> pd.DataFrame:
        """Enrich flow data with price context (optimized with batch queries)"""
        df = flow_rows.copy()
        
        # Add session bucket
        df['session_bucket'] = df['flow_ts_local'].apply(self.get_session_bucket)
        
        if df.empty:
            return df
        
        # Batch fetch all price data using Yahoo Finance
        # Get unique symbols and dates
        unique_symbols = df['symbol_norm'].unique().tolist()
        unique_dates = df['flow_date_cst'].unique().tolist()
        
        # Batch fetch prices at flow times using Yahoo Finance
        # Group by symbol to reduce API calls
        price_data_dict = {}
        for symbol in unique_symbols:
            symbol_flows = df[df['symbol_norm'] == symbol]
            timestamps = symbol_flows['flow_ts_utc'].dropna().unique().tolist()
            
            if not timestamps:
                continue
            
            # For each timestamp, get the latest price <= that timestamp
            # Only fetch unique timestamps per symbol
            for ts in timestamps:
                try:
                    price_data = await self.get_price_at_time(symbol, ts)
                    if price_data:
                        price_data_dict[(symbol.upper(), ts)] = price_data
                    
                    # Also get 5m and 15m ago prices (only if needed)
                    # Skip if we already have this timestamp
                    ts_5m = ts - timedelta(minutes=5)
                    if (symbol.upper(), ts_5m) not in price_data_dict and self.use_yahoo_finance and self.yahoo_service:
                        try:
                            price_5m = self.yahoo_service.get_price_at_time(symbol, ts_5m)
                            if price_5m:
                                price_data_dict[(symbol.upper(), ts_5m)] = {'close': price_5m}
                        except Exception as e:
                            pass
                    
                    ts_15m = ts - timedelta(minutes=15)
                    if (symbol.upper(), ts_15m) not in price_data_dict and self.use_yahoo_finance and self.yahoo_service:
                        try:
                            price_15m = self.yahoo_service.get_price_at_time(symbol, ts_15m)
                            if price_15m:
                                price_data_dict[(symbol.upper(), ts_15m)] = {'close': price_15m}
                        except Exception as e:
                            pass
                except Exception as e:
                    logger.debug(f"Error fetching price for {symbol} at {ts}: {e}")
                    # Add small delay on error to avoid rate limiting
                    import asyncio
                    await asyncio.sleep(0.2)
        
        # Batch fetch RTH opens using Yahoo Finance
        rth_open_dict = {}
        for symbol in unique_symbols:
            for date in unique_dates:
                try:
                    rth_data = await self.get_rth_open(symbol, date)
                    if rth_data:
                        rth_open_dict[(symbol.upper(), date)] = rth_data
                except Exception as e:
                    logger.debug(f"Error fetching RTH open for {symbol} on {date}: {e}")
        
        # Batch fetch prior closes using Yahoo Finance
        prior_close_dict = {}
        for symbol in unique_symbols:
            for date in unique_dates:
                try:
                    prior_close = await self.get_prior_close(symbol, date)
                    if prior_close:
                        prior_close_dict[(symbol.upper(), date)] = prior_close
                except Exception as e:
                    logger.debug(f"Error fetching prior close for {symbol} before {date}: {e}")
        
        # Batch fetch VWAP at signal times using Yahoo Finance
        vwap_dict = {}
        for symbol in unique_symbols:
            symbol_flows = df[df['symbol_norm'] == symbol]
            timestamps = symbol_flows['flow_ts_utc'].dropna().unique().tolist()
            
            for ts in timestamps:
                try:
                    vwap_data = await self.calculate_vwap_at_time(symbol, ts)
                    if vwap_data:
                        vwap_dict[(symbol.upper(), ts)] = vwap_data
                except Exception as e:
                    logger.debug(f"Error calculating VWAP for {symbol} at {ts}: {e}")
            
            # Build price data for each flow row
            price_data = []
            for idx, row in df.iterrows():
                symbol = row['symbol_norm']
                flow_ts_utc = row['flow_ts_utc']
                flow_date_cst = row['flow_date_cst']
                
                # Get price at flow time
                price_at_time = price_data_dict.get((symbol.upper(), flow_ts_utc)) if not pd.isna(flow_ts_utc) else None
                
                # Get RTH open
                rth_open_data = rth_open_dict.get((symbol.upper(), flow_date_cst))
                
                # Get prior close
                prior_close = prior_close_dict.get((symbol.upper(), flow_date_cst))
                
                # Get 5m and 15m ago prices
                price_5m_ago = None
                price_15m_ago = None
                if not pd.isna(flow_ts_utc):
                    ts_5m = flow_ts_utc - timedelta(minutes=5)
                    price_5m_key = (symbol.upper(), ts_5m)
                    price_5m_ago_data = price_data_dict.get(price_5m_key)
                    price_5m_ago = price_5m_ago_data.get('close') if price_5m_ago_data else None
                    
                    ts_15m = flow_ts_utc - timedelta(minutes=15)
                    price_15m_key = (symbol.upper(), ts_15m)
                    price_15m_ago_data = price_data_dict.get(price_15m_key)
                    price_15m_ago = price_15m_ago_data.get('close') if price_15m_ago_data else None
                
                # Get VWAP at signal time
                vwap_data = vwap_dict.get((symbol.upper(), flow_ts_utc)) if not pd.isna(flow_ts_utc) else None
                vwap_at_signal = vwap_data['vwap'] if vwap_data else None
                
                # Calculate price vs VWAP percentage
                price_vs_vwap_pct = None
                if vwap_at_signal and price_at_time and price_at_time.get('close'):
                    try:
                        price_vs_vwap_pct = round(
                            ((price_at_time['close'] - vwap_at_signal) / vwap_at_signal) * 100, 2
                        )
                    except (TypeError, ZeroDivisionError):
                        pass
                
                price_data.append({
                    'u_close': price_at_time['close'] if price_at_time else None,
                    'u_high': price_at_time['high'] if price_at_time else None,
                    'u_low': price_at_time['low'] if price_at_time else None,
                    'u_vol_1m': price_at_time['volume'] if price_at_time else None,
                    'rth_open': rth_open_data['open'] if rth_open_data else None,
                    'prior_close': prior_close,
                    'close_5m_ago': price_5m_ago,
                    'close_15m_ago': price_15m_ago,
                    'vwap_at_signal': vwap_at_signal,
                    'price_vs_vwap_pct': price_vs_vwap_pct,
                })
        
        # Merge price data
        price_df = pd.DataFrame(price_data, index=df.index)
        df = pd.concat([df, price_df], axis=1)
        
        # Calculate price features
        def calc_pct_vs_prior_close(row):
            if row['prior_close'] and row['u_close'] and row['prior_close'] != 0:
                try:
                    return round(((row['u_close'] - row['prior_close']) * 100.0 / row['prior_close']), 2)
                except (TypeError, ValueError):
                    return None
            return None
        
        df['pct_vs_prior_close'] = df.apply(calc_pct_vs_prior_close, axis=1)
        
        def calc_pct_vs_rth_open(row):
            if (row['rth_open'] and row['u_close'] and 
                row['session_bucket'] == 'RTH' and row['rth_open'] != 0):
                try:
                    return round(((row['u_close'] - row['rth_open']) * 100.0 / row['rth_open']), 2)
                except (TypeError, ValueError):
                    return None
            return None
        
        df['pct_vs_rth_open'] = df.apply(calc_pct_vs_rth_open, axis=1)
        
        def calc_pct_5m_momo(row):
            if row['close_5m_ago'] and row['u_close'] and row['close_5m_ago'] != 0:
                try:
                    return round(((row['u_close'] - row['close_5m_ago']) * 100.0 / row['close_5m_ago']), 2)
                except (TypeError, ValueError):
                    return None
            return None
        
        df['pct_5m_momo'] = df.apply(calc_pct_5m_momo, axis=1)
        
        def calc_pct_15m_momo(row):
            if row['close_15m_ago'] and row['u_close'] and row['close_15m_ago'] != 0:
                try:
                    return round(((row['u_close'] - row['close_15m_ago']) * 100.0 / row['close_15m_ago']), 2)
                except (TypeError, ValueError):
                    return None
            return None
        
        df['pct_15m_momo'] = df.apply(calc_pct_15m_momo, axis=1)
        
        # Calculate tape alignment
        df['tape_alignment'] = df.apply(
            lambda row: self.calculate_tape_alignment(row),
            axis=1
        )
        
        return df
    
    def calculate_tape_alignment(self, row) -> int:
        """Calculate tape alignment based on direction and price movement"""
        session = row.get('session_bucket', '')
        direction = row.get('direction', '')
        tol_pct = 0.20  # tolerance percentage
        
        if session == 'PRE' and row.get('pct_vs_prior_close') is not None:
            pct = row['pct_vs_prior_close']
            if (direction == 'BULL' and pct >= tol_pct) or \
               (direction == 'BEAR' and pct <= -tol_pct):
                return 1
        elif session == 'RTH' and row.get('pct_vs_rth_open') is not None:
            pct = row['pct_vs_rth_open']
            if (direction == 'BULL' and pct >= tol_pct) or \
               (direction == 'BEAR' and pct <= -tol_pct):
                return 1
        elif session == 'POST' and row.get('pct_vs_prior_close') is not None:
            pct = row['pct_vs_prior_close']
            if (direction == 'BULL' and pct >= tol_pct) or \
               (direction == 'BEAR' and pct <= -tol_pct):
                return 1
        
        return 0

