// WebSocket sync transport — v3 chirp-sync Phase 2.
//
// The extension is always the WebSocket client; Android (or future iOS /
// desktop sync hosts) is always the server. Each binary WebSocket message
// carries exactly one Noise frame — no length prefix, no fragmentation.
//
// On connect:
//   1. Open ws://<ip>:<port>/
//   2. Run Noise XX as the initiator, using our long-term sync identity
//      static keypair.
//   3. Verify the responder's static public key matches the paired-device
//      registry entry we expect. Mismatch is a hard fail — no implicit
//      pubkey trust on the LAN.
//   4. Hand the caller a `ChirpSyncSession` for the SyncMessage protocol.
//
// Subnet probe falls back to scanning the /24 returned by `getLocalSubnet`
// when the cached last-known IP fails. We use a short connection timeout
// (~500ms per IP) and parallel batches of 16 to keep the worst-case scan
// under ~10 seconds.

import { noiseXXInitiator, type NoiseXXSession } from "./noise-xx";
import { hexToPubkey, pubkeyToHex } from "./sync-identity";

const DEFAULT_PROBE_TIMEOUT_MS = 500;
const DEFAULT_PROBE_BATCH = 16;
const DEFAULT_CONNECT_TIMEOUT_MS = 4_000;
/** v2 brief §5.2 — subnet probe is rate-limited to once per 60 seconds. */
export const SUBNET_PROBE_COOLDOWN_MS = 60_000;

export interface ChirpSyncSession {
  send(plaintext: Uint8Array): void;
  receive(): Promise<Uint8Array>;
  close(): Promise<void>;
  remoteStaticPublicKey: Uint8Array;
}

export interface SyncEndpoint {
  ip: string;
  port: number;
}

/**
 * Open a WebSocket to (ip, port), run Noise XX, verify the responder's
 * long-term static key matches `expectedRemotePubkeyHex`, and return a
 * ready-to-use session.
 *
 * On a successful return, `setLastKnownIp` in the caller can persist
 * (ip, port) for the next session's fast path.
 */
export async function openChirpSyncSession(
  endpoint: SyncEndpoint,
  ourLongTermPrivateKey: Uint8Array,
  ourLongTermPublicKey: Uint8Array,
  expectedRemotePubkeyHex: string,
  options: { connectTimeoutMs?: number } = {},
): Promise<ChirpSyncSession> {
  const ws = await connectWebSocket(endpoint, options.connectTimeoutMs ?? DEFAULT_CONNECT_TIMEOUT_MS);
  const inbox: Uint8Array[] = [];
  const readers: Array<(b: Uint8Array) => void> = [];
  const errors: Error[] = [];
  let closed = false;

  ws.binaryType = "arraybuffer";
  ws.onmessage = (ev: MessageEvent) => {
    const data = ev.data;
    if (data instanceof ArrayBuffer) {
      const bytes = new Uint8Array(data);
      const r = readers.shift();
      if (r) r(bytes);
      else inbox.push(bytes);
    }
  };
  ws.onerror = () => {
    errors.push(new Error("WebSocket error"));
  };
  ws.onclose = () => {
    closed = true;
    while (readers.length > 0) {
      const r = readers.shift()!;
      r(new Uint8Array(0));
    }
  };

  function send(bytes: Uint8Array): Promise<void> {
    return new Promise((resolve, reject) => {
      try {
        ws.send(bytes as Uint8Array<ArrayBuffer>);
        resolve();
      } catch (err) {
        reject(err as Error);
      }
    });
  }
  function receive(): Promise<Uint8Array> {
    if (closed) return Promise.reject(new Error("WebSocket closed before frame"));
    const v = inbox.shift();
    if (v) return Promise.resolve(v);
    return new Promise<Uint8Array>((resolve, reject) => {
      readers.push((b) => {
        if (b.length === 0 && closed) reject(new Error("WebSocket closed"));
        else resolve(b);
      });
    });
  }

  let session: NoiseXXSession;
  try {
    session = await noiseXXInitiator(
      ourLongTermPrivateKey,
      ourLongTermPublicKey,
      send,
      receive,
    );
  } catch (err) {
    try { ws.close(); } catch { /* ignore */ }
    throw err;
  }

  const remoteHex = pubkeyToHex(session.remoteStaticPublicKey);
  if (remoteHex !== expectedRemotePubkeyHex) {
    try { ws.close(); } catch { /* ignore */ }
    throw new Error(
      `paired-device static key mismatch (got ${remoteHex.slice(0, 16)}…, ` +
      `expected ${expectedRemotePubkeyHex.slice(0, 16)}…)`,
    );
  }
  // Bind so the unused expectation-source check helper stays honest.
  void hexToPubkey(expectedRemotePubkeyHex);

  return {
    send(plaintext: Uint8Array) {
      const ciphertext = session.encrypt(plaintext);
      ws.send(ciphertext as Uint8Array<ArrayBuffer>);
    },
    async receive() {
      const ct = await receive();
      return session.decrypt(ct);
    },
    async close() {
      try { ws.close(); } catch { /* ignore */ }
    },
    remoteStaticPublicKey: session.remoteStaticPublicKey,
  };
}

