# app.py

import dash
from dash import html, dcc, Output, Input
import dash_bootstrap_components as dbc
import plotly.graph_objs as go
from plotly.subplots import make_subplots
import pandas as pd
import json
import logging
from pathlib import Path
from typing import Any, List, Tuple

from viz_io import DATA_FILE, DEPTH_FILE, METRICS_FILE

_LAST_DATA: List[list] = []
_LAST_DEPTH: dict = {"bids": [], "asks": []}
_LAST_METRICS: List[list] = []

app = dash.Dash(__name__, external_stylesheets=[dbc.themes.FLATLY], suppress_callback_exceptions=True)

app.layout = dbc.Container([
    dbc.Row([
        dbc.Col([
            dcc.Graph(
                id='ohlc-chart',
                config={'displayModeBar': True, 'displaylogo': False},
                style={'height': '100vh'}
            )
        ], width=9),
        dbc.Col([
            dcc.Graph(
                id='depth-chart',
                config={'displayModeBar': True, 'displaylogo': False},
                style={'height': '100vh'}
            )
        ], width=3)
    ]),
    dcc.Interval(
        id='interval-update', 
        interval=500,
        n_intervals=0
    )
], fluid=True, style={
    'backgroundColor': '#000000',
    'minHeight': '100vh',
    'color': '#ffffff'
})

def build_empty_ohlc_fig() -> go.Figure:
    """Empty OHLC+Volume+OBI+CVD subplot layout to avoid layout jump before data."""
    fig = make_subplots(
        rows=4,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        row_heights=[0.50, 0.18, 0.16, 0.16],
        subplot_titles=('OHLC', 'Volume', 'OBI', 'CVD'),
        specs=[[{"secondary_y": False}],
               [{"secondary_y": False}], 
               [{"secondary_y": False}],
               [{"secondary_y": False}]]
    )
    fig.update_layout(
        xaxis_title='Time',
        yaxis_title='Price',
        template='plotly_dark',
        autosize=True,
        xaxis_rangeslider_visible=False,
        paper_bgcolor='#000000',
        plot_bgcolor='#000000',
        uirevision='ohlc_v2',
        yaxis2_title='Volume',
        yaxis3_title='OBI',
        yaxis4_title='CVD',
        showlegend=False,
        xaxis_rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="minute", stepmode="backward"),
                dict(count=5, label="5m", step="minute", stepmode="backward"),
                dict(count=15, label="15m", step="minute", stepmode="backward"),
                dict(count=1, label="1h", step="hour", stepmode="backward"),
                dict(step="all", label="All")
            ])
        )
    )
    fig.update_xaxes(matches='x', row=1, col=1)
    fig.update_xaxes(matches='x', row=2, col=1)
    fig.update_xaxes(matches='x', row=3, col=1)
    return fig


def build_empty_depth_fig() -> go.Figure:
    fig = go.Figure()
    fig.update_layout(
        xaxis_title='Size',
        yaxis_title='Price',
        template='plotly_dark',
        autosize=True,
        paper_bgcolor='#000000',
        plot_bgcolor='#000000',
        uirevision='depth',
        showlegend=False
    )
    return fig


