#!/usr/bin/env python3
"""Automated Qualcomm EDL Sahara/Firehose Flasher for Bananarchy Linux (MSM8909)."""

import argparse
import os
import subprocess
import sys
import tempfile
import time
import usb.core
import usb.util

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(BASE_DIR, "edl"))

from Library.usb import usb_class
from Library.gpt import gpt
from Library.sahara import qualcomm_sahara
from Library.firehose import qualcomm_firehose, xmlparser


DEVICE_LABELS = {
  "nokia-argon": "Nokia 8110 4G (nokia-argon)",
  "nokia-beatles": "Nokia 2720 Flip (nokia-beatles)",
  "nokia-tough": "Nokia 800 Tough (nokia-coconut)",
  "nokia-sparkler": "Nokia 8000 4G (nokia-sparkler)",
  "nokia-leo": "Nokia 6300 4G (nokia-leo)",
}


def _fb_read_status(ep_in, timeout_ms=5000):
  while True:
    resp = bytes(ep_in.read(64, timeout=timeout_ms))
    if resp.startswith(b"INFO"):
      continue
    return resp


def _fb_send_edl(fb_dev):
  cfg = fb_dev.get_active_configuration()
  intf = cfg[(0, 0)]
  ep_out = usb.util.find_descriptor(
    intf,
    custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_OUT,
  )
  ep_in = usb.util.find_descriptor(
    intf,
    custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_IN,
  )
  if ep_out is None or ep_in is None:
    return False
  for cmd in (b"oem reboot-edl", b"oem edl"):
    ep_out.write(cmd, timeout=1000)
    try:
      resp = _fb_read_status(ep_in, timeout_ms=1500)
      if resp.startswith(b"OKAY"):
        break
    except Exception:
      break
  return True


def _fb_upgrade_lk2nd_and_edl(fb_dev):
  lk2nd_path = os.path.join(BASE_DIR, "bootloader", "lk2nd-msm8909.img")
  if not os.path.isfile(lk2nd_path):
    return
  with open(lk2nd_path, "rb") as f:
    lk2nd_data = f.read()

  print("[*] Updating lk2nd on boot partition via Fastboot to enable EDL reboot...")
  cfg = fb_dev.get_active_configuration()
  intf = cfg[(0, 0)]
  ep_out = usb.util.find_descriptor(
    intf,
    custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_OUT,
  )
  ep_in = usb.util.find_descriptor(
    intf,
    custom_match=lambda e: usb.util.endpoint_direction(e.bEndpointAddress) == usb.util.ENDPOINT_IN,
  )
  if ep_out is None or ep_in is None:
    return

  ep_out.write(f"download:{len(lk2nd_data):08x}".encode("ascii"), timeout=2000)
  resp = _fb_read_status(ep_in, timeout_ms=2000)
  if not resp.startswith(b"DATA"):
    return

  chunk_sz = 16384
  for off in range(0, len(lk2nd_data), chunk_sz):
    ep_out.write(lk2nd_data[off:off + chunk_sz], timeout=5000)
  if not _fb_read_status(ep_in, timeout_ms=5000).startswith(b"OKAY"):
    return

  ep_out.write(b"flash:boot", timeout=2000)
  if not _fb_read_status(ep_in, timeout_ms=10000).startswith(b"OKAY"):
    return

  ep_out.write(b"reboot-bootloader", timeout=2000)
  try:
    _fb_read_status(ep_in, timeout_ms=2000)
  except Exception:
    pass
  usb.util.dispose_resources(fb_dev)

  time.sleep(2.0)
  for _ in range(20):
    new_fb = usb.core.find(idVendor=0x18D1, idProduct=0xD00D)
    if new_fb is not None:
      time.sleep(0.3)
      _fb_send_edl(new_fb)
      usb.util.dispose_resources(new_fb)
      return
    time.sleep(0.5)


_fb_edl_kicked = False


