"""Check an existing U-Boot prompt or configure its volatile environment."""
import argparse
import fcntl
import re
import sys
import time
from pathlib import Path

import serial

# Match the complete final line, never a prompt embedded in command output.
PROMPT = re.compile(r"(?:^|\n)([^\r\n]*=>) ?$")
ERROR = re.compile(
    r"Unknown command|## Error|^Error:|^Usage:|^.*(?:not defined|not found|failed|failure).*$",
    re.IGNORECASE | re.MULTILINE
)


class UBoot:
    def __init__(self, port):
        self.port = port
        self.prompt = None

    def read_prompt(self, timeout):
        output = ""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            data = self.port.read(self.port.in_waiting or 1)
            if data:
                chunk = data.decode("utf-8", errors="replace").replace("\r", "")
                print(chunk, end="", flush=True)
                output += chunk
            match = PROMPT.search(output)
            if match and (self.prompt is None or match[1] == self.prompt):
                # Wait for a quiet prompt rather than a partial output line.
                time.sleep(0.15)
                if self.port.in_waiting:
                    continue
                return output, match[1]
        raise RuntimeError("Timed out waiting for the CLI prompt")

    def detect(self):
        # Observe before requesting a fresh prompt.
        try:
            _, candidate = self.read_prompt(3)
        except RuntimeError:
            candidate = None
        if candidate is None:
            # An idle prompt may have been printed before the port was opened.
            # A newline only requests a fresh prompt; no U-Boot command is sent.
            self.port.write(b"\r")
            self.port.flush()
            _, candidate = self.read_prompt(5)
        if not candidate.endswith("=>"):
            raise RuntimeError("Response is not a recognized U-Boot CLI prompt")
        self.prompt = candidate
        print("\nU-Boot-style CLI prompt detected: " + candidate, flush=True)

    def command(self, command, expected=None, timeout=10, allow_missing=()):
        print("\n>>> " + command, flush=True)
        self.port.write((command + "\r").encode())
        self.port.flush()
        output, _ = self.read_prompt(timeout)
        # Query immediately, before another command changes the return status.
        self.port.write(b"echo __UBOOT_STATUS_ $?\r")
        self.port.flush()
        status, _ = self.read_prompt(5)
        codes = re.findall(r"(?m)^__UBOOT_STATUS_ (\d+) *$", status)
        # Only the initial environment inspection may report unset variables.
        missing_pattern = r'^## Error: "(' + "|".join(map(re.escape, allow_missing)) + r')" not defined$'
        missing = re.findall(missing_pattern, output, re.MULTILINE) if allow_missing else []
        checked_output = re.sub(missing_pattern, "", output, flags=re.MULTILINE) if missing else output
        if codes != ["0"] and not (codes == ["1"] and missing):
            raise RuntimeError("Command failed or its exit status could not be verified: " + command)
        if ERROR.search(checked_output):
            raise RuntimeError("Error in command response: " + command)
        if expected and not re.search(expected, output, re.MULTILINE):
            raise RuntimeError("Expected response missing: " + command)
        if missing:
            print("\nInitially unset: " + ", ".join(missing), flush=True)
        print("\nVerified: " + command, flush=True)
        return output


def transfer(cli, firmware):
    image = firmware.read_bytes()
    if not 32 <= len(image) <= 0x02000000:
        raise RuntimeError("Firmware must contain at least 32 bytes and fit RAM at 0x42000000")
    output = cli.command(
        "tftp 0x42000000 u-boot-sunxi-with-spl.bin",
        r"^Bytes transferred = \d+", timeout=120,
    )
    transferred = re.findall(r"^Bytes transferred = (\d+)", output, re.MULTILINE)
    if transferred != [str(len(image))]:
        raise RuntimeError("TFTP byte count differs from the local firmware file")
    output = cli.command("echo ${filesize}", r"^[0-9a-fA-F]+$")
    sizes = re.findall(r"^([0-9a-fA-F]+)$", output, re.MULTILINE)
    if len(sizes) != 1 or int(sizes[0], 16) != len(image):
        raise RuntimeError("U-Boot filesize differs from the local firmware file")
    output = cli.command("md.b 0x42000000 0x20", r"^42000000:")
    dumped = bytearray()
    for address in ("42000000", "42000010"):
        rows = re.findall(r"^" + address + r":((?: [0-9a-fA-F]{2}){16})(?:\s|$)", output, re.MULTILINE)
        if len(rows) != 1:
            raise RuntimeError("Missing or invalid memory dump at " + address)
        dumped.extend(bytes.fromhex(rows[0]))
    if dumped != image[:32]:
        raise RuntimeError("First 32 bytes in RAM differ from the local firmware file")
    print("\nTransfer verified: size and first 32 bytes match; no flash write performed.")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=("check", "configure", "transfer"), default="check", nargs="?")
    parser.add_argument("--port", default="/dev/ttyUSB0")
    parser.add_argument("--firmware", type=Path, default=Path(__file__).resolve().parent.parent / "firmware" / "u-boot-sunxi-with-spl.bin")
    args = parser.parse_args()
    if args.mode == "transfer" and not args.firmware.is_file():
        raise RuntimeError("Firmware file not found: " + str(args.firmware))
    # Shared by both workflows: fail rather than interleave serial commands.
    with open("/tmp/uboot-serial.lock", "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        with serial.Serial(args.port, 115200, timeout=0.1, write_timeout=2, exclusive=True) as port:
            cli = UBoot(port)
            cli.detect()
            if args.mode == "check":
                return
            if args.mode == "transfer":
                transfer(cli, args.firmware)
                return
            cli.command("version", r"^U-Boot \d")
            cli.command(
                "printenv ipaddr serverip ethaddr bootcmd bootargs",
                allow_missing=("ipaddr", "serverip", "ethaddr", "bootcmd", "bootargs"),
            )
            cli.command("setenv ipaddr 192.168.75.100")
            cli.command("printenv ipaddr", r"^ipaddr=192\.168\.75\.100$")
            cli.command("setenv serverip 192.168.75.247")
            cli.command("printenv serverip", r"^serverip=192\.168\.75\.247$")
            cli.command("ping 192.168.75.247", r"^host 192\.168\.75\.247 is alive$", timeout=30)
            cli.command("sf probe 0:0", r"^SF: Detected .+", timeout=20)
            print("\nAll commands verified; environment changes were not saved to flash.")


if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, OSError, serial.SerialException) as exc:
        print("\nERROR: " + str(exc), file=sys.stderr)
        sys.exit(1)