#!/usr/bin/env python3
"""
protocol_tool.py -- combined upload/download tool for the ESP32 Four
Channel Generator's protocols-N.txt files, meant to be edited in a
spreadsheet program (Excel, LibreOffice Calc, etc.) and saved as CSV/txt.

  protocol_tool.py download --dir ~/Desktop --file protocols-0.txt
  protocol_tool.py upload   --dir ~/Desktop --file protocols-0.txt

"download" = Generator -> PC.  "upload" = PC -> Generator (this is the
direction that runs the format/error-correction pass, since it's the file
a human -- or a spreadsheet program -- may have touched).
"""
import argparse
import os
import re
import sys
import time

import serial

PORT_DEFAULT = "/dev/ttyUSB0"
BAUD = 115200
TIMEOUT = 1.0
LINE_DELAY = 0.01
READY_TIMEOUT_SECONDS = 20
POST_WAIT_SECONDS = 5

BEGIN_MARKER = "BEGIN_FILE"
END_MARKER = "END_FILE"

# Ceiling on frequency records in one protocol file. The generator's EEPROM image
# holds 298 of them at most, because the control record has to fit after the last
# one; 275 is the documented limit, and it is enforced in three places that must
# stay in step -- here, in protocol_editor.py, and in the firmware's loadFloats()
# (MAX_FREQ_RECORDS in Generator_V2_05.ino). Uploading more used to be accepted:
# the file reached the card, and the generator only failed on the next boot.
MAX_RECORDS = 275
UPLOAD_READY_MARKER = "UPLOAD_READY"
DOWNLOAD_READY_MARKER = "DOWNLOAD_READY"


# =================================================================
#  Format / error correction for files coming FROM a spreadsheet
# =================================================================

def clean_field(raw):
    return raw.strip().strip('"').strip("'").strip()


FILE_NAME_RE = re.compile(r"^protocols-(\d{1,2})\.txt$", re.IGNORECASE)


def parse_file_no(filename):
    """Extract the Generator's SD-card file number (0-99) from a local
    filename. Upload/download both now target a specific protocols-N.txt on
    the SD card (not just whatever the device currently has active), and N
    is taken directly from this filename -- so it must conform exactly."""
    m = FILE_NAME_RE.match(filename)
    if not m:
        raise ValueError(
            f"{filename!r} must be named 'protocols-N.txt' (N = 0-99) -- "
            "the Generator uses this name to decide which SD card file to "
            "read or write."
        )
    n = int(m.group(1))
    if not (0 <= n <= 99):
        raise ValueError(f"{filename!r}: file number {n} is out of range (must be 0-99).")
    return n


def canonicalize_line(line, line_no):
    """Reformat one CSV line into exactly what the device expects, or
    raise ValueError with a human-readable reason it couldn't be fixed."""
    fields = [clean_field(f) for f in line.split(",")]
    count = len(fields)

    if count == 9:
        try:
            proto_id = int(float(fields[0]))
            values = [float(f) for f in fields[1:]]
        except ValueError as e:
            raise ValueError(f"line {line_no}: bad number in frequency record: {line!r} ({e})")
        parts = [str(proto_id)] + [f"{v:.2f}" for v in values]
        return ",".join(parts), "freq"

    if count == 11:
        try:
            ints_before = [int(float(f)) for f in fields[0:9]]
            sweep_inc = float(fields[9])
            file_no = int(float(fields[10]))
        except ValueError as e:
            raise ValueError(f"line {line_no}: bad number in control record: {line!r} ({e})")
        # field3/field4 are the frequency/program run times, stored as whole seconds/minutes
        # (matching the device's own screen-5 display and CSV convention) -- bound-check here
        # so a stray leftover raw-millisecond value (e.g. from an old file) is rejected with a
        # clear message instead of silently uploading as a huge run time.
        freq_time_secs, program_time_mins = ints_before[3], ints_before[4]
        if not (0 <= freq_time_secs <= 9999):
            raise ValueError(f"line {line_no}: frequency run time must be 0-9999 seconds, got {freq_time_secs}")
        if not (0 <= program_time_mins <= 9999):
            raise ValueError(f"line {line_no}: program run time must be 0-9999 minutes, got {program_time_mins}")
        parts = [str(v) for v in ints_before] + [f"{sweep_inc:.2f}", str(file_no)]
        return ",".join(parts), "ctrl"

    raise ValueError(f"line {line_no}: expected 9 or 11 fields, got {count}: {line!r}")


