#!/usr/bin/env python3
"""Helpers for BNC raw capture files (rawOutFile / --file format).

File layout (see src/bncrawfile.cpp):
    "1 Version of BNC raw file\n"
    then repeated:  "<ISO time> <staID> <format> <nBytes>\n" <nBytes raw bytes> "\n"

Sub-commands:
    info  <raw> [<raw> ...]          streams, byte counts, time span, RTCM3 message types
    slice <raw> <out> [options]      cut a time window / subset of stations
    crs   <raw>                      decode RTCM 1300/1301/1302 (CRS messages)
"""

import argparse
import collections
import sys

HEADER = b"1 Version of BNC raw file\n"


def read_chunks(path):
    """Yield (time, staID, format, data) for every chunk of a raw file."""
    with open(path, "rb") as f:
        first = f.readline()
        if b"Version of BNC raw file" not in first:
            raise ValueError("%s: not a BNC raw file" % path)
        while True:
            line = f.readline()
            if not line:
                return
            parts = line.split()
            if len(parts) < 4:
                raise ValueError("%s: bad chunk header %r" % (path, line[:80]))
            n = int(parts[3])
            data = f.read(n)
            if len(data) != n:
                raise ValueError("%s: truncated chunk" % path)
            f.read(1)  # trailing '\n'
            yield parts[0].decode(), parts[1].decode(), parts[2].decode(), data


def write_chunks(path, chunks):
    with open(path, "wb") as o:
        o.write(HEADER)
        for t, sta, fmt, data in chunks:
            o.write(("%s %s %s %d\n" % (t, sta, fmt, len(data))).encode())
            o.write(data)
            o.write(b"\n")


# RTCM3 framing -----------------------------------------------------------

def _crc24q_table():
    tab = []
    for i in range(256):
        c = i << 16
        for _ in range(8):
            c <<= 1
            if c & 0x1000000:
                c ^= 0x1864CFB
        tab.append(c & 0xFFFFFF)
    return tab


_CRC_TAB = _crc24q_table()


def crc24q(buf):
    crc = 0
    for b in buf:
        crc = ((crc << 8) & 0xFFFFFF) ^ _CRC_TAB[(crc >> 16) ^ b]
    return crc


class Rtcm3Framer:
    """Incremental CRC-validated RTCM3 frame extractor (one per stream)."""

    def __init__(self):
        self.buf = bytearray()

    def feed(self, data):
        self.buf += data
        frames = []
        i = 0
        buf = self.buf
        while True:
            i = buf.find(0xD3, i)
            if i < 0 or len(buf) - i < 6:
                break
            length = ((buf[i + 1] & 0x03) << 8) | buf[i + 2]
            if len(buf) - i < length + 6:
                break
            frame = bytes(buf[i:i + length + 6])
            if crc24q(frame[:-3]) == int.from_bytes(frame[-3:], "big"):
                frames.append(frame[3:-3])
                i += length + 6
            else:
                i += 1
        self.buf = buf[i:] if i >= 0 else bytearray()
        return frames


def msg_type(payload):
    """Message number; for 4076 return '4076_<subtype>'."""
    if len(payload) < 2:
        return "empty"
    t = (payload[0] << 4) | (payload[1] >> 4)
    if t == 4076 and len(payload) >= 3:
        # 12 bit type, 3 bit version, 8 bit subtype
        sub = ((payload[1] & 0x01) << 7) | (payload[2] >> 1)
        return "4076_%03d" % sub
    return str(t)


