// STK500v1 programmer for Arduino Uno/Nano via WebSerial.
// Implements the subset of the protocol needed to upload Intel HEX firmware
// through the Optiboot bootloader (115200 baud) or older bootloaders (57600).

// ─── STK500 constants ────────────────────────────────────

const Cmnd_STK_GET_SYNC = 0x30;
const Cmnd_STK_SET_DEVICE = 0x42;
const Cmnd_STK_ENTER_PROGMODE = 0x50;
const Cmnd_STK_LEAVE_PROGMODE = 0x51;
const Cmnd_STK_LOAD_ADDRESS = 0x55;
const Cmnd_STK_PROG_PAGE = 0x64;
const Cmnd_STK_READ_PAGE = 0x74;
const Cmnd_STK_READ_SIGN = 0x75;
const Sync_CRC_EOP = 0x20;
const Resp_STK_INSYNC = 0x14;
const Resp_STK_OK = 0x10;

// ─── Board profiles ─────────────────────────────────────

export const BOARDS = {
  uno: {
    name: 'Arduino Uno',
    baudRate: 115200,
    signature: new Uint8Array([0x1e, 0x95, 0x0f]),
    pageSize: 128,
    timeout: 400,
  },
  nano: {
    name: 'Arduino Nano',
    baudRate: 57600,
    signature: new Uint8Array([0x1e, 0x95, 0x0f]),
    pageSize: 128,
    timeout: 400,
  },
  nano_new: {
    name: 'Arduino Nano (new bootloader)',
    baudRate: 115200,
    signature: new Uint8Array([0x1e, 0x95, 0x0f]),
    pageSize: 128,
    timeout: 400,
  },
};

// ─── Intel HEX parser ───────────────────────────────────

export function parseHex(hexString) {
  const lines = hexString.split(/\r?\n/).filter(l => l.startsWith(':'));
  let baseAddress = 0;
  let maxAddress = 0;
  const segments = [];

  for (const line of lines) {
    const byteCount = parseInt(line.substring(1, 3), 16);
    const address = parseInt(line.substring(3, 7), 16);
    const recordType = parseInt(line.substring(7, 9), 16);

    if (recordType === 0x00) {
      const data = new Uint8Array(byteCount);
      for (let i = 0; i < byteCount; i++) {
        data[i] = parseInt(line.substring(9 + i * 2, 11 + i * 2), 16);
      }
      const absAddress = baseAddress + address;
      segments.push({ address: absAddress, data });
      const end = absAddress + byteCount;
      if (end > maxAddress) maxAddress = end;
    } else if (recordType === 0x02) {
      baseAddress = parseInt(line.substring(9, 13), 16) << 4;
    } else if (recordType === 0x04) {
      baseAddress = parseInt(line.substring(9, 13), 16) << 16;
    } else if (recordType === 0x01) {
      break;
    }
  }

  const buffer = new Uint8Array(maxAddress);
  buffer.fill(0xff);
  for (const seg of segments) {
    buffer.set(seg.data, seg.address);
  }
  return buffer;
}

// ─── Low-level serial helpers ───────────────────────────

async function serialWrite(writer, data) {
  if (data instanceof Array) data = new Uint8Array(data);
  await writer.write(data);
}

// Buffered reader: WebSerial delivers variable-size chunks, so we need to
// accumulate bytes and serve exact-length reads from the buffer.
function createBufferedReader(rawReader) {
  let buffer = [];
  let pending = null; // { needed, resolve, reject, timer }

  const pump = async () => {
    try {
      while (true) {
        const { value, done } = await rawReader.read();
        if (done) break;
        if (value) {
          for (let i = 0; i < value.length; i++) buffer.push(value[i]);
          checkPending();
        }
      }
    } catch {
      // port closed / cancelled
      if (pending) pending.reject(new Error('Serial read cancelled'));
    }
  };

  function checkPending() {
    if (pending && buffer.length >= pending.needed) {
      clearTimeout(pending.timer);
      const bytes = new Uint8Array(buffer.splice(0, pending.needed));
      const p = pending;
      pending = null;
      p.resolve(bytes);
    }
  }

  pump();

  return {
    read(length, timeout) {
      // If we already have enough buffered, return immediately
      if (buffer.length >= length) {
        return Promise.resolve(new Uint8Array(buffer.splice(0, length)));
      }
      return new Promise((resolve, reject) => {
        const timer = setTimeout(() => {
          pending = null;
          reject(new Error(`Read timeout (expected ${length} bytes, got ${buffer.length})`));
        }, timeout);
        pending = { needed: length, resolve, reject, timer };
      });
    },
    drain() {
      buffer.length = 0;
    },
    get rawReader() { return rawReader; },
  };
}

