"""Chart data access and analysis tools."""

from typing import Dict, Any, Optional, Tuple
import io
import uuid
import logging
from pathlib import Path
from contextlib import redirect_stdout, redirect_stderr
from langchain_core.tools import tool

logger = logging.getLogger(__name__)


def _get_registry():
    """Get the global registry instance."""
    from . import _registry
    return _registry


def _get_datasource_registry():
    """Get the global datasource registry instance."""
    from . import _datasource_registry
    return _datasource_registry


def _get_indicator_registry():
    """Get the global indicator registry instance."""
    from . import _indicator_registry
    return _indicator_registry


def _get_order_store():
    """Get the global OrderStore instance."""
    registry = _get_registry()
    if registry and "OrderStore" in registry.entries:
        return registry.entries["OrderStore"].model
    return None


def _get_chart_store():
    """Get the global ChartStore instance."""
    registry = _get_registry()
    if registry and "ChartStore" in registry.entries:
        return registry.entries["ChartStore"].model
    return None


async def _get_chart_data_impl(countback: Optional[int] = None):
    """Internal implementation for getting chart data.

    This is a helper function that can be called by both get_chart_data tool
    and analyze_chart_data tool.

    Returns:
        Tuple of (HistoryResult, chart_context dict, source_name)
    """
    registry = _get_registry()
    datasource_registry = _get_datasource_registry()

    if not registry:
        raise ValueError("SyncRegistry not initialized - cannot read ChartStore")

    if not datasource_registry:
        raise ValueError("DataSourceRegistry not initialized - cannot query data")

    # Read current chart state
    chart_store = registry.entries.get("ChartStore")
    if not chart_store:
        raise ValueError("ChartStore not found in registry")

    chart_state = chart_store.model.model_dump(mode="json")
    chart_data = chart_state.get("chart_state", {})

    symbol = chart_data.get("symbol", "")
    interval = chart_data.get("interval", "15")
    start_time = chart_data.get("start_time")
    end_time = chart_data.get("end_time")

    if not symbol or symbol is None:
        raise ValueError(
            "No chart visible - ChartStore symbol is None. "
            "The user is likely on a narrow screen (mobile) where charts are hidden. "
            "Let them know they can view charts on a wider screen, or use get_historical_data() "
            "if they specify a symbol and timeframe."
        )

    # Parse the symbol to extract exchange/source and symbol name
    # Format is "EXCHANGE:SYMBOL" (e.g., "BINANCE:BTC/USDT", "DEMO:BTC/USD")
    if ":" not in symbol:
        raise ValueError(
            f"Invalid symbol format: '{symbol}'. Expected format is 'EXCHANGE:SYMBOL' "
            f"(e.g., 'BINANCE:BTC/USDT' or 'DEMO:BTC/USD')"
        )

    exchange_prefix, symbol_name = symbol.split(":", 1)
    source_name = exchange_prefix.lower()

    # Get the data source
    source = datasource_registry.get(source_name)
    if not source:
        available = datasource_registry.list_sources()
        raise ValueError(
            f"Data source '{source_name}' not found. Available sources: {available}. "
            f"Make sure the exchange in the symbol '{symbol}' matches an available source."
        )

    # Determine time range - REQUIRE it to be set, no defaults
    if start_time is None or end_time is None:
        raise ValueError(
            f"Chart time range not set in ChartStore. start_time={start_time}, end_time={end_time}. "
            f"The user needs to load the chart first, or the frontend may not be sending the visible range. "
            f"Wait for the chart to fully load before analyzing data."
        )

    from_time = int(start_time)
    end_time = int(end_time)
    logger.info(
        f"Using ChartStore time range: from_time={from_time}, end_time={end_time}, "
        f"countback={countback}"
    )

    logger.info(
        f"Querying data source '{source_name}' for symbol '{symbol_name}', "
        f"resolution '{interval}'"
    )

    # Query the data source
    result = await source.get_bars(
        symbol=symbol_name,
        resolution=interval,
        from_time=from_time,
        to_time=end_time,
        countback=countback
    )

    logger.info(
        f"Received {len(result.bars)} bars from data source. "
        f"First bar time: {result.bars[0].time if result.bars else 'N/A'}, "
        f"Last bar time: {result.bars[-1].time if result.bars else 'N/A'}"
    )

    # Build chart context to return along with result
    chart_context = {
        "symbol": symbol,
        "interval": interval,
        "start_time": start_time,
        "end_time": end_time
    }

    return result, chart_context, source_name


