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

import json
from pathlib import Path

import numpy as np
import pandas as pd
import polars as pl

def to_int(x):
    if isinstance(x, (int, np.integer)):
        return int(x)
    if isinstance(x, str):
        try:
            return int(x, 16)
        except ValueError:
            return np.nan
    return np.nan

def extract_id(row: pd.Series) -> str:
    meta = row.get('j1939_metadata')
    if pd.isna(meta):
        return f"ID: {row['ID']}"
    if isinstance(meta, str):
        try:
            meta = json.loads(meta)
        except json.JSONDecodeError:
            return f"ID: {row['ID']}"
    if isinstance(meta, dict) and 'PGN' in meta:
        return f"PGN: {meta['PGN']}"
    return f"ID: {row['ID']}"

def load_data(file_path: Path) -> pd.DataFrame:
    lf = pl.scan_parquet(file_path)
    schema = lf.collect_schema()
    names = schema.names()

    byte_cols = [f"b{i}" for i in range(8) if f"b{i}" in names]
    if byte_cols:
        lf = lf.with_columns([
            pl.col(c).str.to_integer(base=16, strict=False).cast(pl.Int16).alias(c)
            for c in byte_cols
        ])

    id_col = 'ID' if 'ID' in names else 'Identifier'
    id_expr = pl.col(id_col).cast(pl.Utf8)

    if 'j1939_metadata' in names:
        try:
            lf = lf.with_columns(
                pl.when(pl.col('j1939_metadata').is_not_null())
                .then(pl.lit('PGN: ') + pl.col('j1939_metadata').struct.field('PGN').cast(pl.Utf8))
                .otherwise(pl.lit('ID: ') + id_expr)
                .alias('Identifier')
            )
        except Exception:
            lf = lf.with_columns((pl.lit('ID: ') + id_expr).alias('Identifier'))
    else:
        lf = lf.with_columns((pl.lit('ID: ') + id_expr).alias('Identifier'))

    return lf.collect().to_pandas()