def try_kick_into_edl():
  global _fb_edl_kicked
  try:
    res = subprocess.run(
      ["adb", "get-state"],
      stdout=subprocess.PIPE,
      stderr=subprocess.DEVNULL,
      text=True,
      timeout=2,
    )
    if res.returncode == 0 and "device" in res.stdout:
      print("[*] Detected device in ADB mode. Sending 'adb reboot edl'...")
      subprocess.run(["adb", "reboot", "edl"], timeout=5)
      return
  except Exception:
    pass

  # Direct USB Fastboot check (18d1:d00d) via pyusb so host fastboot CLI is optional
  try:
    fb_dev = usb.core.find(idVendor=0x18D1, idProduct=0xD00D)
    if fb_dev is not None:
      print("[*] Detected device in Fastboot mode (18d1:d00d). Sending 'oem reboot-edl'...")
      _fb_send_edl(fb_dev)
      usb.util.dispose_resources(fb_dev)
      return
  except Exception:
    pass

  try:
    res = subprocess.run(
      ["fastboot", "devices"],
      stdout=subprocess.PIPE,
      stderr=subprocess.DEVNULL,
      text=True,
      timeout=2,
    )
    if res.returncode == 0 and res.stdout.strip():
      print("[*] Detected device in Fastboot mode. Sending 'fastboot oem reboot-edl'...")
      rc = subprocess.run(["fastboot", "oem", "reboot-edl"], timeout=5)
      if rc.returncode != 0:
        subprocess.run(["fastboot", "oem", "edl"], timeout=5)
  except Exception:
    pass


def wait_for_edl_usb(vid, pid, device_label, timeout_sec=0):
  cdc = usb_class(vid=vid, pid=pid)
  print(
    f"[*] Waiting for {device_label} in EDL mode "
    f"(USB {vid:04x}:{pid:04x} Qualcomm QDLoader 9008)..."
  )
  print("    To enter EDL mode on the phone, use one of:")
  print("      1) From lk2nd Fastboot menu, select 'EDL' (or connect in Fastboot mode)")
  print("      2) Power off, hold D-Pad Up + D-Pad Down, and plug in USB")
  print("      3) Plug in via an EDL deep-flash cable")
  print("      4) Run 'adb reboot edl' or 'fastboot oem reboot-edl'")

  try_kick_into_edl()

  start_ts = time.time()
  dots = 0
  while not cdc.connected:
    try:
      cdc.connected = cdc.connect()
    except usb.core.USBError as exc:
      if getattr(exc, "errno", None) == 13:
        print(
          "\n[!] Error: USB permission denied when opening 05c6:9008. "
          "Please run via sudo (e.g. sudo ./flash-edl.sh)."
        )
        sys.exit(1)
      time.sleep(0.05)
      continue

    if cdc.connected:
      if dots > 0:
        sys.stdout.write("\n")
      print("[+] USB EDL device connected.")
      return cdc

    if timeout_sec > 0 and (time.time() - start_ts) >= timeout_sec:
      if dots > 0:
        sys.stdout.write("\n")
      print(f"[!] Timed out after {timeout_sec} seconds waiting for EDL device.")
      sys.exit(1)

    if dots > 0 and dots % 100 == 0:
      try_kick_into_edl()

    if dots % 20 == 0:
      sys.stdout.write(".")
      sys.stdout.flush()
    dots += 1
    if dots % 800 == 0:
      sys.stdout.write("\n")
    time.sleep(0.05)

  return cdc


def find_partition(partentries, name):
  for entry in partentries:
    if entry.name == name:
      return entry
  return None


def write_partition(fh, lun, part, filepath):
  size = os.stat(filepath).st_size
  sector_size = fh.cfg.SECTOR_SIZE_IN_BYTES
  sectors = size // sector_size
  if (size % sector_size) > 0:
    sectors += 1

  if sectors > part.sectors:
    print(
      f"[!] Error: {filepath} requires {sectors} sectors, "
      f"but partition '{part.name}' only has {part.sectors} sectors."
    )
    return False

  print(
    f"[*] Flashing '{part.name}' (start sector {part.sector}, "
    f"{sectors}/{part.sectors} sectors, {size / (1024 * 1024):.1f} MiB) "
    f"from {filepath}..."
  )
  ok = fh.cmd_write(lun, part.sector, filepath, Display=True)
  if not ok:
    print(f"[!] Failed writing {filepath} to partition '{part.name}'.")
    return False
  print(f"[+] Successfully flashed '{part.name}'.")
  return True


