import argparse
import polars as pl

def get_j1939_mask() -> pl.Expr:
    """
    Returns a Polars expression representing the strict J1939 filtering rules.
    """
    id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32)
    return (
        (id_int > 0x7FF) &
        ((id_int % 33554432 // 16777216) == 0) &
        (pl.col("DLC") <= 8)
    )

def decode_j1939_metadata(lf: pl.LazyFrame) -> pl.LazyFrame:
    """
    Decodes J1939 fields and bundles them into a Struct column.
    """
    id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32)

    id_shifted_8 = id_int // 256
    id_shifted_16 = id_int // 65536

    priority = ((id_int // 67108864) % 8).cast(pl.UInt8)
    pf = (id_shifted_16 % 256).cast(pl.UInt8)
    ps = (id_shifted_8 % 256).cast(pl.UInt8)
    sa = (id_int % 256).cast(pl.UInt8)

    da = pl.when(pf < 240).then(ps).otherwise(pl.lit(255, dtype=pl.UInt8))

    pgn = pl.when(pf < 240).then(id_shifted_8 % 65536).otherwise(id_shifted_8 % 262144).cast(pl.UInt32)

    return lf.with_columns(
        pl.struct([
            priority.alias("Priority"),
            pf.alias("PF"),
            ps.alias("PS"),
            sa.alias("SA"),
            da.alias("DA"),
            pgn.alias("PGN")
        ]).alias("j1939_metadata")
    )

def decode_j1939_frames(df: pl.DataFrame) -> pl.DataFrame:
    id_int = pl.col("ID").str.to_integer(base=16).cast(pl.UInt32)

    is_j1939 = (id_int > 0x7FF) & ((id_int % 33554432 // 16777216) == 0) & (pl.col("DLC") <= 8)

    priority = (id_int // 67108864) % 8
    pf = (id_int // 65536) % 256
    ps = (id_int // 256) % 256
    sa = id_int % 256
    da = pl.when(pf < 240).then(ps).otherwise(255)
    pgn = pl.when(pf < 240).then((id_int // 256) % 65536).otherwise((id_int // 256) % 262144)

    j1939_meta = pl.when(is_j1939).then(
        pl.struct([
            priority.cast(pl.UInt8).alias("Priority"),
            pf.cast(pl.UInt8).alias("PF"),
            ps.cast(pl.UInt8).alias("PS"),
            sa.cast(pl.UInt8).alias("SA"),
            da.cast(pl.UInt8).alias("DA"),
            pgn.cast(pl.UInt32).alias("PGN")
        ])
    ).otherwise(None)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="J1939 decoder")
    parser.add_argument("input_parquet", help="Path to the raw .parquet file")
    parser.add_argument("output_parquet", help="Path to save the decoded .parquet file")
    args = parser.parse_args()

    df = pl.scan_parquet(args.input_parquet).collect()
    decoded_df = decode_j1939_frames(df)
    decoded_df.write_parquet(args.output_parquet)