"""
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
"""
import argparse
import ctypes
import errno
import fcntl
import glob
import hashlib
import os
import re
import signal
import struct
import sys
import time
import zlib
from pathlib import Path

VID_PIDS = (("add1", "0001"), ("3801", "0001"))
ROM_VID_PIDS = (("174c", "2464"), ("174c", "2463"))
ROM_PRODUCT = "USB 3.2 PCIe TinyEnclosure"
FIRMWARE_PATH = Path(__file__).with_name("firmware_wrapped.bin")
CONFIG_DIR = "/data/egpu_dock_config"
LEGACY_CONFIG_DIR = "/data/chestnut_config"
PM_PATHS = ("/sys/bus/platform/devices/a800000.ssusb", "/sys/bus/platform/devices/a600000.ssusb",
            "/sys/bus/usb/devices/usb4")
VBUS_PATH = "/sys/kernel/debug/regulator/smb2-vbus/enable"
IMAGE_OFFSET = 0x100
SECTOR, PAGE = 4096, 128
MAX_REGISTER_READ_SIZE = 255
MAX_CODE_SIZE = 0x10000
FLASH_BUDGET = 600.0
USBDEVFS_CONTROL = 0xC0185500
USBDEVFS_BULK = 0xC0185502
USBDEVFS_SETINTERFACE = 0x80085504
USBDEVFS_SETCONFIGURATION = 0x80045505
USBDEVFS_CLAIMINTERFACE = 0x8004550F
USBDEVFS_RESET = 0x5514
USBDEVFS_CLEAR_HALT = 0x80045515

_deadline = float("inf")


def check_budget():
  if time.monotonic() > _deadline:
    raise TimeoutError(f"flash did not converge within {FLASH_BUDGET:g}s")


class Ctrl(ctypes.Structure):
  _fields_ = [("request_type", ctypes.c_uint8), ("request", ctypes.c_uint8),
              ("value", ctypes.c_uint16), ("index", ctypes.c_uint16),
              ("length", ctypes.c_uint16), ("timeout", ctypes.c_uint32),
              ("data", ctypes.c_void_p)]


class Bulk(ctypes.Structure):
  _fields_ = [("ep", ctypes.c_uint), ("len", ctypes.c_uint),
              ("timeout", ctypes.c_uint), ("data", ctypes.c_void_p)]


class RomFallback(Exception):
  pass


def find_dock():
  found = []
  for d in glob.glob("/sys/bus/usb/devices/*"):
    try:
      with open(d + "/idVendor") as fv, open(d + "/idProduct") as fp:
        vid_pid = (fv.read().strip(), fp.read().strip())
      if vid_pid in VID_PIDS + ROM_VID_PIDS:
        with open(d + "/product") as fpr:
          found.append((d, vid_pid, fpr.read().strip()))
    except OSError:
      pass
  if len(found) > 1:
    raise RuntimeError(f"expected one eGPU dock, found {len(found)}")
  return found[0] if found else (None, None, None)


def in_rom_bootloader(vid_pid, product):
  return vid_pid in ROM_VID_PIDS or product == ROM_PRODUCT or (product or "").startswith("AS2462")


def disable_runtime_pm(path):
  control = os.path.join(path, "power/control")
  if not os.path.exists(control):
    return
  with open(control, "w") as f:
    f.write("on\n")
  with open(control) as fh:
    applied = fh.read().strip()
  if applied != "on":
    raise RuntimeError(f"could not disable USB runtime PM: {control}")
  delay = os.path.join(path, "power/autosuspend_delay_ms")
  if os.path.exists(delay):
    with open(delay, "w") as f:
      f.write("-1\n")


def unbind_drivers(path):
  for interface in glob.glob(path + ":*"):
    driver = interface + "/driver"
    if os.path.islink(driver):
      with open(os.path.realpath(driver) + "/unbind", "w") as f:
        f.write(os.path.basename(interface))


def open_device(path):
  with open(path + "/busnum") as fb, open(path + "/devnum") as fd_:
    bus, dev = int(fb.read()), int(fd_.read())
  return os.open(f"/dev/bus/usb/{bus:03d}/{dev:03d}", os.O_RDWR)