@tool
async def get_chart_data(countback: Optional[int] = None) -> Dict[str, Any]:
    """Get the candle/bar data for what the user is currently viewing on their chart.

    This is a convenience tool that automatically:
    1. Reads the ChartStore to see what chart the user is viewing
    2. Parses the symbol to determine the data source (exchange prefix)
    3. Queries the appropriate data source for that symbol's data
    4. Returns the data for the visible time range and interval

    This is the preferred way to access chart data when helping the user analyze
    what they're looking at, since it automatically uses their current chart context.

    **IMPORTANT**: This tool will fail if ChartStore.symbol is None (no chart visible).
    This happens when the user is on a narrow screen (mobile) where charts are hidden.
    In that case, let the user know charts are only visible on wider screens, or use
    get_historical_data() if they specify a symbol and timeframe.

    Args:
        countback: Optional limit on number of bars to return. If not specified,
                   returns all bars in the visible time range.

    Returns:
        Dictionary containing:
        - chart_context: Current chart state (symbol, interval, time range)
        - symbol: The trading pair being viewed
        - resolution: The chart interval
        - bars: List of bar data with 'time' and 'data' fields
        - columns: Schema describing available data columns
        - source: Which data source was used

    Raises:
        ValueError: If ChartStore or DataSourceRegistry is not initialized,
                   if no chart is visible (symbol is None), or if the symbol format is invalid

    Example:
        # User is viewing BINANCE:BTC/USDT on 15min chart
        data = get_chart_data()
        # Returns BTC/USDT data from binance source at 15min resolution
        # for the currently visible time range
    """
    result, chart_context, source_name = await _get_chart_data_impl(countback)

    # Return enriched result with chart context
    response = result.model_dump()
    response["chart_context"] = chart_context
    response["source"] = source_name

    return response