class BitReader:
    def __init__(self, data):
        self.data = data
        self.pos = 0

    def u(self, n):
        v = 0
        for _ in range(n):
            if self.pos >= 8 * len(self.data):
                raise ValueError("message too short")
            v = (v << 1) | ((self.data[self.pos // 8] >> (7 - self.pos % 8)) & 1)
            self.pos += 1
        return v

    def s(self, n):
        v = self.u(n)
        return v - (1 << n) if v >> (n - 1) else v

    def text(self):
        n = self.u(5)
        return bytes(self.u(8) for _ in range(n)).decode("latin-1")

    def left(self):
        return 8 * len(self.data) - self.pos


def decode_crs(p):
    """One line per 1300/1301/1302 message (RTCM 10403.4, DF562-DF587)."""
    r = BitReader(p)
    t = r.u(12)
    if t == 1300:
        out = "1300 name %r CE %.2f" % (r.text(), r.u(16) / 100.0)
    elif t == 1301:
        src, tgt = r.text(), r.text()
        out = "1301 source %r target %r sys %d util %d mjd %d par" % (
            src, tgt, r.u(8), r.u(10), r.u(16) + 44244)
        for n in (23, 23, 23, 32, 32, 32, 25, 17, 17, 17, 17, 17, 17, 14):
            out += " %d" % r.s(n)
    else:
        out = "1302 name %r anchor %d plate %d links" % (r.text(), r.u(1), r.u(5))
        out += " %r" % [r.text() for _ in range(r.u(3))]
    if r.left() >= 8:  # more than the padding to the byte boundary
        out += "  TRAILING %d bits" % r.left()
    return out


# Commands ----------------------------------------------------------------

def cmd_info(args):
    for path in args.files:
        stats = collections.OrderedDict()
        t0 = t1 = None
        n = 0
        for t, sta, fmt, data in read_chunks(path):
            n += 1
            t0 = t0 or t
            t1 = t
            s = stats.get(sta)
            if s is None:
                s = stats[sta] = {"fmt": fmt, "bytes": 0, "chunks": 0,
                                  "framer": Rtcm3Framer(),
                                  "types": collections.Counter()}
            s["bytes"] += len(data)
            s["chunks"] += 1
            if fmt.upper().startswith("RTCM_3") or fmt.upper().startswith("RTCM3"):
                for p in s["framer"].feed(data):
                    s["types"][msg_type(p)] += 1
        print("%s: %d chunks  %s -> %s" % (path, n, t0, t1))
        for sta, s in stats.items():
            print("  %-12s %-9s %10d bytes %6d chunks" %
                  (sta, s["fmt"], s["bytes"], s["chunks"]))
            if s["types"]:
                types = sorted(s["types"].items(),
                               key=lambda kv: (len(kv[0].split("_")[0]), kv[0]))
                line = "    "
                for k, v in types:
                    item = "%s:%d " % (k, v)
                    if len(line) + len(item) > 100:
                        print(line)
                        line = "    "
                    line += item
                print(line)


def frame(payload):
    """RTCM3 frame (preamble, length, payload, CRC24Q) of a payload."""
    head = bytes([0xD3, (len(payload) >> 8) & 0x03, len(payload) & 0xFF])
    return head + payload + crc24q(head + payload).to_bytes(3, "big")


def cmd_slice(args):
    stations = set(args.stations.split(",")) if args.stations else None
    types = args.types.split(",") if args.types else None
    parsers = {}

    def selected():
        for t, sta, fmt, data in read_chunks(args.input):
            if stations and sta not in stations:
                continue
            if args.start and t < args.start:
                continue
            if args.end and t >= args.end:
                continue
            if types:
                # keep only the RTCM3 messages of the given types (or type prefixes)
                frames = parsers.setdefault(sta, Rtcm3Framer()).feed(data)
                kept = b"".join(frame(pl) for pl in frames
                                if any(msg_type(pl).startswith(tt) for tt in types))
                if not kept:
                    continue
                data = kept
            yield t, sta, fmt, data

    write_chunks(args.output, selected())


def cmd_crs(args):
    """Distinct CRS message contents per stream, with their number."""
    stats = collections.OrderedDict()
    framers = {}
    for t, sta, fmt, data in read_chunks(args.input):
        for p in framers.setdefault(sta, Rtcm3Framer()).feed(data):
            if msg_type(p) in ("1300", "1301", "1302"):
                try:
                    line = decode_crs(p)
                except ValueError as e:
                    line = "%s %s" % (msg_type(p), e)
                key = (sta, line)
                stats[key] = stats.get(key, 0) + 1
    for (sta, line), n in stats.items():
        print("%s %4dx %s" % (sta, n, line))


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("info", help="summarise raw capture(s)")
    p.add_argument("files", nargs="+")
    p.set_defaults(func=cmd_info)

    p = sub.add_parser("slice", help="cut a raw capture")
    p.add_argument("input")
    p.add_argument("output")
    p.add_argument("--start", help="ISO time, inclusive (string compare, e.g. 2025-11-17T13:49)")
    p.add_argument("--end", help="ISO time, exclusive")
    p.add_argument("--stations", help="comma separated staIDs to keep")
    p.add_argument("--types", help="comma separated RTCM3 message types or prefixes to keep, "
                                   "e.g. 1019,4076_02 (GPS ephemerides and IGS-SSR GPS messages)")
    p.set_defaults(func=cmd_slice)

    p = sub.add_parser("crs", help="decode CRS messages 1300/1301/1302")
    p.add_argument("input")
    p.set_defaults(func=cmd_crs)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    sys.exit(main())