def dump_partition(fh, lun, part, filepath):
  size_mib = (part.sectors * fh.cfg.SECTOR_SIZE_IN_BYTES) / (1024 * 1024)
  print(
    f"[*] Backing up '{part.name}' (sector {part.sector}, "
    f"{part.sectors} sectors, {size_mib:.2f} MiB) -> {filepath}..."
  )
  res = fh.cmd_read(lun, part.sector, part.sectors, filepath, Display=True)
  if res == "" or not os.path.exists(filepath) or os.stat(filepath).st_size == 0:
    print(f"[!] Failed backing up partition '{part.name}'.")
    return False
  print(f"[+] Saved backup: {filepath}")
  return True


def resolve_default_loader():
  for cand in [
    os.path.join(BASE_DIR, "bootloader", "loaders", "8110.mbn"),
    os.path.join(BASE_DIR, "8110.mbn"),
  ]:
    if os.path.isfile(cand):
      return cand
  return os.path.join(BASE_DIR, "bootloader", "loaders", "8110.mbn")


def resolve_default_bootloader():
  for cand in [
    os.path.join(BASE_DIR, "bootloader", "lk2nd-msm8909.img"),
    os.path.join(BASE_DIR, "lk2nd-msm8909.img"),
  ]:
    if os.path.isfile(cand):
      return cand
  return os.path.join(BASE_DIR, "bootloader", "lk2nd-msm8909.img")


def resolve_default_image(device):
  return os.path.join(BASE_DIR, "images", f"bananarchy-{device}.img")


def parse_args():
  parser = argparse.ArgumentParser(
    description="Automated EDL Flasher for Bananarchy Linux (MSM8909 / Sxmo)"
  )
  parser.add_argument(
    "device_pos",
    nargs="?",
    choices=[
      "sxmo",
      "nokia-argon",
      "nokia-beatles",
      "nokia-tough",
      "nokia-sparkler",
      "nokia-leo",
    ],
    default=None,
    help="Target device to flash (default: nokia-argon)",
  )
  parser.add_argument(
    "--device",
    choices=[
      "nokia-argon",
      "nokia-beatles",
      "nokia-tough",
      "nokia-sparkler",
      "nokia-leo",
    ],
    default="nokia-argon",
    help="Target MSM8909 device (default: nokia-argon)",
  )
  parser.add_argument(
    "--image",
    default="",
    help="Explicit path to raw userdata image (overrides --device)",
  )
  parser.add_argument(
    "--bootloader",
    default=resolve_default_bootloader(),
    help="Path to lk2nd boot image",
  )
  parser.add_argument(
    "--loader",
    default=resolve_default_loader(),
    help="Path to Firehose programmer (.mbn)",
  )
  parser.add_argument(
    "--backup-dir",
    default=os.path.join(BASE_DIR, "backup"),
    help="Directory for partition backups (default: ./backup)",
  )
  parser.add_argument(
    "--skip-backup",
    action="store_true",
    help="Skip backing up critical partitions before flashing",
  )
  parser.add_argument(
    "--force-backup",
    action="store_true",
    help="Re-dump backup partitions even if backup files already exist",
  )
  parser.add_argument(
    "--backup-modem",
    action="store_true",
    help="Also back up the 64 MiB modem partition into backup-dir",
  )
  parser.add_argument(
    "--no-wipe-system",
    action="store_true",
    help="Do not zero the stock system partition superblock",
  )
  parser.add_argument(
    "--no-reset",
    action="store_true",
    help="Do not reboot the phone automatically after flashing",
  )
  parser.add_argument(
    "--timeout",
    type=int,
    default=0,
    help="Timeout in seconds while waiting for USB connection (0 = infinite)",
  )
  parser.add_argument(
    "--vid",
    default="0x05C6",
    help="USB Vendor ID (default: 0x05C6)",
  )
  parser.add_argument(
    "--pid",
    default="0x9008",
    help="USB Product ID (default: 0x9008)",
  )
  parser.add_argument(
    "--lun",
    type=int,
    default=0,
    help="Storage LUN number (default: 0)",
  )
  return parser.parse_args()


