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

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

logger = logging.getLogger(__name__)

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. Can be relative to project root or absolute.
            cache_dir (str): Directory to store cached data (currently not implemented).
            use_cache (bool): Whether to use cached data (currently not implemented).
        """
        # Resolve potential relative db_dir path
        if not os.path.isabs(db_dir):
            # Assume db_dir is relative to the project root (two levels up from src/)
            # This might need adjustment depending on where the main script is run
            script_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(script_dir))
            self.db_dir = os.path.abspath(os.path.join(project_root, db_dir))
            logger.info(f"Resolved relative db_dir '{db_dir}' to absolute path: {self.db_dir}")
        else:
            self.db_dir = db_dir

        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):
            logger.warning(f"Database directory does not exist: {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}")
        if not os.path.exists(self.db_dir):
            logger.error(f"Database directory {self.db_dir} does not exist")
            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) -> pd.DataFrame:
        """
        Loads, combines, and optionally resamples 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').

        Returns:
            pd.DataFrame: Combined and resampled OHLCV data, indexed by UTC timestamp.
                         Returns an empty DataFrame on failure.
        """
        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}")

        # 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 