def load_and_correct(path):
    """Read a possibly spreadsheet-exported file and return a list of
    canonical CSV lines ready to send, or raise ValueError describing
    every problem found."""
    with open(path, "rb") as f:
        raw = f.read()

    if raw.startswith(b"\xef\xbb\xbf"):  # UTF-8 BOM some spreadsheets add
        raw = raw[3:]
    text = raw.decode("utf-8", errors="replace")

    freq_lines = []
    ctrl_line = None
    errors = []

    for i, raw_line in enumerate(text.splitlines(), start=1):
        line = raw_line.strip()
        if not line or line.startswith("#") or line in (BEGIN_MARKER, END_MARKER):
            continue
        if "," not in line:
            continue  # header lines like "Protocols-0 #" have no commas

        try:
            canon, kind = canonicalize_line(line, i)
        except ValueError as e:
            errors.append(str(e))
            continue

        if kind == "freq":
            freq_lines.append(canon)
        elif ctrl_line is not None:
            errors.append(f"line {i}: a second control record was found (only one allowed)")
        else:
            ctrl_line = canon

    if errors:
        raise ValueError("Format problems found:\n  " + "\n  ".join(errors))
    if not freq_lines:
        raise ValueError("No frequency records (9-field lines) found.")
    if len(freq_lines) > MAX_RECORDS:
        raise ValueError(
            f"Too many frequency records ({len(freq_lines)}); the generator holds "
            f"at most {MAX_RECORDS}. Split this into more than one protocol file."
        )
    if ctrl_line is None:
        raise ValueError("No control record (11-field line) found -- every protocol file needs exactly one.")

    return freq_lines + [ctrl_line]


# =================================================================
#  Serial helpers
# =================================================================

def rx_line(ser):
    raw = ser.readline()
    if not raw:
        return None
    return raw.decode("utf-8", errors="replace").rstrip("\r\n")


def drain_rx(ser, log=print):
    while ser.in_waiting:
        line = rx_line(ser)
        if line:
            log(f"RX: {line}")


def wait_for_marker(ser, marker, timeout_seconds, log=print):
    log(f"Waiting up to {timeout_seconds}s for {marker}...")
    start_time = time.time()
    while time.time() - start_time < timeout_seconds:
        line = rx_line(ser)
        if not line:
            continue
        log(f"RX: {line}")
        if marker in line.upper():
            return True
    return False


def tx_line(ser, text, log=print):
    ser.write((text + "\n").encode("utf-8"))
    ser.flush()
    log(f"TX: {text}")


# =================================================================
#  Upload: PC -> Generator
# =================================================================