@tool
async def execute_python(code: str, countback: Optional[int] = None) -> Dict[str, Any]:
    """Execute Python code for technical analysis with automatic chart data loading.

    **PRIMARY TOOL for all technical analysis, indicator computation, and chart generation.**

    This is your go-to tool whenever the user asks about indicators, wants to see
    a chart, or needs any computational analysis of market data.

    Pre-loaded Environment:
    - `pd`               : pandas
    - `np`               : numpy
    - `plt`              : matplotlib.pyplot (figures auto-saved to plot_urls)
    - `talib`            : TA-Lib technical analysis library
    - `indicator_registry`: 150+ registered indicators
    - `plot_ohlc(df)`    : Helper function for beautiful candlestick charts
    - `registry`         : SyncRegistry instance - access to all registered stores
    - `datasource_registry`: DataSourceRegistry - access to data sources (binance, etc.)
    - `order_store`      : OrderStore instance - current orders list
    - `chart_store`      : ChartStore instance - current chart state

    Auto-loaded when user has a chart visible (ChartStore.symbol is not None):
    - `df`               : pandas DataFrame with DatetimeIndex and columns:
                           open, high, low, close, volume (OHLCV data ready to use)
    - `chart_context`    : dict with symbol, interval, start_time, end_time

    When NO chart is visible (narrow screen/mobile):
    - `df`               : None
    - `chart_context`    : None

    If `df` is None, you can still load alternative data by:
    - Using chart_store to see what symbol/timeframe is configured
    - Using datasource_registry.get_source('binance') to access data sources
    - Calling datasource.get_history(symbol, interval, start, end) to load any data
    - This allows you to make plots of ANY chart even when not connected to chart view

    The `plot_ohlc()` Helper:
    Create professional candlestick charts instantly:
    - `plot_ohlc(df)` - basic OHLC chart with volume
    - `plot_ohlc(df, title='BTC 15min')` - with custom title
    - `plot_ohlc(df, volume=False)` - price only, no volume
    - Returns a matplotlib Figure that's automatically saved to plot_urls

    Args:
        code: Python code to execute
        countback: Optional limit on number of bars to load (default: all visible bars)

    Returns:
        Dictionary with:
        - script_output    : printed output + last expression result
        - result_dataframe : serialized DataFrame if last expression is a DataFrame
        - plot_urls        : list of image URLs (e.g., ["/uploads/plot_abc123.png"])
        - chart_context    : {symbol, interval, start_time, end_time} or None
        - error            : traceback if execution failed

    Examples:
        # RSI indicator with chart
        execute_python(\"\"\"
        df['RSI'] = talib.RSI(df['close'], 14)
        fig = plot_ohlc(df, title='BTC/USDT with RSI')
        print(f"Current RSI: {df['RSI'].iloc[-1]:.2f}")
        df[['close', 'RSI']].tail(5)
        \"\"\")

        # Multiple indicators
        execute_python(\"\"\"
        df['SMA_20'] = df['close'].rolling(20).mean()
        df['SMA_50'] = df['close'].rolling(50).mean()
        df['BB_upper'] = df['close'].rolling(20).mean() + 2*df['close'].rolling(20).std()
        df['BB_lower'] = df['close'].rolling(20).mean() - 2*df['close'].rolling(20).std()

        fig = plot_ohlc(df, title=f"{chart_context['symbol']} - Bollinger Bands")

        current_price = df['close'].iloc[-1]
        sma20 = df['SMA_20'].iloc[-1]
        print(f"Price: {current_price:.2f}, SMA20: {sma20:.2f}")
        df[['close', 'SMA_20', 'BB_upper', 'BB_lower']].tail(10)
        \"\"\")

        # Pattern detection
        execute_python(\"\"\"
        # Find swing highs
        df['swing_high'] = (df['high'] > df['high'].shift(1)) & (df['high'] > df['high'].shift(-1))
        swing_highs = df[df['swing_high']][['high']].tail(5)

        fig = plot_ohlc(df, title='Swing High Detection')
        print("Recent swing highs:")
        print(swing_highs)
        \"\"\")

        # Load alternative data when df is None or for different symbol/timeframe
        execute_python(\"\"\"
        from datetime import datetime, timedelta

        # Get data source
        binance = datasource_registry.get_source('binance')

        # Load ETH data even if viewing BTC chart
        end_time = datetime.now()
        start_time = end_time - timedelta(days=7)

        result = await binance.get_history(
            symbol='ETH/USDT',
            interval='1h',
            start=int(start_time.timestamp()),
            end=int(end_time.timestamp())
        )

        # Convert to DataFrame
        rows = [{'time': pd.to_datetime(bar.time, unit='s'), **bar.data} for bar in result.bars]
        eth_df = pd.DataFrame(rows).set_index('time')

        # Calculate RSI and plot
        eth_df['RSI'] = talib.RSI(eth_df['close'], 14)
        fig = plot_ohlc(eth_df, title='ETH/USDT 1h - RSI Analysis')
        print(f"ETH RSI: {eth_df['RSI'].iloc[-1]:.2f}")
        \"\"\")

        # Access chart store to see current state
        execute_python(\"\"\"
        print(f"Current symbol: {chart_store.chart_state.symbol}")
        print(f"Current interval: {chart_store.chart_state.interval}")
        print(f"Orders: {len(order_store.orders)}")
        \"\"\")
    """
    import pandas as pd
    import numpy as np
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    try:
        import talib
    except ImportError:
        talib = None
        logger.warning("TA-Lib not available in execute_python environment")

    # --- Attempt to load chart data ---
    df = None
    chart_context = None

    registry = _get_registry()
    datasource_registry = _get_datasource_registry()

    if registry and datasource_registry:
        try:
            result, chart_context, source_name = await _get_chart_data_impl(countback)
            bars = result.bars
            if bars:
                rows = []
                for bar in bars:
                    rows.append({'time': pd.to_datetime(bar.time, unit='s'), **bar.data})
                df = pd.DataFrame(rows).set_index('time')
                for col in ['open', 'high', 'low', 'close', 'volume']:
                    if col in df.columns:
                        df[col] = pd.to_numeric(df[col], errors='coerce')
                logger.info(f"execute_python: loaded {len(df)} bars for {chart_context['symbol']}")
        except Exception as e:
            logger.info(f"execute_python: no chart data loaded ({e})")

    # --- Import chart utilities ---
    from .chart_utils import plot_ohlc

    # --- Get indicator registry ---
    indicator_registry = _get_indicator_registry()

    # --- Get DataStores ---
    order_store = _get_order_store()
    chart_store = _get_chart_store()

    # --- Build globals ---
    script_globals: Dict[str, Any] = {
        'pd': pd,
        'np': np,
        'plt': plt,
        'talib': talib,
        'indicator_registry': indicator_registry,
        'registry': registry,
        'datasource_registry': datasource_registry,
        'order_store': order_store,
        'chart_store': chart_store,
        'df': df,
        'chart_context': chart_context,
        'plot_ohlc': plot_ohlc,
    }

    # --- Execute ---
    uploads_dir = Path(__file__).parent.parent.parent.parent / "data" / "uploads"
    uploads_dir.mkdir(parents=True, exist_ok=True)

    stdout_capture = io.StringIO()
    result_df = None
    error_msg = None
    plot_urls = []

    try:
        with redirect_stdout(stdout_capture), redirect_stderr(stdout_capture):
            exec(code, script_globals)

            # Capture last expression
            lines = code.strip().splitlines()
            if lines:
                last = lines[-1].strip()
                if last and not any(last.startswith(kw) for kw in (
                    'if', 'for', 'while', 'def', 'class', 'import',
                    'from', 'with', 'try', 'return', '#'
                )):
                    try:
                        last_val = eval(last, script_globals)
                        if isinstance(last_val, pd.DataFrame):
                            result_df = last_val
                        elif last_val is not None:
                            stdout_capture.write(str(last_val))
                    except Exception:
                        pass

        # Save plots
        for fig_num in plt.get_fignums():
            fig = plt.figure(fig_num)
            filename = f"plot_{uuid.uuid4()}.png"
            fig.savefig(uploads_dir / filename, format='png', bbox_inches='tight', dpi=100)
            plot_urls.append(f"/uploads/{filename}")
            plt.close(fig)

    except Exception as e:
        import traceback
        error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"

    # --- Build response ---
    response: Dict[str, Any] = {
        'script_output': stdout_capture.getvalue(),
        'chart_context': chart_context,
        'plot_urls': plot_urls,
    }
    if result_df is not None:
        response['result_dataframe'] = {
            'columns': result_df.columns.tolist(),
            'index': result_df.index.astype(str).tolist(),
            'data': result_df.values.tolist(),
            'shape': result_df.shape,
        }
    if error_msg:
        response['error'] = error_msg

    return response


CHART_TOOLS = [
    get_chart_data,
    execute_python
]
