"""
Data Loader for Cryptocurrency Market Data from SQLite Databases.
"""

import os
import logging
import pandas as pd
import numpy as np
import sqlite3
import glob
import re
import sys
from datetime import datetime, timedelta
from typing import List, Optional

# Import IOManager (adjust path if necessary)
from .io_manager import IOManager
from omegaconf import DictConfig # Assuming OmegaConf for config type hint

logger = logging.getLogger(__name__)

# Define helper function (outside the class is fine)
def sample_by_volatility(df: pd.DataFrame, vol_window: int = 30, vol_quantile: float = 0.5) -> pd.Series:
    """
    Creates a boolean mask to sample data points based on rolling volatility.

    Keeps data points where the rolling volatility is above a specified quantile.

    Args:
        df (pd.DataFrame): DataFrame with a 'close' column and DatetimeIndex.
        vol_window (int): Rolling window size for volatility calculation.
        vol_quantile (float): Quantile threshold (0.0 to 1.0). Data with volatility
                              above this quantile will be kept.

    Returns:
        pd.Series: Boolean mask, True for rows to keep.
    """
    if 'close' not in df.columns:
        raise ValueError("DataFrame must contain a 'close' column.")
    if not isinstance(df.index, pd.DatetimeIndex):
        raise ValueError("DataFrame must have a DatetimeIndex.")
    if vol_window <= 1:
        raise ValueError("vol_window must be greater than 1.")

    # Calculate rolling volatility (std dev of returns)
    returns = df['close'].pct_change()
    # Use min_periods=vol_window//2+1 to get some values earlier, but still require significant data
    rolling_vol = returns.rolling(window=vol_window, min_periods=max(2, vol_window // 2 + 1)).std()

    # Calculate the threshold volatility value
    threshold_vol = rolling_vol.quantile(vol_quantile)

    if pd.isna(threshold_vol) or threshold_vol == 0:
        logger.warning(f"Volatility quantile ({vol_quantile}) is NaN or zero ({threshold_vol}). "
                       f"Threshold calculated over {rolling_vol.count()} non-NaN values. "
                       f"Disabling volatility sampling for this segment.")
        # Return a mask that keeps all data if threshold is problematic
        return pd.Series(True, index=df.index)

    # Create the mask
    mask = rolling_vol > threshold_vol

    # Handle initial NaNs from rolling calculation - discard these rows.
    mask.fillna(False, inplace=True) # Discard rows where rolling vol couldn't be calculated

    logger.info(f"Volatility sampling: window={vol_window}, quantile={vol_quantile}, "
                f"threshold={threshold_vol:.6f}. Keeping {mask.sum()} / {len(df)} rows.")

    return mask

class DataLoader:
    """
    Loads historical cryptocurrency market data from SQLite databases.
    Combines functionality from the previous CryptoDBFetcher and data loading logic.
    """
    def __init__(self, db_dir: str, cache_dir: str = "data/cache", use_cache: bool = False):
        """
        Initialize the DataLoader.

        Args:
            db_dir (str): Directory where SQLite database files are stored. Should be an absolute path or resolvable relative to the execution context.
            cache_dir (str): Directory to store cached data (currently not implemented).
            use_cache (bool): Whether to use cached data (currently not implemented).
        """
        # The path resolution should happen *before* calling the DataLoader.
        # We expect db_dir to be the correct path already.
        self.db_dir = os.path.abspath(db_dir) # Ensure absolute path for consistency

        self.cache_dir = cache_dir # Placeholder for future cache implementation
        self.use_cache = use_cache   # Placeholder

        self._db_files = None # Cache discovered DB files

        logger.info(f"Initialized DataLoader with db_dir='{self.db_dir}'")
        if not os.path.exists(self.db_dir):
            # Log a warning, but allow continuation in case the directory is created later
            # or if only specific file checks are relevant later.
            logger.warning(f"Database directory may not exist or is inaccessible: {self.db_dir}")

    def _get_db_files(self) -> List[str]:
        """Get available database files, sorted by date desc (cached). Uses recursive glob."""
        if self._db_files is not None:
            return self._db_files

        logger.info(f"Scanning for DB files recursively in: {self.db_dir}")
        # Check existence *here* right before scanning
        if not os.path.isdir(self.db_dir): # More specific check for directory
            logger.error(f"Database directory does not exist or is not a directory: {self.db_dir}")
            self._db_files = []
            return []

        patterns = ["*.mktdata.ohlcv.db", "*.db", "*.sqlite", "*.sqlite3"]
        db_files = []
        for pattern in patterns:
            # Recursive search
            recursive_pattern = os.path.join(self.db_dir, '**', pattern)
            try:
                files = glob.glob(recursive_pattern, recursive=True)
                if files:
                    logger.debug(f"Found {len(files)} files recursively with pattern '{pattern}'")
                    db_files.extend(files)
            except Exception as e:
                logger.error(f"Error during glob pattern '{recursive_pattern}': {e}")

        if not db_files:
            logger.warning(f"No database files found in '{self.db_dir}' matching patterns: {patterns}")
            self._db_files = []
            return []

        db_files = sorted(list(set(db_files))) # Remove duplicates and sort alphabetically for consistency before date sort

        # Sort by date (newest first) if possible
        date_pattern = re.compile(r'(\d{8})')
        file_dates = []
        for file in db_files:
            basename = os.path.basename(file)
            match = date_pattern.search(basename)
            date_obj = None
            if match:
                try:
                    date_obj = pd.to_datetime(match.group(1), format='%Y%m%d')
                except ValueError:
                    pass
            # Fallback: try modification time
            if date_obj is None:
                 try:
                     date_obj = pd.to_datetime(os.path.getmtime(file), unit='s')
                 except Exception:
                     date_obj = pd.Timestamp.min # Default to oldest if error
            file_dates.append((date_obj, file))

        # Sort by date object, newest first
        file_dates.sort(key=lambda x: x[0], reverse=True)
        self._db_files = [file for _, file in file_dates]

        logger.info(f"Found {len(self._db_files)} DB files. Using newest: {os.path.basename(self._db_files[0]) if self._db_files else 'None'}")
        return self._db_files

    def _get_relevant_db_files(self, start_dt: pd.Timestamp, end_dt: pd.Timestamp) -> List[str]:
        """Find DB files potentially containing data for the date range."""
        all_files = self._get_db_files()
        relevant_files = set()
        date_pattern = re.compile(r'(\d{8})')

        start_date_only = start_dt.date()
        end_date_only = end_dt.date()

        for file in all_files:
            basename = os.path.basename(file)
            match = date_pattern.search(basename)
            file_date = None
            if match:
                try:
                    file_date = pd.to_datetime(match.group(1), format='%Y%m%d').date()
                except ValueError:
                    pass # Ignore files with unparseable dates in name

            # Strategy 1: Check if filename date is within the requested range
            if file_date and start_date_only <= file_date <= end_date_only:
                relevant_files.add(file)
                continue # Found by filename date, no need to check mtime

            # Strategy 2: Check if file modification time falls within the range (less precise)
            # Useful for files without dates in the name or if a single file spans multiple dates
            try:
                mod_time_dt = pd.to_datetime(os.path.getmtime(file), unit='s', utc=True)
                # Check if the file's modification date is within the range or shortly after
                # We add a buffer (e.g., 1 day) because file might contain data slightly past its mod time
                if start_dt <= mod_time_dt <= (end_dt + timedelta(days=1)):
                    relevant_files.add(file)
            except Exception as e:
                logger.debug(f"Could not get or parse modification time for {file}: {e}")

        # If no files found based on date/mtime, use the most recent file as a fallback
        # This is a safety measure, but might lead to incorrect data if the range is old
        if not relevant_files and all_files:
            logger.warning(f"No DB files found matching date range {start_date_only} - {end_date_only}. Using most recent file as fallback: {os.path.basename(all_files[0])}")
            return [all_files[0]]
        elif not relevant_files:
            logger.error("No relevant DB files found and no fallback files available.")
            return []

        # Sort the relevant files chronologically (oldest first for processing)
        # Sorting by basename which often includes date is a reasonable heuristic
        return sorted(list(relevant_files), key=lambda f: os.path.basename(f))

    def _get_table_name(self, conn: sqlite3.Connection, exchange: str, interval: str = "1min") -> Optional[str]:
        """Find the correct table name, trying variations. Prioritizes 1min."""
        cursor = conn.cursor()
        try:
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            tables = [row[0] for row in cursor.fetchall()]
        except sqlite3.Error as e:
            logger.error(f"Failed to list tables in database: {e}")
            return None

        # Standard format check (lowercase exchange, exact interval)
        base_table = f"{exchange.lower()}_ohlcv_{interval}"
        if base_table in tables: return base_table

        # Check for 1min source table specifically (common case)
        one_min_table = f"{exchange.lower()}_ohlcv_1min"
        if one_min_table in tables: return one_min_table

        # Try other variations (case, common interval formats)
        variations = [
            f"{exchange.upper()}_ohlcv_{interval}",
            f"{exchange.upper()}_ohlcv_1min",
            f"{exchange.lower()}_ohlcv_1m", # Common abbreviation
            f"{exchange.upper()}_ohlcv_1m",
        ]
        for var in variations:
            if var in tables:
                logger.debug(f"Found table using variation: {var}")
                return var

        # Fallback: Check if *any* OHLCV table exists for the exchange
        for t in tables:
             if t.lower().startswith(f"{exchange.lower()}_ohlcv_"):
                  logger.warning(f"Using first available OHLCV table found for exchange '{exchange}': {t}. Interval might not match '{interval}'.")
                  return t

        logger.warning(f"No suitable OHLCV table found for exchange '{exchange}' with interval '{interval}' or '1min' in the database.")
        return None

    def _query_data_from_db(self, db_file: str, ticker: str, exchange: str, start_timestamp_ns: int, end_timestamp_ns: int) -> pd.DataFrame:
        """
        Query market data from a single database file for a specific ticker and time range (nanoseconds).
        Always queries the 1-minute interval table.
        """
        instrument_id = f"PAIR-{ticker}" if not ticker.startswith("PAIR-") else ticker
        query_interval = "1min" # Always query base interval from DB

        try:
            logger.debug(f"Querying DB '{os.path.basename(db_file)}' for {instrument_id} between {start_timestamp_ns} and {end_timestamp_ns}")
            with sqlite3.connect(f'file:{db_file}?mode=ro', uri=True) as conn: # Read-only mode
                table_name = self._get_table_name(conn, exchange, query_interval)
                if not table_name:
                    logger.warning(f"No table found for {exchange}/1min in {os.path.basename(db_file)}")
                    return pd.DataFrame()

                cursor = conn.cursor()
                try:
                    cursor.execute(f"PRAGMA table_info({table_name})")
                    columns_info = cursor.fetchall()
                    column_names = [col[1].lower() for col in columns_info]
                except sqlite3.Error as e:
                    logger.error(f"Failed to get column info for table '{table_name}' in {db_file}: {e}")
                    return pd.DataFrame()

                # Check for essential columns
                select_cols = ["tstamp", "open", "high", "low", "close", "volume"]
                if not all(c in column_names for c in select_cols):
                     logger.warning(f"Table '{table_name}' in {db_file} missing one or more standard columns: {select_cols}. Found: {column_names}")
                     return pd.DataFrame()

                # Build query
                select_str = ", ".join(select_cols)
                where_clauses = ["tstamp >= ?", "tstamp <= ?"]
                params: list = [start_timestamp_ns, end_timestamp_ns]

                if 'instrument_id' in column_names:
                    where_clauses.append("instrument_id = ?")
                    params.append(instrument_id)
                # Note: exchange_id filtering is complex due to potential variations; rely on table name for now.

                query = f"SELECT {select_str} FROM {table_name} WHERE {' AND '.join(where_clauses)} ORDER BY tstamp"

                df = pd.read_sql_query(query, conn, params=params)
                if df.empty:
                    logger.debug(f"Query returned no data for {instrument_id} in {os.path.basename(db_file)} for the time range.")
                    return pd.DataFrame()

                # Convert timestamp and set index
                df['date'] = pd.to_datetime(df['tstamp'], unit='ns', utc=True)
                df = df.set_index('date').drop(columns=['tstamp'])
                # Ensure numeric types
                for col in ['open', 'high', 'low', 'close', 'volume']:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

                # Drop rows with NaNs that might result from coerce
                df.dropna(subset=['open', 'high', 'low', 'close'], inplace=True)

                logger.debug(f"Query from {os.path.basename(db_file)} returned {len(df)} rows for {instrument_id}.")
                return df

        except sqlite3.Error as e:
            logger.error(f"SQLite error querying {db_file} table '{table_name if 'table_name' in locals() else 'N/A'}': {e}")
        except Exception as e:
            logger.error(f"Unexpected error querying {db_file}: {e}", exc_info=False)
        return pd.DataFrame()

    def _resample_data(self, df: pd.DataFrame, interval: str) -> pd.DataFrame:
        """Resample 1-minute data to a different interval."""
        if df.empty or not isinstance(df.index, pd.DatetimeIndex):
            logger.warning("Input DataFrame for resampling is empty or has non-DatetimeIndex.")
            return df
        if interval == '1min': # No resampling needed
            return df

        logger.info(f"Resampling data to {interval}...")
        try:
            # Define aggregation rules
            agg_dict = {'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last'}
            if 'volume' in df.columns: agg_dict['volume'] = 'sum'

            # Check for required columns before resampling
            required_cols = ['open', 'high', 'low', 'close']
            missing_cols = [c for c in required_cols if c not in df.columns]
            if missing_cols:
                 logger.error(f"Cannot resample, missing required columns: {missing_cols}")
                 return pd.DataFrame() # Return empty if essential cols missing

            # Perform resampling
            resampled_df = df.resample(interval).agg(agg_dict)

            # Drop rows where essential OHLC data is missing after resampling
            resampled_df = resampled_df.dropna(subset=['open', 'high', 'low', 'close'])

            if resampled_df.empty:
                logger.warning(f"Resampling to {interval} resulted in an empty DataFrame.")
            else:
                logger.info(f"Resampling complete. New shape: {resampled_df.shape}")
            return resampled_df

        except ValueError as e:
            logger.error(f"Invalid interval string for resampling: '{interval}'. Error: {e}")
            return pd.DataFrame()
        except Exception as e:
            logger.error(f"Error during resampling to {interval}: {e}", exc_info=True)
            return pd.DataFrame()

    def load_data(self, ticker: str, exchange: str, start_date: str, end_date: str, interval: str,
                  vol_sampling: bool = False, vol_window: int = 30, vol_quantile: float = 0.5) -> pd.DataFrame:
        """
        Loads, combines, and optionally resamples/filters data from relevant DB files.

        Args:
            ticker (str): The trading pair symbol (e.g., 'SOL-USDT').
            exchange (str): The exchange name (e.g., 'bnbspot').
            start_date (str): Start date string (YYYY-MM-DD).
            end_date (str): End date string (YYYY-MM-DD).
            interval (str): The desired final data interval (e.g., '1min', '5min', '1h').
            vol_sampling (bool): If True, apply volatility-based sampling.
            vol_window (int): Window size for volatility calculation if vol_sampling is True.
            vol_quantile (float): Quantile threshold for volatility sampling.

        Returns:
            pd.DataFrame: Combined and processed OHLCV data.
        """
        logger.info(f"Loading data for {ticker} ({exchange}) from {start_date} to {end_date}, interval {interval}")

        try:
            # Parse dates - add time component to cover full days
            start_dt = pd.to_datetime(start_date, utc=True).replace(hour=0, minute=0, second=0, microsecond=0)
            end_dt = pd.to_datetime(end_date, utc=True).replace(hour=23, minute=59, second=59, microsecond=999999)
            if start_dt >= end_dt:
                raise ValueError("Start date must be before end date")
        except Exception as e:
            logger.error(f"Invalid date format or range: {e}")
            return pd.DataFrame()

        # Timestamps for DB query (nanoseconds)
        start_timestamp_ns = int(start_dt.timestamp() * 1e9)
        end_timestamp_ns = int(end_dt.timestamp() * 1e9)

        # Find relevant database files based on the date range
        db_files_to_query = self._get_relevant_db_files(start_dt, end_dt)
        if not db_files_to_query:
            logger.error("No relevant database files found for the specified date range.")
            return pd.DataFrame()

        logger.info(f"Identified {len(db_files_to_query)} potential DB files: {[os.path.basename(f) for f in db_files_to_query]}")

        # Query each relevant file and collect data
        all_data = []
        for db_file in db_files_to_query:
            df_part = self._query_data_from_db(db_file, ticker, exchange, start_timestamp_ns, end_timestamp_ns)
            if not df_part.empty:
                all_data.append(df_part)

        if not all_data:
            logger.warning(f"No data found in any identified DB files for {ticker} ({exchange}) in the specified range.")
            return pd.DataFrame()

        # Combine data from all files
        try:
            combined_df = pd.concat(all_data)
            # Remove duplicate indices (e.g., from overlapping file queries), keeping the first occurrence
            combined_df = combined_df[~combined_df.index.duplicated(keep='first')]
            # Sort chronologically
            combined_df = combined_df.sort_index()
        except Exception as e:
            logger.error(f"Error concatenating or sorting dataframes: {e}")
            return pd.DataFrame()

        logger.info(f"Combined data shape before final filtering/resampling: {combined_df.shape}")

        # Apply precise date range filtering *after* combining and sorting
        final_df = combined_df[(combined_df.index >= start_dt) & (combined_df.index <= end_dt)]

        if final_df.empty:
            logger.warning(f"Dataframe is empty after final date range filtering ({start_dt} to {end_dt}).")
            return pd.DataFrame()

        logger.info(f"Shape after final date filtering: {final_df.shape}")

        # --- Add future_close for potential leakage analysis upstream ---
        # Placeholder: Needs config access or passed horizon value
        # Assuming horizon = 5 for now as per revisions.txt context.
        # A better implementation would pass cfg.gru.prediction_horizon here.
        prediction_horizon = 5 # TODO: Replace with value from config
        final_df['future_close'] = final_df['close'].shift(-prediction_horizon)
        logger.info(f"Added 'future_close' column shifting by {prediction_horizon} periods.")
        # --- End future_close addition ---

        # --- VOLATILITY SAMPLING ---
        if vol_sampling:
            logger.info("Applying volatility-aware sampling...")
            try:
                vol_mask = sample_by_volatility(final_df, vol_window=vol_window, vol_quantile=vol_quantile)
                rows_before = len(final_df)
                final_df = final_df[vol_mask]
                logger.info(f"Applied volatility sampling. Kept {len(final_df)} of {rows_before} rows.")
                if final_df.empty:
                    logger.warning("DataFrame is empty after volatility sampling.")
                    # Return empty DF immediately if sampling removed everything
                    return pd.DataFrame()
            except Exception as e:
                 logger.error(f"Error during volatility sampling: {e}. Skipping sampling.", exc_info=True)
        # --- END VOLATILITY SAMPLING ---

        # Resample if the requested interval is different from 1min
        if interval != "1min":
            final_df = self._resample_data(final_df, interval)
            if final_df.empty:
                 logger.error(f"Resampling to {interval} resulted in an empty DataFrame. Check resampling logic or input data.")
                 return pd.DataFrame()

        # Final check for NaNs in essential columns
        essential_cols = ['open', 'high', 'low', 'close']
        if final_df[essential_cols].isnull().any().any():
            rows_before = len(final_df)
            final_df.dropna(subset=essential_cols, inplace=True)
            logger.warning(f"Dropped {rows_before - len(final_df)} rows with NaN values in essential OHLC columns after potential resampling.")

        if final_df.empty:
            logger.error(f"Final DataFrame is empty after NaN checks for {ticker}.")
            return pd.DataFrame()

        logger.info(f"Successfully loaded and processed data for {ticker}. Final shape: {final_df.shape}")
        return final_df 

# --- Missing Bar Handling Functions ---

def _consecutive_gaps(missing_indices: pd.DatetimeIndex) -> List[pd.Timedelta]:
    """Calculate the length of consecutive gaps in missing timestamps."""
    if missing_indices.empty:
        return []
    
    # Calculate differences between consecutive missing timestamps
    diffs = missing_indices.to_series().diff()
    
    # Identify the start of each gap (where diff > expected frequency)
    # Assumes missing_indices is sorted, diff() calculates t[i] - t[i-1]
    # We infer the frequency from the first gap if available, otherwise fallback needs consideration.
    # This simple version assumes a constant frequency is implicitly known.
    if len(missing_indices) > 1:
        # Infer frequency from the first difference if possible
        inferred_freq = diffs.iloc[1]
    else:
        inferred_freq = pd.Timedelta(minutes=1) # Default assumption if only one missing point
        # A better approach might need the expected freq passed in.
        
    # A gap starts when the difference is larger than the inferred frequency
    gap_starts = diffs[diffs > inferred_freq].index
    
    # Calculate gap lengths
    gaps = []
    current_start = missing_indices[0]
    count = 0
    for i, ts in enumerate(missing_indices):
        if i > 0:
            # Check if consecutive based on inferred frequency
            if ts - missing_indices[i-1] > inferred_freq:
                # End of a gap
                gaps.append(count) # Store the count of consecutive missing bars
                current_start = ts
                count = 1 # Start a new gap
            else:
                count += 1
        else:
            count = 1 # First timestamp starts a gap
            
    gaps.append(count) # Add the last gap
    
    return gaps

def find_missing_bars(df: pd.DataFrame, freq: str) -> pd.DatetimeIndex:
    """Detects missing timestamps in a DataFrame based on expected frequency."""
    if not isinstance(df.index, pd.DatetimeIndex):
        raise ValueError("DataFrame must have a DatetimeIndex.")
    if df.empty:
        return pd.DatetimeIndex([])

    # Ensure the index is sorted
    df = df.sort_index()

    # Create the expected full date range
    full_range = pd.date_range(start=df.index.min(), end=df.index.max(), freq=freq.replace('T','min'))

    # Find the missing timestamps
    missing = full_range.difference(df.index)
    return missing

def report_missing(missing: pd.DatetimeIndex, cfg: DictConfig, io: IOManager, logger):
    """Reports details about missing bars."""
    total_missing = len(missing)
    strategy = cfg['data']['missing']['strategy'] # Access using dict style
    max_gap_allowed = cfg['data']['missing']['max_gap'] # Access using dict style

    if total_missing == 0:
        logger.info("No missing bars detected.")
        return 0 # Return longest gap = 0

    gaps = _consecutive_gaps(missing)
    longest_gap = max(gaps) if gaps else 0

    # Log warning
    warning_msg = (
        f"Detected {total_missing} missing bars (longest consecutive gap: {longest_gap}). "
        f"Applied strategy: {strategy}."
    )
    logger.warning(warning_msg)

    # Save report
    report_data = {
        "total_missing_bars": total_missing,
        "longest_consecutive_gap": longest_gap,
        "applied_strategy": strategy,
        "gap_distribution": {i: gaps.count(i) for i in sorted(list(set(gaps)))}, # Count occurrences of each gap length
        "missing_timestamps_utc": missing.strftime('%Y-%m-%d %H:%M:%S').tolist() # Store as strings
    }
    try:
        io.save_json(report_data, "missing_bars_summary.json", indent=4)
        logger.info(f"Saved missing bars summary.")
    except Exception as e:
        logger.error(f"Failed to save missing bars summary: {e}")
        
    return longest_gap

def fill_missing_bars(df: pd.DataFrame, cfg: DictConfig, io: IOManager, logger) -> pd.DataFrame:
    """Detects, reports, and fills missing bars based on configuration."""
    freq = cfg['data']['bar_frequency']
    strategy = cfg['data']['missing']['strategy']
    max_gap_allowed = cfg['data']['missing']['max_gap']
    interpolate_cfg = cfg['data']['missing']['interpolate']

    missing_indices = find_missing_bars(df, freq)
    longest_gap = report_missing(missing_indices, cfg, io, logger)

    # Check if longest gap exceeds allowed maximum
    if longest_gap > max_gap_allowed:
        raise ValueError(
            f"Longest consecutive gap ({longest_gap}) exceeds maximum allowed ({max_gap_allowed}). "
            f"Consider adjusting 'data.missing.max_gap' or cleaning the data source."
        )

    if missing_indices.empty:
        df['bar_imputed'] = False # Add column even if no missing data initially
        return df

    # Reindex to include all missing timestamps
    full_range = pd.date_range(start=df.index.min(), end=df.index.max(), freq=freq)
    df_full = df.reindex(full_range)

    # Apply filling strategy
    if strategy == "drop":
        logger.info("Missing bar strategy 'drop': Returning original data (no filling). Missing bars remain as NaNs or gaps.")
        # Note: 'drop' doesn't really fit the fill_missing name, but follows instructions.
        # It implies subsequent steps might handle NaNs or the user accepts gaps.
        # We still need to add the 'bar_imputed' flag based on the *detected* missing indices.
        df_filled = df.copy() # Return original df
        # We can't directly assign based on full_range index to original df if size differs.
        # Instead, add the column and mark True where original index matches a missing one.
        df_filled['bar_imputed'] = df_filled.index.isin(missing_indices)
        return df_filled
    elif strategy == "neutral":
        logger.info("Missing bar strategy 'neutral': Forward filling close, setting OHLC=close, Volume=0.")
        # Ffill everything first to get the last known close
        df_filled = df_full.ffill()
        # Identify rows that were originally NaN (i.e., the missing ones)
        was_nan_mask = df_full['close'].isna()
        # Set Volume to 0 for imputed bars
        df_filled.loc[was_nan_mask, 'volume'] = 0
        # Set Open, High, Low to the forward-filled Close for imputed bars
        df_filled.loc[was_nan_mask, ['open', 'high', 'low']] = df_filled.loc[was_nan_mask, 'close']
        # Ensure no NaNs remain due to ffill at the beginning
        df_filled.bfill(inplace=True)
    elif strategy == "ffill":
        logger.info("Missing bar strategy 'ffill': Forward filling all columns, then backward filling initial NaNs.")
        df_filled = df_full.ffill().bfill()
    elif strategy == "interpolate":
        logger.info(f"Missing bar strategy 'interpolate': Interpolating using method='{interpolate_cfg.method}', limit={interpolate_cfg.limit}.")
        # Interpolate numeric columns (OHLCV)
        numeric_cols = ['open', 'high', 'low', 'close', 'volume']
        df_filled = df_full.copy() # Start with the reindexed df
        for col in numeric_cols:
            if col in df_filled.columns:
                 df_filled[col] = df_filled[col].interpolate(
                     method=interpolate_cfg.method,
                     limit=interpolate_cfg.limit,
                     limit_direction='forward', # Only fill gaps, not leading/trailing NaNs initially
                     limit_area=None # Fill within NaNs only
                 )
        # Still need to handle potential leading/trailing NaNs if any
        df_filled.ffill(inplace=True) # Fill any remaining NaNs at the start using ffill
        df_filled.bfill(inplace=True) # Fill any remaining NaNs at the end using bfill
    else:
        raise ValueError(f"Unknown missing data strategy: {strategy}")

    # Add the 'bar_imputed' column AFTER filling
    df_filled['bar_imputed'] = df_filled.index.isin(missing_indices)

    logger.info(f"Successfully filled missing bars using strategy '{strategy}'. Added 'bar_imputed' column.")
    return df_filled

# --- End Missing Bar Handling --- 