def main():
  args = parse_args()
  device = args.device_pos if args.device_pos and args.device_pos != "sxmo" else args.device
  device_label = DEVICE_LABELS.get(device, device)
  userdata_img = (
    os.path.abspath(args.image)
    if args.image
    else os.path.abspath(resolve_default_image(device))
  )
  bootloader_img = os.path.abspath(args.bootloader)
  loader_mbn = os.path.abspath(args.loader)
  backup_dir = os.path.abspath(args.backup_dir)

  print("============================================================")
  print(" Bananarchy Linux (MSM8909 / Sxmo) Automated EDL Flasher")
  print("============================================================")
  print(f" Target Device   : {device_label}")
  print(f" Target UI       : Sxmo (240x320 native, Multi-Mode T9)")
  print(f" Firehose Loader : {loader_mbn}")
  print(f" Bootloader (boot): {bootloader_img}")
  print(f" System (userdata): {userdata_img}")
  print(
    f" Partition Backup: "
    f"{'disabled (--skip-backup)' if args.skip_backup else backup_dir}"
  )
  print("============================================================")

  for label, path in [
    ("Firehose programmer", loader_mbn),
    ("lk2nd bootloader image", bootloader_img),
    ("System userdata image", userdata_img),
  ]:
    if not os.path.isfile(path) or os.stat(path).st_size == 0:
      print(f"[!] Error: {label} not found or empty at: {path}")
      sys.exit(1)

  vid = int(args.vid, 16)
  pid = int(args.pid, 16)
  cdc = wait_for_edl_usb(vid, pid, device_label, timeout_sec=args.timeout)

  sahara = qualcomm_sahara(cdc)
  with open(loader_mbn, "rb") as rf:
    sahara.programmer = rf.read()

  mode = sahara.connect()
  print(f"[*] Device protocol mode: {mode}")

  if mode == "Sahara":
    sahara.info()
    print(f"[*] Uploading Firehose programmer ({os.path.basename(loader_mbn)})...")
    if not sahara.upload_firehoseloader():
      print("[!] Error: Failed to upload Firehose programmer via Sahara.")
      sys.exit(1)
    time.sleep(0.3)
    mode = "Firehose"
    print("[+] Firehose programmer uploaded and active.")
  elif mode == "Firehose":
    if len(sahara.programmer) > 4 and sahara.programmer[:4] == b"\x7fELF":
      sahara.bit64 = sahara.programmer[4] == 2
    else:
      sahara.bit64 = False
    print("[+] Device is already in Firehose mode.")
  else:
    print(
      f"[!] Error: Could not establish Sahara or Firehose mode (got '{mode}'). "
      "Please power-cycle the phone back into EDL mode and retry."
    )
    sys.exit(1)

  xml = xmlparser()
  cfg = qualcomm_firehose.cfg()
  cfg.MemoryName = "eMMC"
  cfg.ZLPAwareHost = 1
  cfg.SkipStorageInit = False
  cfg.SkipWrite = False
  cfg.MaxPayloadSizeToTargetInBytes = 1048576
  cfg.SECTOR_SIZE_IN_BYTES = 512
  cfg.bit64 = sahara.bit64

  fh = qualcomm_firehose(cdc, xml, cfg)
  print("[*] Connecting to Firehose storage controller...")
  fh.connect(0)

  print("[*] Reading GPT partition table from eMMC...")
  gpt_sectors = 0x4000 // fh.cfg.SECTOR_SIZE_IN_BYTES
  gpt_raw = fh.cmd_read_buffer(args.lun, 0, gpt_sectors, Display=False)
  if not gpt_raw:
    print("[!] Error: Failed to read GPT partition table from device.")
    sys.exit(1)

  guid_gpt = gpt()
  if not guid_gpt.parse(gpt_raw, fh.cfg.SECTOR_SIZE_IN_BYTES):
    print("[!] Error: Failed to parse GPT partition table.")
    sys.exit(1)

  boot_part = find_partition(guid_gpt.partentries, "boot")
  userdata_part = find_partition(guid_gpt.partentries, "userdata")
  system_part = find_partition(guid_gpt.partentries, "system")

  if not boot_part:
    print("[!] Error: Partition 'boot' not found in device GPT.")
    sys.exit(1)
  if not userdata_part:
    print("[!] Error: Partition 'userdata' not found in device GPT.")
    sys.exit(1)

  print(
    f"[+] Found 'boot'     : sector {boot_part.sector}, "
    f"{boot_part.sectors} sectors "
    f"({(boot_part.sectors * fh.cfg.SECTOR_SIZE_IN_BYTES) / (1024 * 1024):.1f} MiB)"
  )
  print(
    f"[+] Found 'userdata' : sector {userdata_part.sector}, "
    f"{userdata_part.sectors} sectors "
    f"({(userdata_part.sectors * fh.cfg.SECTOR_SIZE_IN_BYTES) / (1024 * 1024):.1f} MiB)"
  )

  if not args.skip_backup:
    os.makedirs(backup_dir, exist_ok=True)
    gpt_backup_path = os.path.join(backup_dir, "gpt.bin")
    if args.force_backup or not os.path.exists(gpt_backup_path):
      with open(gpt_backup_path, "wb") as wf:
        wf.write(gpt_raw)
      print(f"[+] Saved GPT backup: {gpt_backup_path}")

    backup_parts = ["fsc", "fsg", "modemst1", "modemst2", "persist", "boot"]
    if args.backup_modem:
      backup_parts.append("modem")

    for pname in backup_parts:
      part = find_partition(guid_gpt.partentries, pname)
      if not part:
        print(f"[!] Warning: Partition '{pname}' not found in GPT; skipping backup.")
        continue
      dest_file = os.path.join(backup_dir, f"{pname}.img")
      if (
        not args.force_backup
        and os.path.exists(dest_file)
        and os.stat(dest_file).st_size > 0
      ):
        print(f"[*] Backup already exists for '{pname}' ({dest_file}); skipping.")
        continue
      if not dump_partition(fh, args.lun, part, dest_file):
        sys.exit(1)

  print("\n[*] Step 1/3: Flashing lk2nd secondary bootloader to 'boot'...")
  if not write_partition(fh, args.lun, boot_part, bootloader_img):
    sys.exit(1)

  print(f"\n[*] Step 2/3: Flashing Bananarchy Linux ({device}) image to 'userdata'...")
  if not write_partition(fh, args.lun, userdata_part, userdata_img):
    sys.exit(1)

  if not args.no_wipe_system and system_part:
    print(
      "\n[*] Step 3/3: Clearing stock 'system' superblock "
      "(first 64 sectors / 32 KiB)..."
    )
    with tempfile.NamedTemporaryFile(prefix="wipe-system-", delete=False) as tf:
      tf.write(b"\x00" * (64 * fh.cfg.SECTOR_SIZE_IN_BYTES))
      zero_path = tf.name
    try:
      if not fh.cmd_write(args.lun, system_part.sector, zero_path, Display=False):
        print("[!] Warning: Could not clear 'system' superblock; continuing.")
      else:
        print("[+] Cleared stock 'system' superblock.")
    finally:
      if os.path.exists(zero_path):
        os.unlink(zero_path)
  else:
    print("\n[*] Step 3/3: Skipping 'system' superblock wipe.")

  if not args.no_reset:
    print(f"\n[*] Resetting and booting {device_label}...")
    fh.cmd_reset()

  print("\n============================================================")
  print(" BANANARCHY EDL FLASHING COMPLETED SUCCESSFULLY")
  print("============================================================")
  print(" The phone will now boot via lk2nd into Bananarchy Linux (Sxmo).")
  print(" Default login: user / 147147  |  USB SSH: ssh user@172.16.42.1")
  print("============================================================")


if __name__ == "__main__":
  main()