#!/usr/bin/env python3
"""Upload a file to the CYD TF card over serial (#PUT + ACK window)."""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

try:
    import serial
    from serial.tools import list_ports
except ImportError:
    print("Need pyserial: pip install pyserial", file=sys.stderr)
    sys.exit(1)

WIN = 4096


def pick_port(preferred: str | None) -> str:
    if preferred:
        return preferred
    ports = list(list_ports.comports())
    for p in ports:
        desc = (p.description or "") + (p.hwid or "")
        if "USB" in desc.upper() or "UART" in desc.upper() or "CP210" in desc.upper() or "CH340" in desc.upper():
            return p.device
    if ports:
        return ports[0].device
    raise SystemExit("No serial port found")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("file", type=Path, help="Local file to upload")
    ap.add_argument("--port", default=None)
    ap.add_argument("--baud", type=int, default=460800)
    ap.add_argument("--remote", default="/paomo.mp3")
    args = ap.parse_args()

    path: Path = args.file
    if not path.is_file():
        print(f"Missing file: {path}", file=sys.stderr)
        return 1

    size = path.stat().st_size
    port = pick_port(args.port)
    print(f"Port {port}  {path.name} -> {args.remote}  ({size} bytes)", flush=True)

    ser = serial.Serial()
    ser.port = port
    ser.baudrate = args.baud
    ser.timeout = 1
    ser.dtr = False
    ser.rts = False
    ser.open()
    # 硬复位 ESP32
    ser.rts = True
    time.sleep(0.1)
    ser.rts = False
    print("waiting BOOT_OK...", flush=True)
    deadline = time.time() + 12
    booted = False
    while time.time() < deadline:
        line = ser.readline().decode("utf-8", errors="ignore").strip()
        if line:
            print("<", line, flush=True)
        if "BOOT_OK" in line:
            booted = True
            break
    if not booted:
        print("warn: BOOT_OK not seen, continuing", flush=True)
    time.sleep(0.5)
    ser.reset_input_buffer()

    cmd = f"#PUT {args.remote} {size}\n".encode("ascii")
    print(f"send {cmd!r}", flush=True)
    ser.write(cmd)
    ser.flush()

    deadline = time.time() + 20
    while time.time() < deadline:
        line = ser.readline().decode("utf-8", errors="ignore").strip()
        if line:
            print("<", line, flush=True)
        if line == "READY":
            break
        if line.startswith("PUT_ERR"):
            return 2
    else:
        print("Timeout waiting READY", file=sys.stderr, flush=True)
        return 3

    sent = 0
    t0 = time.time()
    with path.open("rb") as f:
        while sent < size:
            chunk = f.read(min(WIN, size - sent))
            if not chunk:
                break
            ser.write(chunk)
            ser.flush()
            sent += len(chunk)

            # 等 ACK <sent>
            ack_deadline = time.time() + 30
            acked = False
            while time.time() < ack_deadline:
                line = ser.readline().decode("utf-8", errors="ignore").strip()
                if not line:
                    continue
                print("<", line, flush=True)
                if line.startswith("PUT_ERR"):
                    return 4
                if line.startswith("ACK "):
                    try:
                        got = int(line.split()[1])
                    except (IndexError, ValueError):
                        continue
                    if got >= sent:
                        acked = True
                        break
                if line == "DONE" and sent >= size:
                    print("OK", flush=True)
                    return 0
            if not acked:
                print(f"Timeout waiting ACK after {sent}", file=sys.stderr, flush=True)
                return 4

            if sent % (256 * 1024) < WIN or sent == size:
                rate = sent / max(time.time() - t0, 0.001)
                print(f"sent {sent}/{size} ({rate/1024:.1f} KB/s)", flush=True)

    deadline = time.time() + 30
    while time.time() < deadline:
        line = ser.readline().decode("utf-8", errors="ignore").strip()
        if line:
            print("<", line, flush=True)
        if line == "DONE":
            print("OK", flush=True)
            return 0
        if line.startswith("PUT_ERR"):
            return 5

    print("Timeout waiting DONE", file=sys.stderr, flush=True)
    return 6


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