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

import argparse
from pathlib import Path
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly_resampler import FigureResampler
from stats.utils.extractor import load_data

def _format_can_id_vec(s: pd.Series) -> pd.Series:
    s = s.astype('string').str.strip()
    s = s.str.replace(r'^0x', '', case=False, regex=True)
    s = s.str.upper()
    return s.fillna('UNKNOWN').replace('', 'UNKNOWN')

def prepare_data(df, target_id):
    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)
    target_id_clean = _format_can_id_vec(pd.Series([target_id])).iloc[0]
    filtered = df[df['Formatted_ID'] == target_id_clean]

    byte_cols = [f"b{i}" for i in range(8) if f"b{i}" in filtered.columns]
    if filtered.empty:
        return filtered, byte_cols

    for col in byte_cols:
        if not pd.api.types.is_numeric_dtype(filtered[col]):
            filtered = filtered.assign(**{col: pd.to_numeric(filtered[col], errors='coerce').astype('float32')})

    filtered = filtered.sort_values('Timestamp', kind='stable')
    arr = filtered[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))
        filtered = filtered.iloc[keep]

    return filtered, byte_cols

def plot_bits(df, byte_cols, can_id, title):
    fig = FigureResampler(
        resampled_trace_prefix_suffix=("", ""),
        show_mean_aggregation_size=False
    )
    colors = ['#e41a1c', '#377eb8', '#4daf4a', '#984ea3', '#ff7f00', '#ffff33', '#a65628', '#f781bf']
    n = len(byte_cols)

    x = df['Timestamp'].to_numpy() if not df.empty else np.array([])
    for i, col in enumerate(byte_cols):
        y = df[col].to_numpy(dtype=np.float32, copy=False) if not df.empty else np.array([])

        fig.add_trace(go.Scatter(
            mode='lines',
            line=dict(shape='hv', width=2, color=colors[i % len(colors)]),
            name=col.upper(),
            legendgroup=col.upper(),
            hovertemplate=f"<b>{col.upper()}</b><br>Time: %{{x}}<br>Value: %{{y}}<extra></extra>",
        ), hf_x=x, hf_y=y)

    all_button = dict(label='ALL', method='restyle', args=[{'visible': [True] * n}])
    none_button = dict(label='NONE', method='restyle', args=[{'visible': ['legendonly'] * n}])

    fig.update_layout(
        height=600,
        autosize=True,
        template='plotly_white',
        title=dict(
            text=f"{title} - ID: {can_id}",
            font=dict(size=20, color='#1a1a1a'),
            x=0.5, xanchor='center',
            pad=dict(b=20),
        ),
        font=dict(family="Segoe UI, Arial, sans-serif", size=12, color='#2a2a2a'),
        hoverlabel=dict(bgcolor="white", font_size=13, font_family="Segoe UI", bordercolor='#cccccc'),
        margin=dict(l=60, r=40, t=120, b=140),
        legend=dict(
            orientation='h',
            x=0.5, xanchor='center',
            y=-0.18, yanchor='top',
            title=None,
            bgcolor='white',
            bordercolor='#cccccc',
            borderwidth=1,
            font=dict(size=12, color="#2a2a2a"),
            itemsizing='constant',
            itemclick='toggle',
            itemdoubleclick='toggleothers',
        ),
        xaxis=dict(
            title=dict(text="Timestamp", font=dict(size=13, color="#1a1a1a")),
            showgrid=True, gridwidth=0.5, gridcolor='#e8e8e8',
            zeroline=False, linecolor="#bdbdbd",
            tickfont=dict(size=12, color="#2a2a2a"),
            ticks="outside", ticklen=4, tickcolor="#cccccc",
            minor=dict(showgrid=True, gridcolor='#f4f4f4', gridwidth=0.5),
        ),
        yaxis=dict(
            title=dict(text="Byte Value", font=dict(size=13, color="#1a1a1a")),
            showgrid=True, gridwidth=0.5, gridcolor='#e8e8e8',
            zeroline=False, linecolor="#bdbdbd",
            tickfont=dict(size=12, color="#2a2a2a"),
            ticks="outside", ticklen=4, tickcolor="#cccccc",
        ),
        updatemenus=[
            dict(
                type='buttons',
                direction='right',
                x=0.5, xanchor='center',
                y=-0.06, yanchor='top',
                buttons=[all_button, none_button],
                bgcolor='white',
                bordercolor='#cccccc',
                borderwidth=1,
                font=dict(size=11, color='#2a2a2a'),
                pad=dict(l=5, r=5, t=5, b=5),
            )
        ],
    )
    return fig

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Visualize CAN bus byte changes over time")
    parser.add_argument("input", type=Path, help="Path to the input CAN log file")
    parser.add_argument("can_id", type=str, help="CAN ID to visualize")
    parser.add_argument("output", type=Path, nargs="?", default=Path("bits_report.html"))
    parser.add_argument("title", nargs="?", default="Byte Visualization")
    args = parser.parse_args()

    df = load_data(args.input)
    filtered_df, byte_cols = prepare_data(df, args.can_id)
    fig = plot_bits(filtered_df, byte_cols, args.can_id, title=args.title)

    config = {
        'responsive': True,
        'displaylogo': False,
        'scrollZoom': True,
        'modeBarButtonsToAdd': ['toggleSpikelines'],
        'toImageButtonOptions': {'format': 'png', 'scale': 2},
    }
    fig.write_html(str(args.output), include_plotlyjs='cdn', config=config)