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

import plotly.express as px
import pandas as pd
from dataclasses import dataclass
from typing import Callable
from copy import deepcopy
from vehicle.base import (
    SignalDef, FrameDef, normalize_id, decode_dataframe as _decode_dataframe, plot_signal
)

@dataclass
class CustomPlotDef:
    name: str
    plot_func: Callable

load_state_sig = SignalDef(
    name="Engine Load State",
    bit_start=24,
    bit_length=8,
    factor=1,
    offset=0.0,
    is_signed=False,
    byte_order="big",
    unit="",
)
load_state_sig.skip_plot = True

DECODER_RULES = {
    normalize_id("0x011F"): FrameDef(
        can_id="0x011F",
        description="Engine ECM Main Broadcast",
        color="#e41a1c",
        signals=[
            SignalDef(
                name="Engine",
                bit_start=0,
                bit_length=16,
                factor=0.125,
                offset=0.0,
                is_signed=False,
                byte_order="big",
                unit="RPM",
            ),
            SignalDef(
                name="Pressure / Load",
                bit_start=16,
                bit_length=16,
                factor=0.05,
                offset=0,
                is_signed=False,
                byte_order="little",
                unit="%",
            ),
        ],
    ),
    normalize_id("0x0CFF3300"): FrameDef(
        can_id="0x0CFF3300",
        description="Engine temperature and load state block",
        color="#0080fe",
        signals=[
            SignalDef(
                name="Engine coolant temp",
                bit_start=8,
                bit_length=8,
                factor=1,
                offset=0.0,
                is_signed=False,
                byte_order="big",
                unit="℃",
            ),
            SignalDef(
                name="Engine oil temp",
                bit_start=40,
                bit_length=8,
                factor=1,
                offset=0.0,
                is_signed=False,
                byte_order="big",
                unit="℃",
            ),
            load_state_sig,
            CustomPlotDef(
                name="Engine Load State",
                plot_func=lambda decoded, color: plot_load_state_pie(decoded, color)
            ),
        ],
    ),
}

LOAD_STATE_MAP = {
    0: "Boot up",
    16: "Normal load",
    32: "High load"
}

def plot_load_state_pie(decoded_df, color):
    if decoded_df is None or decoded_df.empty or "Engine Load State" not in decoded_df.columns:
        return px.pie(title="No data for Engine Load State")

    states = pd.to_numeric(decoded_df["Engine Load State"], errors='coerce').dropna().astype(int)

    labels = states.map(LOAD_STATE_MAP).fillna("Unknown")
    counts = labels.value_counts().reset_index()
    counts.columns = ['State', 'Count']

    total = counts['Count'].sum()
    counts['Percentage'] = (counts['Count'] / total * 100).round(1)
    counts['Legend'] = counts['State'] + " (" + counts['Percentage'].astype(str) + "%)"

    fig = px.pie(
        counts,
        values='Count',
        names='Legend',
        color='State',
        title='Engine Load State Distribution',
        color_discrete_map={
            "Boot up": "#ff9900",
            "Normal load": "#00cc00",
            "High load": "#cc0000",
            "Unknown": "#808080"
        }
    )

    fig.update_traces(
        textinfo='none',
        hoverinfo='label+percent+value',
        domain={'x': [0.05, 0.55], 'y': [0.05, 0.95]}
    )
    fig.update_layout(
        margin=dict(l=0, r=10, t=40, b=0),
        legend=dict(x=0.6, y=0.5)
    )
    return fig

def decode_dataframe(df, can_id):
    filtered_rules = {}
    for nid, frame in DECODER_RULES.items():
        filtered_signals = [sig for sig in frame.signals if isinstance(sig, SignalDef)]
        new_frame = deepcopy(frame)
        new_frame.signals = filtered_signals
        filtered_rules[nid] = new_frame

    return _decode_dataframe(df, can_id, filtered_rules)