#!/usr/bin/env python3
"""NEIL.VISION read-only bridge for scientific / tabular tracking containers.

Decode a supported container into UTF-8 JSONL. Coordinates and time keep their
source meanings; declare units, axes and time in the viewer's mapping review.
Never imports pickle, evaluates source code, or opens a source for writing.
Run --help. Optional readers are installed only for formats you use.
"""
from __future__ import annotations
import argparse
import bz2
import gzip
import lzma
import base64
import csv
import datetime as dt
import importlib
import json
import math
import os
from pathlib import Path
import sqlite3
import sys
import tempfile


def require(name, distribution=None):
    try:
        return importlib.import_module(name)
    except ImportError as exc:
        raise ValueError(f"Install this reader first: python -m pip install {distribution or name}") from exc


def clean(value):
    if value is None or isinstance(value, (str, bool)):
        return value
    if isinstance(value, int):
        return str(value) if abs(value) > 2**53 - 1 else value
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, (dt.datetime, dt.date, dt.time)):
        return value.isoformat()
    if isinstance(value, bytes):
        return {"base64": base64.b64encode(value).decode("ascii")}
    if isinstance(value, dict):
        return {str(k): clean(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [clean(v) for v in value]
    if hasattr(value, "tolist"):
        return clean(value.tolist())
    if hasattr(value, "item"):
        return clean(value.item())
    raise ValueError(f"Unsupported value {type(value).__name__}; select a numeric dataset.")


def array_rows(array, args):
    np = require("numpy")
    if array.dtype.hasobject:
        raise ValueError("Object / pickle arrays are disabled. Export numeric or text arrays.")
    if array.dtype.names:
        for row in array.reshape(-1):
            yield {name: clean(row[name]) for name in array.dtype.names}
        return
    if array.ndim == 2 and not args.axes:
        for row in array:
            yield {f"column_{i+1}": clean(v) for i, v in enumerate(row)}
        return
    axes = args.axes
    required = "tejc" if array.ndim == 4 else "tec"
    if not axes or len(axes) != array.ndim or set(axes) != set(required):
        raise ValueError(f"Declare array axes with --axes {required} (t=time, e=entity, j=joint, c=coordinates). Use the actual source order.")
    array = np.transpose(array, [axes.index(c) for c in required])
    if array.shape[-1] not in (2, 3):
        raise ValueError("The coordinate axis must have length 2 or 3. Confidence is a separate field.")
    ids = json.loads(Path(args.ids).read_text()) if args.ids else [str(i) for i in range(array.shape[1])]
    times = json.loads(Path(args.times).read_text()) if args.times else None
    joints = json.loads(Path(args.joints).read_text()) if args.joints else None
    if len(ids) != array.shape[1] or times is not None and len(times) != array.shape[0]:
        raise ValueError("Identity / time array length does not match its declared axis.")
    if len(set(map(str,ids))) != len(ids) or any(value is None for value in ids):
        raise ValueError("Entity IDs must be non-null and unique along the entity axis.")
    if array.ndim == 4 and (not joints or len(joints) != array.shape[2]):
        raise ValueError("A pose tensor needs --joints names.json matching its joint axis.")
    if joints and (any(name is not None and (not isinstance(name,str) or not name) for name in joints) or len(set(name for name in joints if name is not None)) != len([name for name in joints if name is not None])):
        raise ValueError("Joint names must be unique non-empty strings (null explicitly skips an index).")
    for t in range(array.shape[0]):
        for e in range(array.shape[1]):
            base = {"frame": t, "entity_id": str(ids[e])}
            if times is not None:
                base["time"] = times[t]
            if array.ndim == 4:
                pose = {}
                for j, name in enumerate(joints):
                    if name is None:
                        continue
                    q = array[t, e, j]
                    pose[name] = {"x": clean(q[0]), "y": clean(q[1]), "z": clean(q[2]) if len(q) == 3 else None}
                yield {**base, "joints": pose}
            else:
                q = array[t, e]
                yield {**base, "x": clean(q[0]), "y": clean(q[1]), "z": clean(q[2]) if len(q) == 3 else None}


def records(source, args):
    ext = source.suffix.lower()
    if ext in (".gz", ".bz2", ".xz", ".zst", ".zstd", ".lz4"):
        if ext == ".gz":
            stream = gzip.open(source, "rb")
        elif ext == ".bz2":
            stream = bz2.open(source, "rb")
        elif ext == ".xz":
            stream = lzma.open(source, "rb")
        elif ext in (".zst", ".zstd"):
            stream = require("zstandard").ZstdDecompressor().stream_reader(source.open("rb"))
        else:
            stream = require("lz4.frame", "lz4").open(source, "rb")
        suffix = Path(source.stem).suffix
        with tempfile.TemporaryDirectory() as folder, stream as reader:
            decoded = Path(folder)/("decoded"+suffix)
            total = 0
            with decoded.open("wb") as writer:
                while chunk := reader.read(1024*1024):
                    total += len(chunk)
                    if total > args.max_bytes:
                        raise ValueError("Decompression exceeds --max-bytes. Export segments.")
                    writer.write(chunk)
            yield from records(decoded, args)
        return
    if ext in (".pkl", ".pickle", ".pt", ".pth", ".rds", ".rda", ".rdata"):
        raise ValueError("Executable or runtime-specific serialization is disabled. Re-export with the producing application.")
    if ext in (".csv", ".tsv", ".txt", ".dat", ".fwf"):
        with source.open(encoding=args.encoding, newline="") as handle:
            if args.widths or args.whitespace:
                names = json.loads(args.columns) if args.columns else None
                widths = [int(x) for x in args.widths.split(",")] if args.widths else None
                if widths and any(w < 1 for w in widths):
                    raise ValueError("Fixed-width column sizes must be positive.")
                for n, line in enumerate(handle, 1):
                    if not line.strip() or line.lstrip().startswith("#"):
                        continue
                    if widths:
                        cells, offset = [], 0
                        for width in widths:
                            cells.append(line[offset:offset+width].strip())
                            offset += width
                    else:
                        cells = line.split()
                    if names is None:
                        names = cells
                        continue
                    if len(cells) != len(names):
                        raise ValueError(f"Line {n}: expected {len(names)} cells; found {len(cells)}")
                    yield dict(zip(names, cells))
            else:
                yield from csv.DictReader(handle, delimiter=args.delimiter or ("\t" if ext == ".tsv" else ","))
        return
    if ext in (".sqlite", ".sqlite3", ".db"):
        with sqlite3.connect(source.resolve().as_uri()+"?mode=ro&immutable=1", uri=True) as db:
            db.execute("PRAGMA query_only=ON")
            db.execute("PRAGMA trusted_schema=OFF")
            names = [r[0] for r in db.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")]
            if args.list:
                print(json.dumps({"tables": names}, indent=2)); return
            if not args.table or args.table not in names:
                raise ValueError("Choose an existing --table. Use --list to see names.")
            cursor = db.execute('SELECT * FROM "'+args.table.replace('"', '""')+'"')
            columns = [r[0] for r in cursor.description]
            for row in cursor:
                yield dict(zip(columns, row))
        return
    if ext in (".parquet", ".arrow", ".feather", ".ipc"):
        if ext == ".parquet":
            batches = require("pyarrow.parquet", "pyarrow").ParquetFile(source).iter_batches(batch_size=10000)
        else:
            pa = require("pyarrow")
            data = pa.memory_map(str(source), "r")
            try:
                reader = pa.ipc.open_file(data)
                batches = (reader.get_batch(i) for i in range(reader.num_record_batches))
            except pa.ArrowInvalid:
                data.seek(0)
                batches = pa.ipc.open_stream(data)
        for batch in batches:
            yield from batch.to_pylist()
        return
    if ext == ".avro":
        reader = require("fastavro")
        with source.open("rb") as handle:
            yield from reader.reader(handle)
        return
    if ext in (".npy", ".npz"):
        np = require("numpy")
        data = np.load(source, allow_pickle=False, mmap_mode="r" if ext == ".npy" else None)
        if ext == ".npz":
            try:
                if args.list:
                    print(json.dumps({"arrays": data.files}, indent=2)); return
                if not args.dataset or args.dataset not in data.files:
                    raise ValueError("Choose --dataset. Use --list to see array names.")
                yield from array_rows(data[args.dataset], args)
            finally:
                data.close()
        else:
            if args.list:
                print(json.dumps({"shape": data.shape, "dtype": str(data.dtype)})); return
            yield from array_rows(data, args)
        return
    if ext in (".h5", ".hdf5", ".mat"):
        h5 = require("h5py")
        if h5.is_hdf5(source):
            with h5.File(source, "r") as file:
                if args.list:
                    found = {}
                    file.visititems(lambda name, obj: found.update({name: {"shape": obj.shape, "dtype": str(obj.dtype)}}) if isinstance(obj, h5.Dataset) else None)
                    print(json.dumps(found, indent=2)); return
                if not args.dataset or args.dataset not in file or not isinstance(file[args.dataset], h5.Dataset):
                    raise ValueError("Choose one numeric --dataset. Use --list first.")
                for i in range(1,len(args.dataset.strip("/").split("/"))+1):
                    prefix="/".join(args.dataset.strip("/").split("/")[:i])
                    if isinstance(file.get(prefix,getlink=True),h5.ExternalLink):
                        raise ValueError("External HDF links require a self-contained export.")
                dataset = file[args.dataset]
                if dataset.is_virtual or dataset.external:
                    raise ValueError("External / virtual HDF datasets require a self-contained export.")
                # HDF is read-only. This bridge uses bounded source files, not arbitrary server execution.
                if dataset.nbytes > args.max_bytes:
                    raise ValueError("Dataset exceeds --max-bytes. Export time segments with the producing tool.")
                yield from array_rows(dataset[()], args)
            return
        if ext != ".mat":
            raise ValueError("Unrecognized HDF5 container.")
        scipy = require("scipy.io", "scipy")
        if args.list:
            print(json.dumps(scipy.whosmat(source), indent=2)); return
        if not args.dataset:
            raise ValueError("Choose a numeric --dataset variable. Use --list first.")
        variables = scipy.loadmat(source, variable_names=[args.dataset], struct_as_record=True)
        if args.dataset not in variables:
            raise ValueError("MATLAB variable not found.")
        yield from array_rows(variables[args.dataset], args)
        return
    raise ValueError("No dedicated reader for this container. Export JSONL / CSV / Parquet with its producer. Protobuf needs its .proto or descriptor and message framing.")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("source", type=Path)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--list", action="store_true", help="List datasets / tables without exporting")
    parser.add_argument("--dataset", help="HDF dataset, MATLAB variable, or NPZ array")
    parser.add_argument("--table", help="Existing SQLite table")
    parser.add_argument("--axes", help="Explicit tensor-axis order, e.g. tec or etjc")
    parser.add_argument("--ids", help="JSON array of entity IDs, matching entity axis")
    parser.add_argument("--times", help="JSON array of time values, matching time axis")
    parser.add_argument("--joints", help="JSON array of joint names, matching joint axis")
    parser.add_argument("--encoding", default="utf-8-sig")
    parser.add_argument("--delimiter")
    parser.add_argument("--whitespace", action="store_true", help="Split whitespace fields without CSV quoting")
    parser.add_argument("--widths", help="Fixed-width fields, e.g. 8,16,12,12")
    parser.add_argument("--columns", help='Headerless text column names as JSON, e.g. ["time","id","x","y"]')
    parser.add_argument("--max-bytes", type=int, default=2*1024**3)
    parser.add_argument("--max-rows", type=int, default=10_000_000)
    args = parser.parse_args()
    temp = None
    try:
        if not args.source.is_file() or args.source.stat().st_size > args.max_bytes:
            raise ValueError("Source is missing or exceeds --max-bytes.")
        if args.list:
            for _ in records(args.source, args):
                pass
            return
        if args.out is None:
            raise ValueError("Provide --out tracking.jsonl, or use --list.")
        if args.out.resolve() == args.source.resolve() or args.out.exists():
            raise ValueError("Output already exists or matches the source. Choose a new output path.")
        args.out.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", newline="\n", dir=args.out.parent, delete=False) as handle:
            temp = Path(handle.name)
            count = 0
            for row in records(args.source, args):
                count += 1
                if count > args.max_rows:
                    raise ValueError("Output exceeds --max-rows. Export time segments.")
                handle.write(json.dumps(clean(row), ensure_ascii=False, allow_nan=False, separators=(",", ":"))+"\n")
            handle.flush()
            os.fsync(handle.fileno())
        if not count:
            raise ValueError("No rows found; no output written.")
        # Hard-link creation fails if another process created the output in the meantime.
        os.link(temp, args.out)
        temp.unlink(); temp = None
        print(f"Wrote {count:,} rows to {args.out}. Review the source units, time basis and coordinate frame in NEIL.VISION.")
    except (ValueError, OSError, KeyError, TypeError) as exc:
        print(f"Conversion stopped: {exc}", file=sys.stderr)
        return 2
    finally:
        if temp:
            temp.unlink(missing_ok=True)
    return 0


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