def link_up() -> bool:
  try:
    path, _, _ = find_dock()
    if path is None:
      return False
    fd = open_device(path)
  except (OSError, RuntimeError):
    return False
  try:
    fcntl.ioctl(fd, USBDEVFS_CONTROL, Ctrl(0x40, 0xF3, 1, 0, 0, 2000, None))
    buf = (ctypes.c_ubyte * 1)()
    fcntl.ioctl(fd, USBDEVFS_CONTROL, Ctrl(0xC0, 0xE4, 0xB450, 0, 1, 1000, ctypes.cast(buf, ctypes.c_void_p)))
    return buf[0] == 0x78
  except OSError:
    return False
  finally:
    os.close(fd)


def claim_interface(path, setup=False):
  disable_runtime_pm(path)
  unbind_drivers(path)
  fd = open_device(path)
  try:
    if setup:
      fcntl.ioctl(fd, USBDEVFS_SETCONFIGURATION, struct.pack("I", 1))
    fcntl.ioctl(fd, USBDEVFS_CLAIMINTERFACE, struct.pack("I", 0))
    if setup:
      fcntl.ioctl(fd, USBDEVFS_SETINTERFACE, struct.pack("II", 0, 0))
  except OSError as e:
    os.close(fd)
    if e.errno == errno.EBUSY:
      raise RuntimeError("eGPU dock is in use, stop the model/GPU processes before flashing") from e
    raise
  return fd