function connectWebSocket(endpoint: SyncEndpoint, timeoutMs: number): Promise<WebSocket> {
  return new Promise((resolve, reject) => {
    let ws: WebSocket;
    try {
      ws = new WebSocket(`ws://${endpoint.ip}:${endpoint.port}/`);
    } catch (err) {
      reject(err as Error);
      return;
    }
    const t = setTimeout(() => {
      try { ws.close(); } catch { /* ignore */ }
      reject(new Error(`WebSocket connect to ${endpoint.ip}:${endpoint.port} timed out`));
    }, timeoutMs);
    ws.onopen = () => {
      clearTimeout(t);
      resolve(ws);
    };
    ws.onerror = () => {
      clearTimeout(t);
      reject(new Error(`WebSocket connect to ${endpoint.ip}:${endpoint.port} errored`));
    };
  });
}

/**
 * Per-pubkey cooldown tracker for the subnet probe. v2 brief §5.2 caps
 * full /24 scans at once per minute so an extension whose cached IP has
 * gone stale doesn't hammer the LAN with handshake attempts every time
 * the popup opens.
 *
 * Keyed by the expected remote pubkey because each paired peer is probed
 * independently; the rate limit applies to "scan looking for THIS peer",
 * not "scan the subnet at all".
 */
const lastProbeAt = new Map<string, number>();

export function subnetProbeOnCooldown(
  expectedRemotePubkeyHex: string,
  now: number = Date.now(),
  cooldownMs: number = SUBNET_PROBE_COOLDOWN_MS,
): boolean {
  const last = lastProbeAt.get(expectedRemotePubkeyHex);
  return last !== undefined && now - last < cooldownMs;
}

/** Test-only — clear the per-pubkey probe cooldown ledger. */
export function _resetSubnetProbeCooldown(): void {
  lastProbeAt.clear();
}

/**
 * Subnet probe — scan a /24 looking for a host that completes Noise XX
 * with the expected long-term static key. Returns the first match, or
 * null if the cooldown is still active or no host responded.
 *
 * Rate-limited to once per `SUBNET_PROBE_COOLDOWN_MS` per peer pubkey
 * (v2 brief §5.2). Callers should still attempt the cached
 * `last_known_ip`/`last_known_port` directly first — the subnet probe is
 * fallback only.
 *
 * For each IP we attempt the full Noise XX handshake, but with a tight
 * timeout: dead hosts fail at the WebSocket connect step in <1ms (browser
 * surfaces ECONNREFUSED immediately) or by the connectTimeout otherwise.
 */
export async function probeSubnetForPeer(
  subnet: string, // e.g. "192.168.1"
  port: number,
  ourLongTermPrivateKey: Uint8Array,
  ourLongTermPublicKey: Uint8Array,
  expectedRemotePubkeyHex: string,
  options: { batch?: number; timeoutMs?: number; cooldownMs?: number; now?: number } = {},
): Promise<SyncEndpoint | null> {
  const now = options.now ?? Date.now();
  const cooldownMs = options.cooldownMs ?? SUBNET_PROBE_COOLDOWN_MS;
  if (subnetProbeOnCooldown(expectedRemotePubkeyHex, now, cooldownMs)) {
    return null;
  }
  lastProbeAt.set(expectedRemotePubkeyHex, now);
  const batchSize = options.batch ?? DEFAULT_PROBE_BATCH;
  const timeoutMs = options.timeoutMs ?? DEFAULT_PROBE_TIMEOUT_MS;
  for (let start = 1; start <= 254; start += batchSize) {
    const end = Math.min(start + batchSize - 1, 254);
    const ips: string[] = [];
    for (let i = start; i <= end; i++) ips.push(`${subnet}.${i}`);

    const results = await Promise.all(ips.map((ip) => probeOne(
      { ip, port },
      ourLongTermPrivateKey,
      ourLongTermPublicKey,
      expectedRemotePubkeyHex,
      timeoutMs,
    )));
    const hit = results.find((r): r is SyncEndpoint => r !== null);
    if (hit) return hit;
  }
  return null;
}

async function probeOne(
  endpoint: SyncEndpoint,
  ourLongTermPrivateKey: Uint8Array,
  ourLongTermPublicKey: Uint8Array,
  expectedRemotePubkeyHex: string,
  timeoutMs: number,
): Promise<SyncEndpoint | null> {
  try {
    const session = await openChirpSyncSession(
      endpoint,
      ourLongTermPrivateKey,
      ourLongTermPublicKey,
      expectedRemotePubkeyHex,
      { connectTimeoutMs: timeoutMs },
    );
    await session.close();
    return endpoint;
  } catch {
    return null;
  }
}