#!/opt/srasenv/bin/python3
"""
SRAS Scan Manager
Command-line / interactive TUI for inspecting v6 .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) 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
and stream-copy only the selected angles' waveform data, producing a file
that is itself a valid v6 .sras readable by sras_viewer.py-style tools
(once updated for v6) or sc3_aui_app.py.

Only format version 6 is supported.
"""

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 GEOM_FMT, HDR_FMT, MAGIC, VERSION as BLOB_VERSION, 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
    data_offset: int             # byte offset into the file where this angle's data starts
    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 v6 .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.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.background_raw  = sras.background
        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,
                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 in zip(sras.per_angle, sras.angle_status(), 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 v6 .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, BLOB_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)

        dst.write(struct.pack(">I", len(sf.background_raw)))
        dst.write(sf.background_raw)

        for e in selected:
            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()
    print(f"File: {sf.path}  (v{BLOB_VERSION}, {_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 v6 .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()