def do_upload(path, port, log=print):
    try:
        file_no = parse_file_no(os.path.basename(path))
    except ValueError as e:
        log(str(e))
        return 1

    try:
        lines = load_and_correct(path)
    except FileNotFoundError:
        log(f"Source file not found: {path}")
        return 1
    except ValueError as e:
        log(f"Cannot upload -- file needs fixing:\n{e}")
        return 1

    log(f"Loaded and validated {len(lines)} data lines from {path}")
    log(f"Opening {port} at {BAUD} baud...")

    try:
        with serial.Serial(port, BAUD, timeout=TIMEOUT) as ser:
            time.sleep(2.0)
            ser.reset_input_buffer()
            ser.reset_output_buffer()

            if not wait_for_marker(ser, UPLOAD_READY_MARKER, READY_TIMEOUT_SECONDS, log=log):
                log(f"Timed out waiting for {UPLOAD_READY_MARKER}. Is the generator in RECEIVE mode?")
                return 1

            tx_line(ser, f"{BEGIN_MARKER}:{file_no}", log=log)
            time.sleep(LINE_DELAY)
            drain_rx(ser, log=log)

            for index, line in enumerate(lines, start=1):
                tx_line(ser, line, log=log)
                time.sleep(LINE_DELAY)
                drain_rx(ser, log=log)

            tx_line(ser, END_MARKER, log=log)

            log("Waiting briefly for ESP32 response...")
            ok = False
            end_time = time.time() + POST_WAIT_SECONDS
            while time.time() < end_time:
                line = rx_line(ser)
                if not line:
                    continue
                log(f"RX: {line}")
                if "UPLOAD_OK" in line.upper():
                    ok = True

    except serial.SerialException as e:
        log(f"Serial error: {e}")
        if "Errno 16" in str(e) or "resource busy" in str(e).lower():
            log("Hint: close the IDE serial monitor, minicom, or any other app using the serial port.")
        return 1

    log("Upload finished." if ok else "Upload finished, but no UPLOAD_OK seen -- check RX output above.")
    return 0 if ok else 1


# =================================================================
#  Download: Generator -> PC
# =================================================================

def do_download(path, port, log=print):
    try:
        file_no = parse_file_no(os.path.basename(path))
    except ValueError as e:
        log(str(e))
        return 1

    log(f"Opening {port} at {BAUD} baud...")
    captured = []

    try:
        with serial.Serial(port, BAUD, timeout=TIMEOUT) as ser:
            time.sleep(2.0)
            ser.reset_input_buffer()
            ser.reset_output_buffer()

            if not wait_for_marker(ser, DOWNLOAD_READY_MARKER, READY_TIMEOUT_SECONDS, log=log):
                log(f"Timed out waiting for {DOWNLOAD_READY_MARKER}. Is the generator in SEND mode?")
                return 1

            tx_line(ser, f"READY:{file_no}", log=log)

            if not wait_for_marker(ser, BEGIN_MARKER, READY_TIMEOUT_SECONDS, log=log):
                log(f"Timed out waiting for {BEGIN_MARKER}.")
                return 1

            while True:
                line = rx_line(ser)
                if line is None:
                    continue
                log(f"RX: {line}")
                if END_MARKER in line.upper():
                    break
                if line:
                    captured.append(line)

            drain_rx(ser, log=log)  # pick up trailing DOWNLOAD_OK if it's already buffered

    except serial.SerialException as e:
        log(f"Serial error: {e}")
        if "Errno 16" in str(e) or "resource busy" in str(e).lower():
            log("Hint: close the IDE serial monitor, minicom, or any other app using the serial port.")
        return 1

    if not captured:
        log("No file data was captured.")
        return 1

    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(captured) + "\n")

    log(f"Saved {len(captured)} data lines to {path}")
    return 0


# =================================================================
#  CLI
# =================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Send/receive ESP32 Four Channel Generator protocol files over USB."
    )
    parser.add_argument("mode", choices=["upload", "download"],
                         help="upload: PC -> Generator (validates/corrects first). "
                              "download: Generator -> PC.")
    parser.add_argument("--file", required=True,
                         help="Protocol file name, e.g. protocols-0.txt")
    parser.add_argument("--dir", default=".",
                         help="Directory the file lives in (upload) or will be written to "
                              "(download). Default: current directory.")
    parser.add_argument("--port", default=PORT_DEFAULT,
                         help=f"Serial port (default: {PORT_DEFAULT})")
    args = parser.parse_args()

    path = os.path.join(args.dir, args.file)

    if args.mode == "upload":
        return do_upload(path, args.port)
    else:
        return do_download(path, args.port)


if __name__ == "__main__":
    raise SystemExit(main())