#!/usr/bin/env python3
"""Minimal SCPI helper for Hantek DPO7204C over /dev/usbtmc2."""
import os
import sys
import time

DEV = "/dev/usbtmc2"


class Scope:
    def __init__(self, dev=DEV):
        self.fd = os.open(dev, os.O_RDWR)

    def write(self, cmd: str):
        os.write(self.fd, cmd.encode() + b"\n")

    def read(self, n=1 << 20, timeout=3.0) -> bytes:
        # usbtmc read returns one transfer chunk; loop until short read
        chunks = []
        end = time.time() + timeout
        while time.time() < end:
            try:
                data = os.read(self.fd, n)
            except OSError:
                break
            chunks.append(data)
            if not data or len(data) < n:
                break
        return b"".join(chunks)

    def query(self, cmd: str, timeout=3.0) -> str:
        self.write(cmd)
        return self.read(timeout=timeout).decode(errors="replace").strip()

    def query_raw(self, cmd: str, timeout=5.0) -> bytes:
        self.write(cmd)
        return self.read(timeout=timeout)

    def close(self):
        os.close(self.fd)


if __name__ == "__main__":
    s = Scope()
    for cmd in sys.argv[1:]:
        if cmd.endswith("?"):
            print(f"{cmd:40s} -> {s.query(cmd)}")
        else:
            s.write(cmd)
            print(f"{cmd:40s} [sent]")
    s.close()