class Flash:
  def __init__(self):
    self.fd = -1
    self.max_register_read_size = MAX_REGISTER_READ_SIZE

  def close(self):
    if self.fd >= 0:
      os.close(self.fd)
      self.fd = -1

  def connect(self, timeout=5.0):
    self.close()
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
      path, vid_pid, product = find_dock()
      if in_rom_bootloader(vid_pid, product):
        raise RomFallback("eGPU dock fell back to the ROM bootloader")
      if path is not None:
        try:
          with open(path + "/speed") as fs:
            speed = int(fs.read())
        except (OSError, ValueError):
          speed = 0
        self.max_register_read_size = 64 if speed < 5000 else MAX_REGISTER_READ_SIZE
        self.fd = claim_interface(path)
        return
      time.sleep(0.1)
    raise RuntimeError(f"eGPU dock did not enumerate within {timeout:g}s")

  def reg_write(self, addr, value):
    fcntl.ioctl(self.fd, USBDEVFS_CONTROL,
                Ctrl(0x40, 0xE5, addr & 0xFFFF, value & 0xFFFF, 0, 2000, None))

  def reg_read(self, addr, length=1):
    buf = (ctypes.c_ubyte * length)()
    fcntl.ioctl(self.fd, USBDEVFS_CONTROL,
                Ctrl(0xC0, 0xE4, addr & 0xFFFF, 0, length, 2000, ctypes.cast(buf, ctypes.c_void_p)))
    return bytes(buf)

  def write_buffer(self, data):
    for i, value in enumerate(data):
      self.reg_write(0x7000 + i, value)

  def wait_controller(self, timeout=2.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
      if not self.reg_read(0xC8A9)[0] & 1:
        return
    raise TimeoutError("flash controller timeout")

  def transaction(self, command, addr=0, length=0, addr_len=0x07, mode=0):
    for reg, value in ((0xC8AD, mode), (0xC8AE, 0), (0xC8AF, 0), (0xC8AA, command), (0xC8AC, addr_len),
                       (0xC8A1, addr), (0xC8A2, addr >> 8), (0xC8AB, addr >> 16), (0xC8A3, length >> 8), (0xC8A4, length)):
      self.reg_write(reg, value & 0xFF)
    self.reg_write(0xC8A9, 1)
    self.wait_controller()
    for _ in range(4):
      self.reg_write(0xC8AD, 0)

  def write_enable(self):
    for reg, value in ((0xC8AD, 0), (0xC8AA, 0x06), (0xC8AC, 0x04), (0xC8A3, 0), (0xC8A4, 0), (0xC8A9, 1)):
      self.reg_write(reg, value)
    self.wait_controller()

  def status(self):
    self.transaction(0x05, length=1, addr_len=0x04)
    return self.reg_read(0x7000)[0]

  def wait_write_done(self, timeout=10.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
      if not self.status() & 1:
        return
      time.sleep(0.005)
    raise TimeoutError("SPI flash WIP timeout")

  def init(self):
    self.reg_write(0xCC33, 0x04)
    self.reg_write(0xCA81, self.reg_read(0xCA81)[0] | 1)
    self.reg_write(0xC805, 0x02)
    self.reg_write(0xC8A6, 0x04)
    for _ in range(5):
      self.write_enable()
      self.write_buffer(bytes(4))
      self.transaction(0x01, length=1, addr_len=0x04, mode=1)
      time.sleep(0.01)
      if not self.status() & 0x1C:
        return
    raise RuntimeError("could not clear SPI block protection")

  def read(self, addr, length):
    out = bytearray()
    while len(out) < length:
      n = min(4096, length - len(out))
      self.transaction(0x03, addr + len(out), max(4096, n))
      for off in range(0, n, self.max_register_read_size):
        out += self.reg_read(0x7000 + off, min(self.max_register_read_size, n - off))
    return bytes(out)

  def erase_sector(self, addr):
    self.write_enable()
    self.transaction(0x20, addr)
    self.wait_write_done()

  def program(self, addr, data):
    self.write_buffer(data + bytes((-len(data)) % 4))
    self.write_enable()
    self.transaction(0x02, addr, len(data), mode=1)
    self.wait_write_done()


def validate_image(data):
  if len(data) < 10:
    raise ValueError("wrapped firmware is too short")
  body_len = int.from_bytes(data[:4], "little")
  if body_len > MAX_CODE_SIZE:
    raise ValueError(f"wrapped firmware body exceeds {MAX_CODE_SIZE} bytes")
  if len(data) != body_len + 10 or data[4 + body_len] != 0xA5:
    raise ValueError("invalid wrapped firmware length or magic")
  body = data[4:4 + body_len]
  if data[5 + body_len] != sum(body) & 0xFF:
    raise ValueError("invalid wrapped firmware checksum")
  if data[6 + body_len:] != zlib.crc32(body).to_bytes(4, "little"):
    raise ValueError("invalid wrapped firmware CRC")


def image_product(image):
  match = re.search(rb"custom [0-9a-f]{8}-CLEAN", image)
  if match is None:
    raise ValueError("no product string in wrapped firmware")
  return match.group().decode()


def reconnect(flash):
  attempt = 0
  while True:
    attempt += 1
    check_budget()
    try:
      flash.connect()
      flash.init()
      return
    except (OSError, TimeoutError, RuntimeError) as e:
      print(f"waiting for eGPU dock (attempt {attempt}): {e}", flush=True)
      time.sleep(1)


def with_retries(flash, label, operation):
  attempt = 0
  while True:
    attempt += 1
    try:
      return operation()
    except (OSError, TimeoutError, RuntimeError) as e:
      check_budget()
      print(f"{label} attempt {attempt}: {e}", flush=True)
      reconnect(flash)


def stable_read(flash, addr, length, count=2):
  def read():
    reads = [flash.read(addr, length) for _ in range(count)]
    if any(x != reads[0] for x in reads[1:]):
      raise RuntimeError(f"unstable flash read at 0x{addr:05x}")
    return reads[0]
  return with_retries(flash, f"read 0x{addr:05x}", read)


def program_sector(flash, addr, target):
  def program():
    flash.erase_sector(addr)
    if flash.read(addr, SECTOR) != bytes([0xFF]) * SECTOR:
      raise RuntimeError("sector erase verification failed")
    for off in range(0, SECTOR, PAGE):
      chunk = target[off:off + PAGE]
      if chunk != bytes([0xFF]) * len(chunk):
        flash.program(addr + off, chunk)
        if flash.read(addr + off, len(chunk)) != chunk:
          raise RuntimeError(f"page verify failed at 0x{addr + off:05x}")
    if flash.read(addr, SECTOR) != target:
      raise RuntimeError("sector verification failed")
  with_retries(flash, f"sector 0x{addr:05x}", program)


def config_path():
  return os.path.join(CONFIG_DIR, f"{os.uname().nodename}.bin")


def legacy_config_path():
  return os.path.join(LEGACY_CONFIG_DIR, f"{os.uname().nodename}.bin")


def saved_config(path, data):
  os.makedirs(os.path.dirname(path), exist_ok=True)
  try:
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
  except FileExistsError as e:
    with open(path, "rb") as fh:
      backup = fh.read()
    if len(backup) != 0x100:
      raise RuntimeError(f"invalid config backup: {path}") from e
    if backup != data:
      print(f"restoring config from {path}", flush=True)
    return backup
  with os.fdopen(fd, "wb") as f:
    f.write(data)
    f.flush()
    os.fsync(f.fileno())
  return data


def rom_write(image, config):
  path, _, _ = find_dock()
  if path is None:
    raise RuntimeError("eGPU dock disappeared before recovery")
  unbind_drivers(path)
  fd = open_device(path)
  try:
    fcntl.ioctl(fd, USBDEVFS_RESET)
  finally:
    os.close(fd)
  time.sleep(3)
  path, _, _ = find_dock()
  if path is None:
    raise RuntimeError("eGPU dock did not re-enumerate after reset")
  fd = claim_interface(path, setup=True)
  for ep in (0x02, 0x81):
    fcntl.ioctl(fd, USBDEVFS_CLEAR_HALT, struct.pack("I", ep))
  tag = 0

  def bulk(ep, payload, timeout):
    buf = ctypes.create_string_buffer(bytes(payload), len(payload))
    fcntl.ioctl(fd, USBDEVFS_BULK, Bulk(ep, len(payload), timeout, ctypes.cast(buf, ctypes.c_void_p)))
    return buf.raw

  def cmd(cdb, data=b"", timeout=30000):
    nonlocal tag
    tag += 1
    bulk(0x02, struct.pack("<IIIBBB16s", 0x43425355, tag, len(data), 0, 0, len(cdb), cdb), timeout)
    if data:
      bulk(0x02, data, timeout)
    try:
      csw = bulk(0x81, bytes(13), timeout)
    except OSError as e:
      if e.errno != errno.EPIPE:
        raise
      fcntl.ioctl(fd, USBDEVFS_CLEAR_HALT, struct.pack("I", 0x81))
      csw = bulk(0x81, bytes(13), timeout)
    if csw[:4] != b"USBS" or csw[12] != 0:
      raise RuntimeError(f"ROM flash command {cdb[0]:02x} {cdb[1]:02x} failed")

  print("recovering from the ROM bootloader", flush=True)
  try:
    cmd(struct.pack(">BBB12x", 0xE1, 0x50, 0), config[:0x80])
    cmd(struct.pack(">BBB12x", 0xE1, 0x50, 1), config[0x80:])
    cmd(struct.pack(">BBI", 0xE3, 0x50, min(len(image), 0xFF00)), image[:0xFF00])
    if len(image) > 0xFF00:
      cmd(struct.pack(">BBI", 0xE3, 0xD0, len(image) - 0xFF00), image[0xFF00:])
    cmd(struct.pack(">BB13x", 0xE8, 0x51))
  finally:
    os.close(fd)
  print("recovery flash done", flush=True)


def vbus_write(value):
  try:
    with open(VBUS_PATH, "w") as f:
      f.write(value + "\n")
  except OSError:
    pass


def vbus_cycle():
  if os.path.exists(VBUS_PATH):
    vbus_write("0")
    time.sleep(2)
    vbus_write("1")
    time.sleep(5)


def activate(expected_product):
  if not os.path.exists(VBUS_PATH):
    print("no VBUS control, firmware activates on the next dock power cycle", flush=True)
    return
  print("power-cycling the eGPU dock VBUS", flush=True)
  vbus_write("0")
  disconnected = False
  deadline = time.monotonic() + 5.0
  while time.monotonic() < deadline:
    path, _, _ = find_dock()
    if path is None:
      disconnected = True
      break
    time.sleep(0.2)
  time.sleep(1)
  vbus_write("1")
  if not disconnected:
    print("dock stayed powered, firmware activates on its next power cycle", flush=True)
    return
  deadline = time.monotonic() + 15.0
  while time.monotonic() < deadline:
    _, _, product = find_dock()
    if product is not None:
      if product == expected_product:
        print(f"activated {expected_product}", flush=True)
      else:
        print(f"dock re-enumerated with {product!r}, firmware activates on its next power cycle", flush=True)
      return
    time.sleep(0.2)
  print("dock did not re-enumerate, firmware activates on its next power cycle", flush=True)


def defer_signal(signum, _frame):
  os.write(1, f"signal {signum} deferred until the dock is powered back up\n".encode())


def flash_dock(expected_version=None, force=False):
  global _deadline

  image = FIRMWARE_PATH.read_bytes()
  validate_image(image)
  expected_product = image_product(image)
  if expected_version is not None and expected_product != f"custom {expected_version}-CLEAN":
    raise RuntimeError(f"bundled firmware is {expected_product!r}, expected version {expected_version}")

  path, vid_pid, product = find_dock()
  if path is None:
    print("no eGPU dock connected", flush=True)
    return
  if product == expected_product and not force:
    print(f"eGPU dock firmware is up to date ({expected_product})", flush=True)
    return

  _deadline = time.monotonic() + FLASH_BUDGET
  for pm_path in PM_PATHS:
    disable_runtime_pm(pm_path)

  previous = {sig: signal.signal(sig, defer_signal) for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP)}
  try:
    if in_rom_bootloader(vid_pid, product):
      if not recover_from_rom(image, expected_product):
        return
      force, product = True, None
    write_image(image, expected_product, product, force)
  finally:
    for sig, handler in previous.items():
      signal.signal(sig, handler)


def recover_from_rom(image, expected_product):
  backup = config_path()
  if not os.path.isfile(backup) and os.path.isfile(legacy_config_path()):
    backup = legacy_config_path()
  if not os.path.isfile(backup):
    raise RuntimeError(f"cannot recover from the ROM bootloader without a config backup at {config_path()}")
  with open(backup, "rb") as fh:
    config = fh.read()
  if len(config) != 0x100:
    raise RuntimeError(f"invalid config backup: {backup}")

  committed = False
  while True:
    check_budget()
    path, vid_pid, product = find_dock()
    if path is None:
      if committed:
        print("dock is offline, recovered firmware boots on its next power cycle", flush=True)
        return False
      vbus_cycle()
      continue
    if not in_rom_bootloader(vid_pid, product):
      return True
    if committed:
      print("dock stayed powered, recovered firmware boots on its next power cycle", flush=True)
      return False
    try:
      rom_write(image, config)
      committed = True
    except (OSError, TimeoutError, RuntimeError) as e:
      print(f"ROM recovery failed, retrying: {e}", flush=True)
      vbus_cycle()
      continue
    activate(expected_product)


def write_image(image, expected_product, product, force):
  if force:
    print(f"forced reflash of {expected_product}", flush=True)
  else:
    print(f"eGPU dock firmware mismatch: {product!r}; expected {expected_product!r}", flush=True)

  flash = Flash()
  try:
    reconnect(flash)
    config = stable_read(flash, 0, 0x100, 3)
    config = saved_config(config_path(), config)
    image_end = IMAGE_OFFSET + len(image)
    first_sector = IMAGE_OFFSET & ~(SECTOR - 1)
    span = (image_end + SECTOR - 1) & ~(SECTOR - 1)
    current = stable_read(flash, first_sector, span - first_sector)
    target = bytearray(current)
    target[:len(config)] = config
    target[IMAGE_OFFSET - first_sector:image_end - first_sector] = image
    target = bytes(target)
    print(f"target {len(image)} bytes at 0x{IMAGE_OFFSET:05x}, sha256={hashlib.sha256(image).hexdigest()}", flush=True)

    if not _still_offroad():
      raise RuntimeError("device went onroad before any sector was written; aborting flash")

    for addr in range(first_sector, span, SECTOR):
      off = addr - first_sector
      wanted = target[off:off + SECTOR]
      if current[off:off + SECTOR] == wanted:
        print(f"sector 0x{addr:05x}: unchanged", flush=True)
      else:
        print(f"sector 0x{addr:05x}: programming", flush=True)
        program_sector(flash, addr, wanted)

    verified = stable_read(flash, first_sector, span - first_sector, 3)
    if verified != target:
      raise RuntimeError("final full-image verification failed")
    print(f"verified sha256={hashlib.sha256(verified).hexdigest()}", flush=True)
  finally:
    flash.close()

  activate(expected_product)


def bundled_version() -> str:
  return image_product(FIRMWARE_PATH.read_bytes())


def _still_offroad() -> bool:
  try:
    from iqpilot.common.params import Params
    return bool(Params().get_bool("IsOffroad"))
  except Exception:
    return True


def dock_needs_flash(usb_devices: list[dict]) -> bool:
  try:
    expected = bundled_version()
  except (OSError, ValueError):
    return False
  ids = tuple(tuple(int(x, 16) for x in p) for p in VID_PIDS + ROM_VID_PIDS)
  return any((d.get("vendorId"), d.get("productId")) in ids and d.get("product") != expected
             for d in usb_devices)


def main():
  parser = argparse.ArgumentParser(description="check and flash the bundled eGPU dock firmware")
  parser.add_argument("version", nargs="?", help="expected firmware version hash")
  parser.add_argument("--force", action="store_true", help="reflash even when the version matches")
  args = parser.parse_args()
  if os.geteuid() != 0:
    raise RuntimeError("flash.py must run as root")
  flash_dock(expected_version=args.version, force=args.force)


if __name__ == "__main__":
  try:
    main()
  except Exception as e:
    print(f"FAIL: {type(e).__name__}: {e}", file=sys.stderr)
    sys.exit(1)