def build_ohlc_fig(df: pd.DataFrame) -> go.Figure:
    """Build a four-row subplot with OHLC candlesticks, volume, OBI, and CVD.

    The four subplots share the same X axis (time). Volume bars are directionally
    colored (green for up, red for down). Candlestick colors use Plotly defaults.
    """
    fig = make_subplots(
        rows=4,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        row_heights=[0.50, 0.18, 0.16, 0.16],
        subplot_titles=('OHLC', 'Volume', 'OBI', 'CVD'),
        specs=[[{"secondary_y": False}],
               [{"secondary_y": False}], 
               [{"secondary_y": False}],
               [{"secondary_y": False}]]
    )

    fig.add_trace(
        go.Candlestick(
            x=df['datetime'],
            open=df['open'],
            high=df['high'],
            low=df['low'],
            close=df['close'],
            name='OHLC'
        ),
        row=1,
        col=1
    )

    volume_colors = [
        'rgba(0,200,0,0.7)' if close_val >= open_val else 'rgba(200,0,0,0.7)'
        for open_val, close_val in zip(df['open'], df['close'])
    ]

    fig.add_trace(
        go.Bar(
            x=df['datetime'],
            y=df['volume'],
            marker=dict(color=volume_colors, line=dict(width=0)),
            name='Volume',
            width=50000
        ),
        row=2,
        col=1
    )

    fig.update_layout(
        template='plotly_dark',
        autosize=True,
        xaxis_rangeslider_visible=False,
        paper_bgcolor='#000000',
        plot_bgcolor='#000000',
        uirevision='ohlc_v2',
        yaxis_title='Price',
        yaxis2_title='Volume',
        yaxis3_title='OBI',
        yaxis4_title='CVD',
        showlegend=False,
        xaxis_rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="minute", stepmode="backward"),
                dict(count=5, label="5m", step="minute", stepmode="backward"),
                dict(count=15, label="15m", step="minute", stepmode="backward"),
                dict(count=1, label="1h", step="hour", stepmode="backward"),
                dict(step="all", label="All")
            ])
        )
    )
    
    # All x-axes share the same range selector buttons
    fig.update_xaxes(matches='x', row=1, col=1)
    fig.update_xaxes(matches='x', row=2, col=1)
    fig.update_xaxes(matches='x', row=3, col=1)
    fig.update_xaxes(matches='x', row=4, col=1)

    return fig


def add_obi_subplot(fig: go.Figure, metrics: pd.DataFrame) -> go.Figure:
    """Add raw OBI candlestick data to the existing row 3 subplot."""
    if metrics.empty:
        return fig

    fig.add_trace(
        go.Candlestick(
            x=metrics['datetime'],
            open=metrics['obi_open'],
            high=metrics['obi_high'],
            low=metrics['obi_low'],
            close=metrics['obi_close'],
            increasing_line_color='rgba(0, 120, 255, 1.0)',
            decreasing_line_color='rgba(0, 80, 180, 1.0)',
            increasing_fillcolor='rgba(0, 120, 255, 0.6)',
            decreasing_fillcolor='rgba(0, 80, 180, 0.6)',
            name='OBI'
        ),
        row=3, col=1
    )
    # Zero line
    fig.add_hline(y=0, line=dict(color='rgba(100,100,150,0.5)', width=1), row=3, col=1)
    return fig


def add_cvd_subplot(fig: go.Figure, metrics: pd.DataFrame) -> go.Figure:
    """Add CVD line chart to the existing row 4 subplot."""
    if metrics.empty or 'cvd_value' not in metrics.columns:
        return fig

    fig.add_trace(
        go.Scatter(
            x=metrics['datetime'],
            y=metrics['cvd_value'],
            mode='lines+markers',
            line=dict(color='rgba(255, 165, 0, 1.0)', width=3),
            marker=dict(color='rgba(255, 165, 0, 0.8)', size=4),
            name='CVD'
        ),
        row=4, col=1
    )
    # Zero line
    fig.add_hline(y=0, line=dict(color='rgba(100,100,150,0.5)', width=1), row=4, col=1)
    return fig


def _cumulate_levels(levels: List[List[float]], reverse: bool, limit: int) -> List[Tuple[float, float]]:
    try:
        lv = sorted(levels, key=lambda x: float(x[0]), reverse=reverse)[:limit]
        cumulative: List[Tuple[float, float]] = []
        total = 0.0
        for price, size in lv:
            total += float(size)
            cumulative.append((float(price), total))
        return cumulative
    except Exception:
        return []


def build_depth_fig(depth: Any, levels_per_side: int = 50) -> go.Figure:
    fig = build_empty_depth_fig()
    if not isinstance(depth, dict):
        return fig
    bids = depth.get('bids', [])
    asks = depth.get('asks', [])

    cum_bids = _cumulate_levels(bids, reverse=True, limit=levels_per_side)
    cum_asks = _cumulate_levels(asks, reverse=False, limit=levels_per_side)

    if cum_bids:
        fig.add_trace(go.Scatter(
            x=[s for _, s in cum_bids],
            y=[p for p, _ in cum_bids],
            mode='lines',
            name='Bids',
            line=dict(color='rgba(0,200,0,1)'),
            fill='tozerox',
            fillcolor='rgba(0,200,0,0.2)',
            line_shape='hv'
        ))
    if cum_asks:
        fig.add_trace(go.Scatter(
            x=[s for _, s in cum_asks],
            y=[p for p, _ in cum_asks],
            mode='lines',
            name='Asks',
            line=dict(color='rgba(200,0,0,1)'),
            fill='tozerox',
            fillcolor='rgba(200,0,0,0.2)',
            line_shape='hv'
        ))
    return fig


