"""
Price Context Service
Handles price data enrichment for options flow
"""
import pandas as pd
import asyncpg
from datetime import datetime, timedelta
from typing import Dict, Optional
import pytz


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')
    
    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
    ) -> Optional[Dict]:
        """Get price data at or before a specific timestamp"""
        async with pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT close, high, low, volume, ts
                FROM prices_intraday_1m
                WHERE UPPER(symbol) = $1
                  AND ts <= $2
                ORDER BY ts DESC
                LIMIT 1
            """, symbol.upper(), timestamp)
            
            if row:
                return dict(row)
            return None
    
    async def get_rth_open(
        self,
        symbol: str,
        flow_date_cst: datetime.date,
        pool: asyncpg.Pool
    ) -> Optional[Dict]:
        """Get first RTH bar for a given date"""
        async with pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT open, ts
                FROM prices_intraday_1m
                WHERE UPPER(symbol) = $1
                  AND (timezone('America/Chicago', ts))::date = $2
                  AND (timezone('America/Chicago', ts))::time >= time '09:30:00'
                ORDER BY ts ASC
                LIMIT 1
            """, symbol.upper(), flow_date_cst)
            
            if row:
                return dict(row)
            return None
    
    async def get_prior_close(
        self,
        symbol: str,
        flow_date_cst: datetime.date,
        pool: asyncpg.Pool
    ) -> Optional[float]:
        """Get prior day's close"""
        async with pool.acquire() as conn:
            prior_date = flow_date_cst - timedelta(days=1)
            row = await conn.fetchrow("""
                SELECT close
                FROM prices_daily
                WHERE UPPER(symbol) = $1
                  AND "Date" = $2
                ORDER BY "Date" DESC
                LIMIT 1
            """, symbol.upper(), prior_date)
            
            if row:
                return row['close']
            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
        async with pool.acquire() as conn:
            # 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 (for each symbol, get prices <= each timestamp)
            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
                for ts in timestamps:
                    row = await conn.fetchrow("""
                        SELECT close, high, low, volume, ts
                        FROM prices_intraday_1m
                        WHERE UPPER(symbol) = $1
                          AND ts <= $2
                        ORDER BY ts DESC
                        LIMIT 1
                    """, symbol.upper(), ts)
                    
                    if row:
                        price_data_dict[(symbol.upper(), ts)] = dict(row)
                    
                    # Also get 5m and 15m ago prices (with error handling)
                    try:
                        ts_5m = ts - timedelta(minutes=5)
                        row_5m = await conn.fetchrow("""
                            SELECT close
                            FROM prices_intraday_1m
                            WHERE UPPER(symbol) = $1
                              AND ts <= $2
                            ORDER BY ts DESC
                            LIMIT 1
                        """, symbol.upper(), ts_5m)
                        
                        if row_5m:
                            price_data_dict[(symbol.upper(), ts_5m)] = {'close': row_5m['close']}
                    except Exception as e:
                        # Log but don't fail on 5m price lookup
                        pass
                    
                    try:
                        ts_15m = ts - timedelta(minutes=15)
                        row_15m = await conn.fetchrow("""
                            SELECT close
                            FROM prices_intraday_1m
                            WHERE UPPER(symbol) = $1
                              AND ts <= $2
                            ORDER BY ts DESC
                            LIMIT 1
                        """, symbol.upper(), ts_15m)
                        
                        if row_15m:
                            price_data_dict[(symbol.upper(), ts_15m)] = {'close': row_15m['close']}
                    except Exception as e:
                        # Log but don't fail on 15m price lookup
                        pass
            
            # Batch fetch RTH opens
            rth_open_dict = {}
            for symbol in unique_symbols:
                for date in unique_dates:
                    row = await conn.fetchrow("""
                        SELECT open, ts
                        FROM prices_intraday_1m
                        WHERE UPPER(symbol) = $1
                          AND (timezone('America/Chicago', ts))::date = $2
                          AND (timezone('America/Chicago', ts))::time >= time '09:30:00'
                        ORDER BY ts ASC
                        LIMIT 1
                    """, symbol.upper(), date)
                    
                    if row:
                        rth_open_dict[(symbol.upper(), date)] = dict(row)
            
            # Batch fetch prior closes
            prior_close_dict = {}
            for symbol in unique_symbols:
                for date in unique_dates:
                    # Handle date conversion
                    if isinstance(date, datetime):
                        prior_date = (date - timedelta(days=1)).date()
                    elif isinstance(date, pd.Timestamp):
                        prior_date = (date - timedelta(days=1)).date()
                    else:
                        # Assume it's already a date object
                        prior_date = date - timedelta(days=1) if hasattr(date, '__sub__') else date
                    
                    row = await conn.fetchrow("""
                        SELECT close
                        FROM prices_daily
                        WHERE UPPER(symbol) = $1
                          AND "Date" = $2
                        ORDER BY "Date" DESC
                        LIMIT 1
                    """, symbol.upper(), prior_date)
                    
                    if row:
                        prior_close_dict[(symbol.upper(), date)] = row['close']
            
            # 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
                
                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,
                })
        
        # 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