# File: main.py
# Copyright (C) 2026 Erick Ahmed
# SPDX-License-Identifier: AGPL-3.0-or-later

import os
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

import polars as pl
import dash
from dash import dcc, html, Input, Output, State
import dash_bootstrap_components as dbc
import numpy as np

from parser import parse_log, parse_csv
from decoder import decode_j1939_frames
from stats.utils.extractor import load_data
from stats.id_viewer import _format_can_id_vec, plot_bits
from stats.frequency import calculate_frequency, plot_frequency
from stats.correlation import calculate_correlation, plot_correlation_heatmap
from stats.entropy import calculate_byte_entropy, plot_entropy_heatmap
from logs.view import get_logs_table_component, prepare_logs_data

RAW_LOG_DIR = "data/logs"

def parse_vehicle_from_filename(filename: str):
    stem = Path(filename).stem
    if '-' in stem:
        brand, model_part = stem.split('-', 1)
    else:
        brand, model_part = stem, "Unknown"
    model = model_part.replace('_', ' ')
    vehicle = f"{brand} {model}".strip()
    return vehicle, brand, model

def run_pipeline():
    os.makedirs(RAW_LOG_DIR, exist_ok=True)
    os.makedirs("data/csv", exist_ok=True)
    os.makedirs("data/parquet", exist_ok=True)

    for log_file in Path(RAW_LOG_DIR).glob("*.txt"):
        vehicle, brand, model = parse_vehicle_from_filename(log_file.name)

        bus1_csv = f"data/csv/{vehicle}_bus1.csv"
        bus2_csv = f"data/csv/{vehicle}_bus2.csv"
        bus1_parquet = f"data/parquet/{vehicle}_bus1.parquet"
        bus2_parquet = f"data/parquet/{vehicle}_bus2.parquet"
        bus1_decoded = f"data/parquet/{vehicle}_bus1_decoded.parquet"
        bus2_decoded = f"data/parquet/{vehicle}_bus2_decoded.parquet"

        if not Path(bus1_decoded).exists() or not Path(bus2_decoded).exists():
            print(f"Parsing raw log: {log_file.name}...")
            parse_log(str(log_file), bus1_csv, bus2_csv)

            print("Converting to parquet...")
            parse_csv(bus1_csv).sink_parquet(bus1_parquet)
            parse_csv(bus2_csv).sink_parquet(bus2_parquet)

            print("Decoding J1939...")
            df1 = pl.read_parquet(bus1_parquet)
            df2 = pl.read_parquet(bus2_parquet)
            dec1 = decode_j1939_frames(df1)
            dec2 = decode_j1939_frames(df2)
            dec1.write_parquet(bus1_decoded)
            dec2.write_parquet(bus2_decoded)

run_pipeline()

print("Loading data into memory...")
DATA = {}
VEHICLE_META = {}

for log_file in Path(RAW_LOG_DIR).glob("*.txt"):
    vehicle, brand, model = parse_vehicle_from_filename(log_file.name)
    VEHICLE_META[vehicle] = {"brand": brand, "model": model}

    bus1_decoded = f"data/parquet/{vehicle}_bus1_decoded.parquet"
    bus2_decoded = f"data/parquet/{vehicle}_bus2_decoded.parquet"

    if Path(bus1_decoded).exists() and Path(bus2_decoded).exists():
        DATA[vehicle] = {
            "Bus 1": load_data(bus1_decoded),
            "Bus 2": load_data(bus2_decoded)
        }

PRECOMPUTED_FIGURES = {}
DATA_BY_ID = {}
CORR_CACHE = {}

def process_bus_data(vehicle, bus, df):
    precomp = {}
    precomp[f"{vehicle}_{bus}_freq"] = plot_frequency(calculate_frequency(df), title=f"{vehicle} {bus} Frequency")
    precomp[f"{vehicle}_{bus}_entropy"] = plot_entropy_heatmap(calculate_byte_entropy(df), title=f"{vehicle} {bus} Byte-Level Entropy")

    can_id_col = 'ID' if 'ID' in df.columns else 'Identifier'
    formatted = _format_can_id_vec(df[can_id_col])
    df = df.assign(Formatted_ID=formatted)
    df = df.sort_values(['Formatted_ID', 'Timestamp'], kind='stable')

    grouped = {}
    for can_id, group in df.groupby(by='Formatted_ID'):
        byte_cols = [f"b{i}" for i in range(8) if f"b{i}" in group.columns]
        if not group.empty and len(byte_cols) > 0:
            arr = group[byte_cols].to_numpy(dtype=np.float32, copy=False)
            if len(arr) > 1:
                changed = np.any(arr[1:] != arr[:-1], axis=1)
                keep = np.concatenate(([True], changed))
                group = group.iloc[keep]
        grouped[can_id] = (group, byte_cols)

    return vehicle, bus, precomp, grouped

