"""
Data Collection Tool - Core functionality for fetching market data and SMC indicators

This module integrates with Openfund-core to provide:
- K-line data fetching across multiple timeframes
- SMC-ICT technical indicators (Struct, FVG, OrderBlock, PDArray)
- Parameter validation for symbols and timeframes
"""

import logging
from typing import Dict, Any, List, Optional
import pandas as pd
import sys
import os

# Add openfund-core to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../openfund-core/src"))

from core.Exchange import Exchange
from core.smc.SMCStruct import SMCStruct
from core.smc.SMCFVG import SMCFVG
from core.smc.SMCOrderBlock import SMCOrderBlock
from core.smc.SMCPDArray import SMCPDArray

from openfund_agent.datas.cache import DataCache
from openfund_agent.datas.exceptions import (
    DataCollectionError,
    RetryableError,
    NetworkError,
    ValidationError,
    IndicatorCalculationError,
)
from openfund_agent.datas.retry import with_retry


class DataCollectionTool:
    """
    Data collection tool that encapsulates Openfund-core data fetching functionality

    This tool provides Agent with access to market data and technical indicators.
    """

    # Supported timeframes
    SUPPORTED_TIMEFRAMES = ["1m", "5m", "15m", "1h", "4h", "1d"]

    # Supported indicator types
    SUPPORTED_INDICATORS = ["struct", "fvg", "orderblock", "pdarray"]

    def __init__(self, exchange_config: Dict[str, Any], cache_ttl: int = 60):
        """
        Initialize data collection tool

        Args:
            exchange_config: Exchange configuration (API keys, etc.)
                Format: {
                    'apiKey': str,
                    'secret': str,
                    'password': str (optional),
                    'enableRateLimit': bool,
                    'options': dict (optional)
                }
            cache_ttl: Cache time-to-live in minutes
        """
        self.logger = logging.getLogger(__name__)
        self.cache_ttl = cache_ttl

        # Initialize cache
        self.cache = DataCache(ttl_minutes=cache_ttl)

        # Initialize exchange client
        exchange_key = exchange_config.get("exchange", "okx")
        self.exchange = Exchange(config=exchange_config, exchangeKey=exchange_key)

        # Initialize SMC indicator calculators
        self.smc_struct = SMCStruct()
        self.smc_fvg = SMCFVG()
        self.smc_orderblock = SMCOrderBlock()
        self.smc_pdarray = SMCPDArray()

        self.logger.info(f"DataCollectionTool initialized with exchange: {exchange_key}")

    def fetch_kline_data(self, symbol: str, timeframe: str, limit: int = 500) -> pd.DataFrame:
        """
        Fetch K-line (candlestick) data

        Args:
            symbol: Trading pair symbol in CCXT format (e.g., "BTC/USDT:USDT")
            timeframe: Timeframe, must be one of: 1m, 5m, 15m, 1h, 4h, 1d
            limit: Number of data points to fetch

        Returns:
            DataFrame containing OHLCV data with columns:
            - timestamp: Unix timestamp
            - open: Opening price
            - high: Highest price
            - low: Lowest price
            - close: Closing price
            - volume: Trading volume

        Raises:
            ValueError: If symbol format or timeframe is invalid
            Exception: If data fetching fails
        """
        # Validate parameters
        self._validate_symbol(symbol)
        self._validate_timeframe(timeframe)

        if limit <= 0:
            raise ValidationError(f"limit must be positive, got {limit}")

        # Check cache first
        cache_key = DataCache.make_key("kline", symbol, timeframe, limit=limit)
        cached_data = self.cache.get(cache_key)

        if cached_data is not None:
            self.logger.info(f"Using cached K-line data for {symbol}")
            return cached_data

        self.logger.info(
            f"Fetching K-line data: symbol={symbol}, " f"timeframe={timeframe}, limit={limit}"
        )

        # Fetch data with retry logic
        df = self._fetch_kline_with_retry(symbol, timeframe, limit)

        # Cache the result
        self.cache.set(cache_key, df)

        self.logger.info(f"Successfully fetched {len(df)} K-line records for {symbol}")

        return df

    def fetch_smc_indicators(
        self, symbol: str, timeframe: str, indicator_types: List[str], limit: int = 500
    ) -> Dict[str, Any]:
        """
        Fetch SMC-ICT technical indicators

        Args:
            symbol: Trading pair symbol
            timeframe: Timeframe
            indicator_types: List of indicator types to fetch.
                Supported types: ['struct', 'fvg', 'orderblock', 'pdarray']
            limit: Number of K-line data points to use for calculation

        Returns:
            Dictionary containing requested indicators:
            {
                'struct': {
                    'swing_highs': List[Dict],  # Swing high points
                    'swing_lows': List[Dict],   # Swing low points
                    'trend': str,               # 'Bullish' | 'Bearish' | 'Neutral'
                    'direction': int            # 1 (bullish) | -1 (bearish) | 0 (neutral)
                },
                'fvg': {
                    'bullish': List[Dict],      # Bullish Fair Value Gaps
                    'bearish': List[Dict]       # Bearish Fair Value Gaps
                },
                'orderblock': {
                    'bullish': List[Dict],      # Bullish Order Blocks
                    'bearish': List[Dict]       # Bearish Order Blocks
                },
                'pdarray': {
                    'premium': List[Dict],      # Premium zones
                    'discount': List[Dict],     # Discount zones
                    'equilibrium': float        # Equilibrium price
                }
            }

        Raises:
            ValueError: If indicator types are invalid
            Exception: If indicator calculation fails
        """
        # Validate parameters
        self._validate_symbol(symbol)
        self._validate_timeframe(timeframe)
        self._validate_indicator_types(indicator_types)

        # Check cache first
        cache_key = DataCache.make_key(
            "indicators", symbol, timeframe, *sorted(indicator_types), limit=limit
        )
        cached_data = self.cache.get(cache_key)

        if cached_data is not None:
            self.logger.info(f"Using cached SMC indicators for {symbol}")
            return cached_data

        try:
            self.logger.info(
                f"Fetching SMC indicators: symbol={symbol}, "
                f"timeframe={timeframe}, indicators={indicator_types}"
            )

            # First fetch K-line data
            df = self.fetch_kline_data(symbol, timeframe, limit)

            # Calculate ATR (required for many indicators)
            df = self.smc_struct.calculate_atr(df, period=200, multiplier=2)

            result = {}

            # Calculate requested indicators
            if "struct" in indicator_types:
                result["struct"] = self._calculate_struct(df, timeframe)

            if "fvg" in indicator_types:
                # FVG requires struct data
                if "struct" not in result:
                    struct_df = self._calculate_struct_df(df, timeframe)
                else:
                    struct_df = df
                result["fvg"] = self._calculate_fvg(struct_df)

            if "orderblock" in indicator_types:
                # OrderBlock requires struct data
                if "struct" not in result:
                    struct_df = self._calculate_struct_df(df, timeframe)
                else:
                    struct_df = df
                result["orderblock"] = self._calculate_orderblock(struct_df)

            if "pdarray" in indicator_types:
                # PDArray requires struct data
                if "struct" not in result:
                    struct_df = self._calculate_struct_df(df, timeframe)
                else:
                    struct_df = df
                result["pdarray"] = self._calculate_pdarray(struct_df)

            # Cache the result
            self.cache.set(cache_key, result)

            self.logger.info(
                f"Successfully calculated {len(indicator_types)} indicators for {symbol}"
            )

            return result

        except ValidationError:
            # Re-raise validation errors without wrapping
            raise
        except RetryableError:
            # Re-raise retryable errors without wrapping
            raise
        except Exception as e:
            error_msg = f"Failed to calculate SMC indicators for {symbol}: {str(e)}"
            self.logger.error(error_msg)
            raise IndicatorCalculationError(
                error_msg,
                details={"symbol": symbol, "timeframe": timeframe, "indicators": indicator_types},
            )

    def _calculate_struct_df(self, df: pd.DataFrame, timeframe: str) -> pd.DataFrame:
        """Calculate structure and return DataFrame with struct columns"""
        return self.smc_struct.build_struct(
            df, tf=timeframe, is_struct_body_break=True, struct_timeout_bars=100
        )

    def _calculate_struct(self, df: pd.DataFrame, timeframe: str) -> Dict[str, Any]:
        """
        Calculate market structure (swing highs/lows and trend)

        Returns:
            Dictionary with structure information
        """
        struct_df = self._calculate_struct_df(df, timeframe)

        # Extract swing highs and lows
        swing_highs = []
        swing_lows = []

        for idx, row in struct_df.iterrows():
            if row.get("struct_high") and row["struct_high"] > 0:
                swing_highs.append(
                    {
                        "index": int(row.get("struct_high_index", idx)),
                        "price": float(row["struct_high"]),
                        "timestamp": row["timestamp"],
                    }
                )

            if row.get("struct_low") and row["struct_low"] > 0:
                swing_lows.append(
                    {
                        "index": int(row.get("struct_low_index", idx)),
                        "price": float(row["struct_low"]),
                        "timestamp": row["timestamp"],
                    }
                )

        # Determine trend from last direction
        last_direction = struct_df["direction"].iloc[-1] if "direction" in struct_df.columns else 0

        if last_direction > 0:
            trend = "Bullish"
        elif last_direction < 0:
            trend = "Bearish"
        else:
            trend = "Neutral"

        return {
            "swing_highs": swing_highs,
            "swing_lows": swing_lows,
            "trend": trend,
            "direction": int(last_direction),
        }

    def _calculate_fvg(self, struct_df: pd.DataFrame) -> Dict[str, Any]:
        """
        Calculate Fair Value Gaps

        Returns:
            Dictionary with bullish and bearish FVGs
        """
        # Find bullish FVGs
        bullish_fvg_df = self.smc_fvg.find_FVGs(struct_df, side="buy", check_balanced=True)

        # Find bearish FVGs
        bearish_fvg_df = self.smc_fvg.find_FVGs(struct_df, side="sell", check_balanced=True)

        # Convert to list of dictionaries
        bullish_fvgs = []
        for idx, row in bullish_fvg_df.iterrows():
            bullish_fvgs.append(
                {
                    "index": int(idx),
                    "top": float(row["fvg_top"]),
                    "bottom": float(row["fvg_bot"]),
                    "mid": float(row["fvg_mid"]),
                    "timestamp": row["timestamp"],
                }
            )

        bearish_fvgs = []
        for idx, row in bearish_fvg_df.iterrows():
            bearish_fvgs.append(
                {
                    "index": int(idx),
                    "top": float(row["fvg_top"]),
                    "bottom": float(row["fvg_bot"]),
                    "mid": float(row["fvg_mid"]),
                    "timestamp": row["timestamp"],
                }
            )

        return {"bullish": bullish_fvgs, "bearish": bearish_fvgs}

    def _calculate_orderblock(self, struct_df: pd.DataFrame) -> Dict[str, Any]:
        """
        Calculate Order Blocks

        Returns:
            Dictionary with bullish and bearish order blocks
        """
        # Find bullish order blocks
        bullish_ob_df = self.smc_orderblock.find_order_blocks(struct_df, side="buy")

        # Find bearish order blocks
        bearish_ob_df = self.smc_orderblock.find_order_blocks(struct_df, side="sell")

        # Convert to list of dictionaries
        bullish_obs = []
        for idx, row in bullish_ob_df.iterrows():
            bullish_obs.append(
                {
                    "index": int(idx),
                    "top": float(row.get("ob_top", 0)),
                    "bottom": float(row.get("ob_bot", 0)),
                    "timestamp": row["timestamp"],
                }
            )

        bearish_obs = []
        for idx, row in bearish_ob_df.iterrows():
            bearish_obs.append(
                {
                    "index": int(idx),
                    "top": float(row.get("ob_top", 0)),
                    "bottom": float(row.get("ob_bot", 0)),
                    "timestamp": row["timestamp"],
                }
            )

        return {"bullish": bullish_obs, "bearish": bearish_obs}

    def _calculate_pdarray(self, struct_df: pd.DataFrame) -> Dict[str, Any]:
        """
        Calculate Premium/Discount Arrays

        Returns:
            Dictionary with premium and discount zones
        """
        # Calculate PD Array
        pdarray_result = self.smc_pdarray.calculate_pd_array(struct_df)

        if pdarray_result is None or pdarray_result.empty:
            return {"premium": [], "discount": [], "equilibrium": 0.0}

        # Extract premium and discount zones
        premium_zones = []
        discount_zones = []
        equilibrium = 0.0

        # Get the latest PD array data
        if "premium_high" in pdarray_result.columns:
            for idx, row in pdarray_result.iterrows():
                if row.get("premium_high") and row["premium_high"] > 0:
                    premium_zones.append(
                        {
                            "index": int(idx),
                            "high": float(row["premium_high"]),
                            "low": float(row.get("premium_low", 0)),
                            "timestamp": row["timestamp"],
                        }
                    )

                if row.get("discount_high") and row["discount_high"] > 0:
                    discount_zones.append(
                        {
                            "index": int(idx),
                            "high": float(row["discount_high"]),
                            "low": float(row.get("discount_low", 0)),
                            "timestamp": row["timestamp"],
                        }
                    )

                if row.get("equilibrium"):
                    equilibrium = float(row["equilibrium"])

        return {"premium": premium_zones, "discount": discount_zones, "equilibrium": equilibrium}

    @with_retry(max_retries=3, delay=1.0, backoff=2.0)
    def _fetch_kline_with_retry(self, symbol: str, timeframe: str, limit: int) -> pd.DataFrame:
        """
        Fetch K-line data with automatic retry on failure

        This method wraps the exchange API call with retry logic.

        Args:
            symbol: Trading pair symbol
            timeframe: Timeframe
            limit: Number of data points

        Returns:
            DataFrame with K-line data

        Raises:
            RetryableError: If fetching fails after all retries
        """
        try:
            df = self.exchange.get_historical_klines_df(symbol=symbol, bar=timeframe, limit=limit)

            if df is None or df.empty:
                raise RetryableError(
                    f"Empty data returned for {symbol}",
                    details={"symbol": symbol, "timeframe": timeframe, "limit": limit},
                )

            return df

        except Exception as e:
            # Wrap exceptions as RetryableError for retry logic
            error_msg = f"Failed to fetch K-line data: {str(e)}"
            self.logger.error(error_msg)
            raise RetryableError(
                error_msg, details={"symbol": symbol, "timeframe": timeframe, "limit": limit}
            )

    def _validate_symbol(self, symbol: str) -> None:
        """
        Validate trading pair symbol format

        Args:
            symbol: Symbol to validate

        Raises:
            ValueError: If symbol format is invalid
        """
        if not symbol or not isinstance(symbol, str):
            raise ValidationError(f"Invalid symbol: must be a non-empty string, got {symbol}")

        # Check CCXT format: BASE/QUOTE or BASE/QUOTE:SETTLE
        if "/" not in symbol:
            raise ValidationError(
                f"Invalid symbol format: must be in CCXT format "
                f"(e.g., 'BTC/USDT' or 'BTC/USDT:USDT'), got '{symbol}'"
            )

    def _validate_timeframe(self, timeframe: str) -> None:
        """
        Validate timeframe

        Args:
            timeframe: Timeframe to validate

        Raises:
            ValidationError: If timeframe is not supported
        """
        if timeframe not in self.SUPPORTED_TIMEFRAMES:
            raise ValidationError(
                f"Unsupported timeframe: '{timeframe}'. "
                f"Must be one of {self.SUPPORTED_TIMEFRAMES}"
            )

    def _validate_indicator_types(self, indicator_types: List[str]) -> None:
        """
        Validate indicator types

        Args:
            indicator_types: List of indicator types to validate

        Raises:
            ValidationError: If any indicator type is not supported
        """
        if not indicator_types:
            raise ValidationError("indicator_types cannot be empty")

        invalid_types = [t for t in indicator_types if t not in self.SUPPORTED_INDICATORS]
        if invalid_types:
            raise ValidationError(
                f"Unsupported indicator types: {invalid_types}. "
                f"Must be from {self.SUPPORTED_INDICATORS}"
            )
