#!/opt/srasenv/bin/python3
"""
SRAS Scan Manager
Command-line / interactive TUI for inspecting .sras files.

A .sras file (see scan_format.md) holds one acquisition run across several
GR rotation angles, each with its own geometry (x_start, x_delta, n_frames,
n_rows), background waveform, and waveform data block. This tool lists those
per-angle sub-scans and lets you export a subset to a new .sras file, or
delete a subset from the file in place — both operations rewrite the
angle/geometry/row tables, carry each kept angle's background across, and
stream-copy only the selected angles' waveform data, producing a file that is
itself a valid .sras readable by sras_viewer.py or sc3_aui_app.py.

Format versions 7 (full scan) and 11 (middle-row SAW check) are supported,
as are their pre-per-angle-background predecessors 6 and 10. A subset keeps
the version — and therefore the background layout — of the file it came
from: a v11 check exports as a v11 check, since dropping angles from one
leaves it one row per angle.
"""

import argparse
import struct
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from core.sras_format import (
    BG_LEN_FMT, GEOM_FMT, HDR_FMT, MAGIC, SrasFile,
)


@dataclass
class AngleEntry:
    index: int
    angle_deg: float
    x_start: float
    x_delta: float
    n_frames: int
    n_rows_declared: int
    y_positions: list  # declared length; may exceed what's actually on disk
    row_bytes: int
    background: bytes            # this angle's own background (v7/v11)
    data_offset: int             # byte offset into the file where this angle's rows start
    n_rows_available: int = 0
    data_size_available: int = 0
    complete: bool = True

    @property
    def data_size_declared(self) -> int:
        return self.row_bytes * self.n_rows_declared


class SrasScanFile:
    """Parsed view of a .sras file's header/tables plus per-angle data offsets."""

    def __init__(self, path: Path):
        self.path = Path(path)
        self._parse()

    def _parse(self):
        sras = SrasFile(self.path)
        h = sras.header
        self.version = sras.version
        self.is_saw_check = sras.is_saw_check
        self.x_start_nominal = h.x_start_nominal
        self.y_start_nominal = h.y_start_nominal
        self.x_delta_nominal = h.x_delta_nominal
        self.y_delta_nominal = h.y_delta_nominal
        self.row_spacing_mm  = h.row_spacing
        self.velocity_mm_s   = h.velocity
        self.laser_freq_hz   = h.laser_freq
        self.samples_per_frame = h.samples_per_frame
        self.sample_rate_hz  = h.sample_rate
        self.bytes_per_sample = h.bytes_per_sample
        self.n_channels       = h.n_channels
        self.preambles_raw   = sras.preambles_raw
        self.legacy_layout   = sras.is_legacy_layout
        # v6/v10 keep one background ahead of the data block; v7/v11 keep one
        # per angle inside it.  Either way sras.backgrounds is per angle.
        self.shared_background = sras.backgrounds[0] if sras.is_legacy_layout else b""
        self.data_start_offset = sras.data_start_offset
        self.file_size       = sras.file_size

        self.angles = [
            AngleEntry(
                index=st.index, angle_deg=st.angle_deg,
                x_start=pa.x_start, x_delta=pa.x_delta,
                n_frames=pa.n_frames, n_rows_declared=pa.n_rows,
                y_positions=pa.y_positions, row_bytes=st.row_bytes,
                background=bg, data_offset=st.data_offset,
                n_rows_available=st.n_rows_available,
                data_size_available=st.n_rows_available * st.row_bytes,
                complete=st.complete,
            )
            for pa, st, bg in zip(sras.per_angle, sras.angle_status(),
                                  sras.backgrounds, strict=True)
        ]

    def get(self, index: int) -> AngleEntry:
        return self.angles[index]


# ---------------------------------------------------------------------------
# Export / delete
# ---------------------------------------------------------------------------