with ThreadPoolExecutor() as executor:
    futures = []
    for vehicle, buses in DATA.items():
        for bus, df in buses.items():
            futures.append(executor.submit(process_bus_data, vehicle, bus, df))
    for future in futures:
        v, b, precomp, grouped = future.result()
        PRECOMPUTED_FIGURES.update(precomp)
        DATA_BY_ID[(v, b)] = grouped

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

app.layout = dbc.Container([
    html.H1("CANveyor", className="my-4"),
    dbc.Tabs([
        dbc.Tab(label="Overview", tab_id="overview", children=[
            html.Div(id="overview-content")
        ]),
        dbc.Tab(label="Logs", tab_id="logs", children=[
            dbc.Row([
                dbc.Col(html.Label("Select Vehicle:", className="mt-2"), width="auto"),
                dbc.Col(dcc.Dropdown(
                    id='logs-vehicle-selector',
                    options=[{'label': v, 'value': v} for v in DATA.keys()],
                    value=list(DATA.keys())[0] if DATA else None,
                    clearable=False
                ), width=3, className="me-4"),
                dbc.Col(html.Label("Select Bus:", className="mt-2"), width="auto"),
                dbc.Col(dcc.Dropdown(
                    id='logs-bus-selector',
                    options=[{'label': 'Bus 1', 'value': 'Bus 1'}, {'label': 'Bus 2', 'value': 'Bus 2'}],
                    value='Bus 1',
                    clearable=False
                ), width=2),
            ], className="mb-3 mt-3", align="end"),
            get_logs_table_component()
        ]),
        dbc.Tab(label="Statistics", tab_id="statistics", children=[
            dbc.Row([
                dbc.Col(html.Label("Select Vehicle:", className="mt-2"), width="auto"),
                dbc.Col(dcc.Dropdown(
                    id='vehicle-selector',
                    options=[{'label': v, 'value': v} for v in DATA.keys()],
                    value=list(DATA.keys())[0] if DATA else None,
                    clearable=False
                ), width=3, className="me-4"),
                dbc.Col(html.Label("Select Bus:", className="mt-2"), width="auto"),
                dbc.Col(dcc.Dropdown(
                    id='bus-selector',
                    options=[{'label': 'Bus 1', 'value': 'Bus 1'}, {'label': 'Bus 2', 'value': 'Bus 2'}],
                    value='Bus 1',
                    clearable=False
                ), width=2),
            ], className="mb-3 mt-3", align="end"),
            dbc.Tabs([
                dbc.Tab(label="Frequency", tab_id="freq"),
                dbc.Tab(label="ID Viewer", tab_id="id_viewer"),
                dbc.Tab(label="Correlation", tab_id="corr"),
                dbc.Tab(label="Entropy", tab_id="entropy"),
            ], id="tabs", active_tab="freq"),
            html.Div(id="tab-content", className="mt-3")
        ])
    ], id="main-tabs", active_tab="statistics")
], fluid=True)

app.clientside_callback(
    """
    function(data) {
        if (!data) return '';
        requestAnimationFrame(() => {
            const btn = document.getElementById('load-more-logs-btn');
            const info = document.getElementById('logs-info-text');
            if (!btn || !info || info.innerText.toLowerCase().includes('all')) {
                return;
            }

            btn.dataset.loading = "false";

            const containers = document.querySelectorAll('.dash-spreadsheet-container, .dash-spreadsheet-inner');
            if (containers.length === 0) return;

            containers.forEach(container => {
                container.onscroll = function() {
                    if (container.scrollHeight - container.scrollTop - container.clientHeight < 1500) {
                        if (btn && btn.dataset.loading === "false") {
                            btn.dataset.loading = "true";
                            btn.click();
                        }
                    }
                };

                if (container.scrollHeight - container.scrollTop - container.clientHeight < 1500) {
                    if (btn && btn.dataset.loading === "false") {
                        btn.dataset.loading = "true";
                        btn.click();
                    }
                }
            });
        });
        return '';
    }
    """,
    Output('dummy-output', 'children'),
    Input('logs-table', 'data')
)

