import { ESPLoader, Transport } from 'esptool-js';

const DEFAULT_FIRMWARE_URL = import.meta.env.BASE_URL + 'firmware/ESP32_GENERIC-20260406-v1.28.0.bin';

function arrayBufferToBinaryString(buffer) {
  const bytes = new Uint8Array(buffer);
  const chunks = [];
  const chunkSize = 8192;
  for (let i = 0; i < bytes.length; i += chunkSize) {
    const slice = bytes.subarray(i, i + chunkSize);
    chunks.push(String.fromCharCode.apply(null, slice));
  }
  return chunks.join('');
}

export async function flashFirmware(onLog, onProgress, options = {}) {
  const firmwareUrl = options.firmwareUrl || DEFAULT_FIRMWARE_URL;
  const targetLabel = options.targetLabel || 'ESP32';
  const flashAddress = options.flashAddress ?? 0x0;
  const resetMode = options.resetMode || 'default_reset';
  const port = await navigator.serial.requestPort();

  // Buffer to track message state for proper newline insertion
  let messageBuffer = '';

  function normalizeAndLog(data) {
    messageBuffer += data;
    
    // Detect patterns that indicate message boundaries and insert newlines
    // Pattern 1: "Writing at" should always start on a new line (unless already at start/newline)
    messageBuffer = messageBuffer.replace(/([^\n])(Writing at)/g, '$1\n$2');
    
    // Pattern 2: Messages ending with "%)" followed by "Writing at" need a newline
    messageBuffer = messageBuffer.replace(/(%\))(Writing at)/g, '$1\n$2');
    
    // Pattern 3: Messages ending with "." or "..." followed by capital letters (like "Leaving...")
    messageBuffer = messageBuffer.replace(/(\.\.?)([A-Z])/g, '$1\n$2');
    
    // Pattern 4: "Wrote" messages should be on their own line
    messageBuffer = messageBuffer.replace(/([^\n])(Wrote \d+ bytes)/g, '$1\n$2');
    
    // Split on newlines and process complete lines
    const lines = messageBuffer.split('\n');
    // Keep the last potentially incomplete line in buffer
    messageBuffer = lines.pop() || '';
    
    // Log all complete lines
    for (const line of lines) {
      if (line.trim()) {
        onLog?.(line + '\n');
      }
    }
  }
  
  function flushBuffer() {
    if (messageBuffer.trim()) {
      onLog?.(messageBuffer + '\n');
      messageBuffer = '';
    }
  }

  // Some USB-native ESP32 ports (e.g. S2 CDC) don't support DTR/RTS control
  // signals. Patch setSignals so failures are non-fatal — the reset sequence
  // becomes a no-op and the chip must already be in bootloader mode.
  const origSetSignals = port.setSignals.bind(port);
  let signalWarningShown = false;
  port.setSignals = async (signals) => {
    try {
      return await origSetSignals(signals);
    } catch (_) {
      if (!signalWarningShown) {
        onLog?.('(DTR/RTS not supported on this port — skipping reset)\n');
        signalWarningShown = true;
      }
    }
  };

  const transport = new Transport(port);
  const esploader = new ESPLoader({
    transport,
    baudrate: 115200,
    romBaudrate: 115200,
    terminal: {
      clean() {},
      writeLine(data) {
        flushBuffer();
        onLog?.(data + (data.endsWith('\n') ? '' : '\n'));
      },
      write(data) {
        normalizeAndLog(data);
      },
    },
  });

  onLog?.(`Connecting to ${targetLabel}...\n`);
  if (resetMode === 'no_reset') {
    onLog?.('(no_reset mode — device must be in bootloader already)\n');
  }
  const chip = await esploader.main(resetMode);
  flushBuffer();
  onLog?.(`Detected: ${chip}\n`);

  onLog?.('Fetching firmware...\n');
  const resp = await fetch(firmwareUrl);
  if (!resp.ok) throw new Error(`Failed to fetch firmware (${firmwareUrl}): ${resp.status}`);
  const firmwareBuf = await resp.arrayBuffer();
  const firmwareBin = arrayBufferToBinaryString(firmwareBuf);
  onLog?.(`Firmware size: ${(firmwareBuf.byteLength / 1024).toFixed(0)} KB\n`);

  onLog?.('Erasing flash...\n');
  await esploader.eraseFlash();
  flushBuffer();
  onLog?.('Erase complete.\n');

  const addrHex = '0x' + flashAddress.toString(16);
  onLog?.(`Writing firmware at ${addrHex}...\n`);
  await esploader.writeFlash({
    fileArray: [{ data: firmwareBin, address: flashAddress }],
    flashSize: 'keep',
    flashMode: 'keep',
    flashFreq: 'keep',
    eraseAll: false,
    compress: true,
    reportProgress(fileIndex, written, total) {
      const pct = Math.round((written / total) * 100);
      onProgress?.(pct);
    },
  });
  flushBuffer();

  onLog?.('\nFirmware written successfully!\n');
  onLog?.('Hard resetting device...\n');
  await esploader.after('hard_reset');
  flushBuffer();

  try {
    await transport.disconnect();
  } catch (_) {
    // already closed
  }

  onLog?.('Done! Device is ready.\n');
}