def _write_subset(sf: SrasScanFile, indices: list, dst_path: Path) -> list:
    """Write a new .sras file containing only the given angle indices
    (in the given order). Returns a list of warning strings (e.g. for
    angles that were truncated on disk and thus exported with fewer rows
    than declared).
    """
    warnings = []
    selected = [sf.get(i) for i in indices]

    header = struct.pack(
        HDR_FMT, MAGIC, sf.version, len(selected),
        sf.x_start_nominal, sf.y_start_nominal,
        sf.x_delta_nominal, sf.y_delta_nominal,
        sf.row_spacing_mm, sf.velocity_mm_s, sf.laser_freq_hz,
        sf.samples_per_frame, sf.sample_rate_hz,
        sf.bytes_per_sample, sf.n_channels,
    )

    with open(sf.path, "rb") as src, open(dst_path, "wb") as dst:
        dst.write(header)
        dst.write(struct.pack(f">{len(selected)}f", *[e.angle_deg for e in selected]))

        for e in selected:
            if e.n_rows_available != e.n_rows_declared:
                warnings.append(
                    f"angle[{e.index}] ({e.angle_deg:.2f} deg): declared "
                    f"{e.n_rows_declared} rows but only {e.n_rows_available} "
                    f"present on disk — exporting truncated")
            dst.write(struct.pack(GEOM_FMT, e.x_start, e.x_delta, e.n_frames,
                                   e.n_rows_available))

        for e in selected:
            ys = e.y_positions[:e.n_rows_available]
            dst.write(struct.pack(f">{len(ys)}f", *ys))

        for praw in sf.preambles_raw:
            dst.write(struct.pack(">H", len(praw)))
            dst.write(praw)

        if sf.legacy_layout:
            dst.write(struct.pack(BG_LEN_FMT, len(sf.shared_background)))
            dst.write(sf.shared_background)

        for e in selected:
            if not sf.legacy_layout:
                if not e.background:
                    warnings.append(
                        f"angle[{e.index}] ({e.angle_deg:.2f} deg): no background "
                        "on disk — exported with an empty background block")
                dst.write(struct.pack(BG_LEN_FMT, len(e.background)))
                dst.write(e.background)
            src.seek(e.data_offset)
            remaining = e.data_size_available
            chunk_size = 1 << 20
            while remaining > 0:
                chunk = src.read(min(chunk_size, remaining))
                if not chunk:
                    break
                dst.write(chunk)
                remaining -= len(chunk)

    return warnings


def export_angles(sf: SrasScanFile, indices: list, dst_path: Path) -> list:
    if not indices:
        raise ValueError("no angles selected to export")
    return _write_subset(sf, indices, dst_path)


def delete_angles(sf: SrasScanFile, indices_to_delete: list, backup: bool = True) -> Path | None:
    """Rewrite sf.path in place, keeping every angle NOT in indices_to_delete.

    Returns the backup file path if one was made, else None.
    """
    keep = [e.index for e in sf.angles if e.index not in set(indices_to_delete)]
    if not keep:
        raise ValueError("refusing to delete every angle — a .sras file needs at least one")

    tmp_path = sf.path.with_suffix(sf.path.suffix + ".tmp")
    _write_subset(sf, keep, tmp_path)

    backup_path = None
    if backup:
        stamp = datetime.now().strftime("%Y%m%d%H%M%S")
        backup_path = sf.path.with_name(f"{sf.path.stem}.bak-{stamp}{sf.path.suffix}")
        sf.path.rename(backup_path)

    tmp_path.replace(sf.path)
    return backup_path


# ---------------------------------------------------------------------------
# Formatting helpers
# ---------------------------------------------------------------------------

def _human_size(n: int) -> str:
    size = float(n)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if size < 1024.0:
            return f"{size:.1f} {unit}"
        size /= 1024.0
    return f"{size:.1f} PB"