// ─── STK500 protocol steps ──────────────────────────────

async function sync(writer, reader, timeout) {
  const cmd = new Uint8Array([Cmnd_STK_GET_SYNC, Sync_CRC_EOP]);
  for (let attempt = 0; attempt < 10; attempt++) {
    reader.drain();
    try {
      await serialWrite(writer, cmd);
      const resp = await reader.read(2, timeout);
      if (resp[0] === Resp_STK_INSYNC && resp[1] === Resp_STK_OK) {
        return true;
      }
    } catch {
      // retry
    }
    await new Promise(r => setTimeout(r, 30));
  }
  throw new Error('Failed to sync with bootloader after 10 attempts');
}

async function readSignature(writer, reader, timeout) {
  const cmd = new Uint8Array([Cmnd_STK_READ_SIGN, Sync_CRC_EOP]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(5, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[4] !== Resp_STK_OK) {
    throw new Error('Failed to read device signature');
  }
  return resp.slice(1, 4);
}

async function setDeviceParameters(writer, reader, pageSize, timeout) {
  const cmd = new Uint8Array([
    Cmnd_STK_SET_DEVICE,
    0x86, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x03,
    0xff, 0xff, 0xff, 0xff,
    (pageSize >> 8) & 0xff, pageSize & 0xff,
    0x00, 0x00,
    0x00, 0x00, 0x80, 0x00,
    Sync_CRC_EOP,
  ]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[1] !== Resp_STK_OK) {
    throw new Error('Failed to set device parameters');
  }
}

async function enterProgMode(writer, reader, timeout) {
  const cmd = new Uint8Array([Cmnd_STK_ENTER_PROGMODE, Sync_CRC_EOP]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[1] !== Resp_STK_OK) {
    throw new Error('Failed to enter programming mode');
  }
}

async function leaveProgMode(writer, reader, timeout) {
  const cmd = new Uint8Array([Cmnd_STK_LEAVE_PROGMODE, Sync_CRC_EOP]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[1] !== Resp_STK_OK) {
    throw new Error('Failed to leave programming mode');
  }
}

async function loadAddress(writer, reader, addr, timeout) {
  const cmd = new Uint8Array([
    Cmnd_STK_LOAD_ADDRESS,
    addr & 0xff,
    (addr >> 8) & 0xff,
    Sync_CRC_EOP,
  ]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[1] !== Resp_STK_OK) {
    throw new Error('Failed to load address');
  }
}

async function programPage(writer, reader, data, timeout) {
  const sizeHigh = (data.length >> 8) & 0xff;
  const sizeLow = data.length & 0xff;
  const cmd = new Uint8Array([
    Cmnd_STK_PROG_PAGE,
    sizeHigh, sizeLow,
    0x46, // 'F' for flash memory
    ...data,
    Sync_CRC_EOP,
  ]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[1] !== Resp_STK_OK) {
    throw new Error('Failed to program page');
  }
}

async function readPage(writer, reader, size, timeout) {
  const cmd = new Uint8Array([
    Cmnd_STK_READ_PAGE,
    (size >> 8) & 0xff,
    size & 0xff,
    0x46, // 'F' for flash
    Sync_CRC_EOP,
  ]);
  await serialWrite(writer, cmd);
  const resp = await reader.read(size + 2, timeout);
  if (resp[0] !== Resp_STK_INSYNC || resp[resp.length - 1] !== Resp_STK_OK) {
    throw new Error('Failed to read page');
  }
  return resp.slice(1, 1 + size);
}

// ─── DTR reset ──────────────────────────────────────────
// Arduino boards reset when DTR transitions HIGH → LOW.  The RC circuit
// (100nF cap + 10k pull-up) converts the falling edge into a brief LOW
// pulse on the RESET pin.  We hold DTR HIGH first to guarantee a clean
// falling edge, then wait for the bootloader to start (~100-250ms).

async function resetBoard(port) {
  await port.setSignals({ dataTerminalReady: false, requestToSend: false });
  await new Promise(r => setTimeout(r, 50));
  await port.setSignals({ dataTerminalReady: true, requestToSend: true });
  await new Promise(r => setTimeout(r, 100));
  await port.setSignals({ dataTerminalReady: false, requestToSend: false });
  await new Promise(r => setTimeout(r, 100));
}

// ─── Main upload function ───────────────────────────────

/**
 * Upload an Intel HEX file to an Arduino via WebSerial.
 *
 * @param {string} hexString - The Intel HEX file contents as a string
 * @param {object} [options] - Options
 * @param {string} [options.board='uno'] - Board key from BOARDS
 * @param {function} [options.onProgress] - Callback: (status: string, percent: number) => void
 * @param {SerialPort} [options.port] - An already-opened SerialPort. If not provided, will prompt user.
 * @returns {Promise<void>}
 */
export async function uploadHex(hexString, options = {}) {
  const boardKey = options.board || 'uno';
  const board = typeof boardKey === 'object' ? boardKey : BOARDS[boardKey];
  if (!board) throw new Error(`Unknown board: ${boardKey}`);
  const onProgress = options.onProgress || (() => {});

  const hexData = parseHex(hexString);
  if (hexData.length === 0) throw new Error('Hex file is empty');

  onProgress('Connecting...', 0);

  // Open a raw binary serial port (separate from the REPL text connection)
  let port = options.port || null;
  let weOpened = false;

  if (!port) {
    port = await navigator.serial.requestPort();
  }

  // Close if already open, then open fresh with correct baud rate
  try { await port.close(); } catch { /* not open */ }
  await port.open({ baudRate: board.baudRate });
  weOpened = true;

  let writer, rawReader, reader;
  try {
    writer = port.writable.getWriter();
    rawReader = port.readable.getReader();
    reader = createBufferedReader(rawReader);

    onProgress('Resetting board...', 5);
    await resetBoard(port);

    // Start syncing immediately — don't waste time draining
    reader.drain();

    onProgress('Syncing with bootloader...', 10);
    await sync(writer, reader, board.timeout);

    onProgress('Reading signature...', 20);
    const sig = await readSignature(writer, reader, board.timeout);
    if (board.signature[0] !== sig[0] || board.signature[1] !== sig[1] || board.signature[2] !== sig[2]) {
      const got = Array.from(sig).map(b => '0x' + b.toString(16).padStart(2, '0')).join(' ');
      const expected = Array.from(board.signature).map(b => '0x' + b.toString(16).padStart(2, '0')).join(' ');
      throw new Error(`Signature mismatch: got ${got}, expected ${expected}`);
    }

    onProgress('Setting device parameters...', 25);
    await setDeviceParameters(writer, reader, board.pageSize, board.timeout);

    onProgress('Entering programming mode...', 30);
    await enterProgMode(writer, reader, board.timeout);

    // Program pages
    const totalPages = Math.ceil(hexData.length / board.pageSize);
    for (let page = 0; page < totalPages; page++) {
      const addr = (page * board.pageSize) >> 1; // word address
      const chunk = hexData.subarray(
        page * board.pageSize,
        Math.min((page + 1) * board.pageSize, hexData.length)
      );

      await loadAddress(writer, reader, addr, board.timeout);
      await programPage(writer, reader, chunk, board.timeout);

      const pct = 35 + Math.round((page / totalPages) * 40);
      onProgress(`Writing page ${page + 1}/${totalPages}...`, pct);
    }

    // Verify pages
    onProgress('Verifying...', 75);
    for (let page = 0; page < totalPages; page++) {
      const addr = (page * board.pageSize) >> 1;
      const chunk = hexData.subarray(
        page * board.pageSize,
        Math.min((page + 1) * board.pageSize, hexData.length)
      );

      await loadAddress(writer, reader, addr, board.timeout);
      const readBack = await readPage(writer, reader, chunk.length, board.timeout);

      for (let i = 0; i < chunk.length; i++) {
        if (chunk[i] !== readBack[i]) {
          throw new Error(`Verification failed at byte ${page * board.pageSize + i}`);
        }
      }

      const pct = 75 + Math.round((page / totalPages) * 20);
      onProgress(`Verifying page ${page + 1}/${totalPages}...`, pct);
    }

    onProgress('Finishing...', 97);
    await leaveProgMode(writer, reader, board.timeout);

    onProgress('Upload complete!', 100);
  } finally {
    try { rawReader.releaseLock(); } catch { /* */ }
    try { writer.releaseLock(); } catch { /* */ }
    if (weOpened) {
      try { await port.close(); } catch { /* */ }
    }
  }
}