# Setup callback
@app.callback(
    [Output('ohlc-chart', 'figure'), Output('depth-chart', 'figure')],
    [Input('interval-update', 'n_intervals')]
)
def update_chart(n):
    try:
        if not DATA_FILE.exists():
            data = _LAST_DATA
        else:
            try:
                with open(DATA_FILE, "r") as f:
                    data = json.load(f)

                if isinstance(data, list):
                    _LAST_DATA.clear()
                    _LAST_DATA.extend(data)
                else:
                    data = _LAST_DATA
            except json.JSONDecodeError:
                logging.warning("JSON decode error while reading OHLC data; using cached data")
                data = _LAST_DATA
        
        depth_data = _LAST_DEPTH
        try:
            if DEPTH_FILE.exists():
                with open(DEPTH_FILE, "r") as f:
                    loaded = json.load(f)
                    if isinstance(loaded, dict):
                        _LAST_DEPTH.clear()
                        _LAST_DEPTH.update(loaded)
                        depth_data = _LAST_DEPTH
        except json.JSONDecodeError:
            logging.warning("JSON decode error while reading depth data; using cached depth data")
        except Exception as e:
            logging.debug(f"Depth read skipped: {e}")
        depth_fig = build_depth_fig(depth_data, levels_per_side=50)

        # Read metrics
        metrics_raw = _LAST_METRICS
        try:
            if METRICS_FILE.exists():
                with open(METRICS_FILE, "r") as f:
                    loaded = json.load(f)
                    if isinstance(loaded, list):
                        _LAST_METRICS.clear()
                        _LAST_METRICS.extend(loaded)
                        metrics_raw = _LAST_METRICS
        except json.JSONDecodeError:
            logging.warning("JSON decode error while reading metrics data; using cached metrics data")
        except Exception as e:
            logging.debug(f"Metrics read skipped: {e}")

        if not data:
            ohlc_fig = build_empty_ohlc_fig()
            return [ohlc_fig, depth_fig]
        
        df = pd.DataFrame(data, columns=['timestamp', 'open', 'high', 'low', 'close', 'volume'])
        df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
        ohlc_fig = build_ohlc_fig(df)

        # Build metrics df and append OBI and CVD subplots
        if metrics_raw:
            # Handle both 5-element (legacy) and 6-element (with CVD) schemas
            if metrics_raw and len(metrics_raw[0]) == 6:
                md = pd.DataFrame(metrics_raw, columns=[
                    'timestamp', 'obi_open', 'obi_high', 'obi_low', 'obi_close', 'cvd_value'
                ])
            else:
                md = pd.DataFrame(metrics_raw, columns=[
                    'timestamp', 'obi_open', 'obi_high', 'obi_low', 'obi_close'
                ])
                md['cvd_value'] = 0.0  # Default CVD for backward compatibility
            
            md['datetime'] = pd.to_datetime(md['timestamp'], unit='ms')
            ohlc_fig = add_obi_subplot(ohlc_fig, md)
            ohlc_fig = add_cvd_subplot(ohlc_fig, md)
        
        return [ohlc_fig, depth_fig]
        
    except Exception as e:
        logging.error(f"Error updating chart: {e}")

        ohlc_fig = go.Figure()
        ohlc_fig.update_layout(
            title=f'Error: {str(e)}',
            xaxis_title='Time',
            yaxis_title='Price',
            template='plotly_white',
            height=600
        )

        depth_fig = go.Figure()
        depth_fig.update_layout(
            title=f'Error: {str(e)}',
            xaxis_title='Spread',
            yaxis_title='Spread Size',
            template='plotly_white',
            height=600
        )
        return [ohlc_fig, depth_fig]
        
if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    logging.info(f"Starting OHLC visualizer on http://localhost:8050")
    logging.info(f"Registered callbacks: {list(app.callback_map.keys())}")
    app.run(debug=False, port=8050, host='0.0.0.0')