def parse_index_spec(spec: str, max_index: int) -> list:
    """Parse '0,2,4-6' or 'all' into a sorted list of unique in-range indices."""
    spec = spec.strip().lower()
    if spec in ("all", "*"):
        return list(range(max_index + 1))
    if not spec:
        return []
    out = set()
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            lo, hi = part.split("-", 1)
            lo, hi = int(lo), int(hi)
            if lo > hi:
                lo, hi = hi, lo
            for i in range(lo, hi + 1):
                out.add(i)
        else:
            out.add(int(part))
    bad = [i for i in out if i < 0 or i > max_index]
    if bad:
        raise ValueError(f"index out of range (0-{max_index}): {sorted(bad)}")
    return sorted(out)


def print_summary(sf: SrasScanFile, selected: set):
    print()
    kind = " SAW check" if sf.is_saw_check else ""
    print(f"File: {sf.path}  (v{sf.version}{kind}, {_human_size(sf.file_size)})")
    print(f"Nominal ROI: x_start={sf.x_start_nominal:.4f} x_delta={sf.x_delta_nominal:.4f} "
          f"y_start={sf.y_start_nominal:.4f} y_delta={sf.y_delta_nominal:.4f} mm  "
          f"row_spacing={sf.row_spacing_mm:.4f} mm")
    print(f"Velocity={sf.velocity_mm_s:.2f} mm/s  Laser={sf.laser_freq_hz:.1f} Hz  "
          f"Samples/frame={sf.samples_per_frame}  Sample rate={sf.sample_rate_hz:.3e} Hz  "
          f"Channels={sf.n_channels}  Bytes/sample={sf.bytes_per_sample}")
    print()
    hdr = f"{'':>2} {'#':>3} {'Angle(deg)':>10} {'Rows':>7} {'Frames/row':>10} {'x_start':>9} {'x_delta':>9} {'Data size':>11}  Status"
    print(hdr)
    print("-" * len(hdr))
    for e in sf.angles:
        mark = "*" if e.index in selected else " "
        if e.complete:
            status = "OK"
        elif e.n_rows_available == 0:
            status = "MISSING (no data on disk)"
        else:
            status = f"TRUNCATED ({e.n_rows_available}/{e.n_rows_declared} rows on disk)"
        print(f"{mark:>2} {e.index:>3} {e.angle_deg:>10.2f} {e.n_rows_declared:>7} "
              f"{e.n_frames:>10} {e.x_start:>9.3f} {e.x_delta:>9.3f} "
              f"{_human_size(e.data_size_available):>11}  {status}")
    print()


# ---------------------------------------------------------------------------
# Interactive TUI
# ---------------------------------------------------------------------------