@app.callback(
    Output('logs-table', 'data'),
    Output('logs-table', 'columns'),
    Output('logs-info-text', 'children'),
    Input('logs-vehicle-selector', 'value'),
    Input('logs-bus-selector', 'value'),
    Input('load-more-logs-btn', 'n_clicks'),
    State('logs-table', 'data'),
)
def update_logs_table(vehicle, bus, n_clicks, current_data):
    ctx = dash.callback_context
    trigger_id = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else ''

    if not vehicle or not bus or vehicle not in DATA or bus not in DATA[vehicle]:
        return [], [], "No data available"

    df = DATA[vehicle][bus]
    prepared_df = prepare_logs_data(df)

    total_rows = len(prepared_df)
    columns = [{"name": i, "id": i} for i in prepared_df.columns]

    if trigger_id in ['logs-vehicle-selector', 'logs-bus-selector', '']:
        current_data = []

    offset = len(current_data) if current_data else 0
    chunk_size = 50000

    if offset >= total_rows:
        return current_data, columns, f"Displaying all {total_rows} total frames."

    next_chunk = prepared_df.iloc[offset:offset + chunk_size].to_dict('records')
    new_data = current_data + next_chunk
    new_offset = len(new_data)

    info_text = f"Displaying {new_offset} of {total_rows} total frames."

    return new_data, columns, info_text

@app.callback(
    Output('tab-content', 'children'),
    Input('tabs', 'active_tab'),
    Input('vehicle-selector', 'value'),
    Input('bus-selector', 'value')
)
def render_content(tab, vehicle, bus):
    if not vehicle or not bus or vehicle not in DATA or bus not in DATA[vehicle]:
        return html.Div("No data available")

    df = DATA[vehicle][bus]

    if tab == 'freq':
        return dcc.Graph(figure=PRECOMPUTED_FIGURES[f"{vehicle}_{bus}_freq"], style={'height': '80vh'})

    elif tab == 'id_viewer':
        ids = sorted(DATA_BY_ID.get((vehicle, bus), {}).keys())
        return html.Div([
            html.Label("Select CAN ID:"),
            dcc.Dropdown(
                id='id-selector',
                options=[{'label': i, 'value': i} for i in ids],
                value=ids[0] if ids else None,
                clearable=False,
                style={'width': '50%', 'marginBottom': '10px'}
            ),
            dcc.Graph(id='id-viewer-graph', style={'height': '70vh'})
        ])

    elif tab == 'corr':
        ids = sorted(DATA_BY_ID.get((vehicle, bus), {}).keys())
        return html.Div([
            dbc.Row([
                dbc.Col(html.Label("Method:"), width=1, className="mt-2"),
                dbc.Col(dcc.Dropdown(
                    id='corr-method',
                    options=[{'label': 'Pearson', 'value': 'pearson'}, {'label': 'Spearman', 'value': 'spearman'}],
                    value='pearson',
                    clearable=False
                ), width=2),
                dbc.Col(html.Label("Target ID:"), width=1, className="mt-2"),
                dbc.Col(dcc.Dropdown(
                    id='corr-target',
                    options=[{'label': 'All IDs (Max Corr)', 'value': 'all'}] + [{'label': i, 'value': i} for i in ids],
                    value='all',
                    clearable=True
                ), width=4),
            ], className="mb-3"),
            dcc.Graph(id='corr-graph', style={'height': '80vh'})
        ])

    elif tab == 'entropy':
        return dcc.Graph(figure=PRECOMPUTED_FIGURES[f"{vehicle}_{bus}_entropy"], style={'height': '80vh'})

    return html.Div("Tab not found")

@app.callback(
    Output('id-viewer-graph', 'figure'),
    Input('id-selector', 'value'),
    Input('vehicle-selector', 'value'),
    Input('bus-selector', 'value'),
    Input('tabs', 'active_tab'),
)
def update_id_viewer(selected_id, vehicle, bus, tab):
    if tab != 'id_viewer' or not selected_id or not vehicle or not bus:
        return dash.no_update

    grouped_data = DATA_BY_ID.get((vehicle, bus), {})
    if selected_id not in grouped_data:
        return dash.no_update

    filtered_df, byte_cols = grouped_data[selected_id]
    return plot_bits(filtered_df, byte_cols, selected_id, title=f"{vehicle} {bus} Byte Visualization")

@app.callback(
    Output('corr-graph', 'figure'),
    Input('corr-method', 'value'),
    Input('corr-target', 'value'),
    Input('vehicle-selector', 'value'),
    Input('bus-selector', 'value'),
    Input('tabs', 'active_tab'),
)
def update_corr(method, target, vehicle, bus, tab):
    if tab != 'corr' or not vehicle or not bus:
        return dash.no_update

    target_id = None if target == 'all' or not target else target
    cache_key = (vehicle, bus, method, target_id)

    if cache_key not in CORR_CACHE:
        df = DATA[vehicle][bus]
        corr_df = calculate_correlation(df, method=method, target_id=target_id)
        CORR_CACHE[cache_key] = corr_df
    else:
        corr_df = CORR_CACHE[cache_key]

    title = f"{vehicle} {bus} Correlation"
    if target_id:
        title += f" ({target_id})"

    return plot_correlation_heatmap(corr_df, target_id=target_id, title=title)

if __name__ == '__main__':
    app.run(debug=False)