def interactive_loop(path: Path):
    sf = SrasScanFile(path)
    selected: set = set()

    help_text = (
        "  [l] list                 show the angle table again\n"
        "  [s] select <spec>        set selection, e.g. '0,2,4-6' or 'all' or 'none'\n"
        "  [e] export <path>        write selected angles to a new .sras file\n"
        "  [d] delete               remove selected angles from this file in place\n"
        "  [r] reload               re-read the file from disk (after external changes)\n"
        "  [h] help                 show this help\n"
        "  [q] quit"
    )

    print_summary(sf, selected)
    print(help_text)

    while True:
        try:
            cmd_line = input("\nsras> ").strip()
        except EOFError:
            print()
            break
        if not cmd_line:
            continue
        parts = cmd_line.split(None, 1)
        cmd = parts[0].lower()
        arg = parts[1].strip() if len(parts) > 1 else ""

        try:
            if cmd in ("q", "quit", "exit"):
                break
            elif cmd in ("h", "help", "?"):
                print(help_text)
            elif cmd in ("l", "list"):
                print_summary(sf, selected)
            elif cmd in ("s", "select"):
                if not arg:
                    arg = input("Angles to select (e.g. 0,2,4-6 / all / none): ").strip()
                if arg.lower() == "none":
                    selected = set()
                else:
                    selected = set(parse_index_spec(arg, len(sf.angles) - 1))
                print(f"Selected {len(selected)} angle(s): {sorted(selected)}")
            elif cmd in ("e", "export"):
                if not selected:
                    print("Nothing selected — use 's' first.")
                    continue
                dst = arg or input("Output path: ").strip()
                if not dst:
                    print("Export cancelled — no path given.")
                    continue
                dst_path = Path(dst)
                if dst_path.exists():
                    ans = input(f"{dst_path} exists — overwrite? [y/N] ").strip().lower()
                    if ans != "y":
                        print("Export cancelled.")
                        continue
                warnings = export_angles(sf, sorted(selected), dst_path)
                print(f"Exported {len(selected)} angle(s) -> {dst_path}")
                for w in warnings:
                    print(f"  warning: {w}")
            elif cmd in ("d", "delete"):
                if not selected:
                    print("Nothing selected — use 's' first.")
                    continue
                print(f"About to delete {len(selected)} angle(s) from {sf.path}: {sorted(selected)}")
                ans = input("Type 'yes' to confirm (a timestamped .bak copy will be kept): ").strip()
                if ans != "yes":
                    print("Delete cancelled.")
                    continue
                backup_path = delete_angles(sf, sorted(selected), backup=True)
                print(f"Deleted. Backup saved to {backup_path}")
                sf = SrasScanFile(sf.path)
                selected = set()
                print_summary(sf, selected)
            elif cmd in ("r", "reload"):
                sf = SrasScanFile(sf.path)
                selected = set()
                print_summary(sf, selected)
            else:
                print(f"Unknown command: {cmd!r} (type 'h' for help)")
        except Exception as exc:
            print(f"Error: {exc}")


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(
        description="Inspect, export, or delete per-angle sub-scans in a .sras file.")
    ap.add_argument("file", type=Path, help="path to a .sras file")
    ap.add_argument("--list", action="store_true", help="print the angle table and exit")
    ap.add_argument("--export", metavar="SPEC", help="angle index spec to export, e.g. '0,2,4-6' or 'all'")
    ap.add_argument("--output", metavar="PATH", type=Path, help="destination path for --export")
    ap.add_argument("--delete", metavar="SPEC", help="angle index spec to delete in place, e.g. '1,3'")
    ap.add_argument("--no-backup", action="store_true", help="skip the .bak copy when using --delete")
    ap.add_argument("--yes", action="store_true", help="don't prompt for confirmation on --delete")
    args = ap.parse_args()

    if not args.file.exists():
        print(f"error: {args.file} does not exist", file=sys.stderr)
        sys.exit(1)

    try:
        sf = SrasScanFile(args.file)
    except ValueError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(1)

    non_interactive = args.list or args.export or args.delete

    if args.list:
        print_summary(sf, set())

    try:
        if args.export:
            indices = parse_index_spec(args.export, len(sf.angles) - 1)
            if not args.output:
                print("error: --export requires --output", file=sys.stderr)
                sys.exit(1)
            if args.output.exists() and not args.yes:
                ans = input(f"{args.output} exists — overwrite? [y/N] ").strip().lower()
                if ans != "y":
                    print("Export cancelled.")
                    sys.exit(1)
            warnings = export_angles(sf, indices, args.output)
            print(f"Exported {len(indices)} angle(s) -> {args.output}")
            for w in warnings:
                print(f"  warning: {w}")

        if args.delete:
            indices = parse_index_spec(args.delete, len(sf.angles) - 1)
            if not indices:
                print("error: --delete requires a non-empty angle spec", file=sys.stderr)
                sys.exit(1)
            if not args.yes:
                print(f"About to delete {len(indices)} angle(s) from {sf.path}: {indices}")
                ans = input("Type 'yes' to confirm: ").strip()
                if ans != "yes":
                    print("Delete cancelled.")
                    sys.exit(1)
            backup_path = delete_angles(sf, indices, backup=not args.no_backup)
            if backup_path:
                print(f"Deleted. Backup saved to {backup_path}")
            else:
                print("Deleted (no backup kept).")
    except ValueError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(1)

    if not non_interactive:
        interactive_loop(args.file)


if __name__ == "__main__":
    main()