// WebSocket sync transport — standardized on Noise NK (see
// HANDSHAKE_STANDARDIZE_ON_NK.md). Android's ongoing-sync listener
// (SyncReconnectListener.kt) is a Noise NK responder on the same fixed port
// (49142) used for pairing; there is one listener and one handshake pattern.
//
// The extension is always the WebSocket client; Android (or future iOS /
// native 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 NK as the initiator against the paired device's known
//      long-term static public key. NK authenticates only the
//      responder during the handshake — there is no initiator static
//      keypair in the crypto itself.
//   3. Exchange PAIR_IDENTITY frames over the encrypted channel (same frame
//      Android's SyncReconnectListener expects before it will treat the
//      connection as sync traffic rather than an unknown caller). Our
//      identity is asserted, not proven, at this layer — this mirrors
//      Android's own implementation, not a gap introduced here.
//   4. Hand the caller a `SyncSession` 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 { noiseNKInitiator, type NoiseNKSession } from "./noise-js";
import { SYNC_LIMITS } from "./sync-limits";
import { hexToPubkey } from "./sync-identity";
import type { DeviceType } from "./sync-types";

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 SyncSession {
  send(plaintext: Uint8Array): void;
  receive(): Promise<Uint8Array>;
  close(): Promise<void>;
  remoteStaticPublicKey: Uint8Array;
}

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

/** Our identity, asserted in the post-handshake PAIR_IDENTITY frame — the
 *  same shape `code-pairing.ts` sends during pairing (`device_id` etc.). */
export interface SyncPeerIdentity {
  deviceId: string;
  deviceName: string;
  deviceType: DeviceType;
  longTermPublicKeyHex: string;
  signingPublicKeyHex?: string;
}

export class SyncFrameTooLargeError extends Error {
  constructor(message = "sync frame too large (WebSocket close 1009)") {
    super(message);
    this.name = "SyncFrameTooLargeError";
  }
}

function encodePairIdentity(id: SyncPeerIdentity): Uint8Array {
  return new TextEncoder().encode(JSON.stringify({
    type: "PAIR_IDENTITY",
    device_id: id.deviceId,
    device_name: id.deviceName,
    device_type: id.deviceType,
    long_term_public_key_hex: id.longTermPublicKeyHex,
    signing_public_key_hex: id.signingPublicKeyHex,
  }));
}

/**
 * Open a WebSocket to (ip, port), run Noise NK against the paired device's
 * known static key, exchange PAIR_IDENTITY, 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 openSyncSession(
  endpoint: SyncEndpoint,
  ourIdentity: SyncPeerIdentity,
  expectedRemotePubkeyHex: string,
  options: { connectTimeoutMs?: number } = {},
): Promise<SyncSession> {
  const ws = await connectWebSocket(endpoint, options.connectTimeoutMs ?? DEFAULT_CONNECT_TIMEOUT_MS);
  const inbox: Uint8Array[] = [];
  const readers: Array<(b: Uint8Array) => void> = [];
  let closed = false;
  let closeError: Error = new Error("WebSocket closed");

  ws.binaryType = "arraybuffer";
  ws.onmessage = (ev: MessageEvent) => {
    const data = ev.data;
    if (data instanceof ArrayBuffer) {
      // MV-3: refuse oversized frames before they are buffered or decrypted.
      // A peer that sends one is misbehaving — tear the session down rather
      // than letting the inbox grow toward OOM.
      if (data.byteLength > SYNC_LIMITS.MAX_FRAME_BYTES) {
        // MV-3: oversized frame — tear the session down. close() routes through
        // onclose, which rejects any parked receive(), so no error array is
        // needed (PCH-01-026 #9).
        console.warn(
          `[Peach] sync frame exceeds max size (${SYNC_LIMITS.MAX_FRAME_BYTES} bytes); closing`,
        );
        closeError = new SyncFrameTooLargeError(
          `sync frame exceeds ${SYNC_LIMITS.MAX_FRAME_BYTES} bytes`,
        );
        try { ws.close(1009, "sync frame too large"); } catch { /* ignore */ }
        return;
      }
      const bytes = new Uint8Array(data);
      const r = readers.shift();
      if (r) r(bytes);
      else inbox.push(bytes);
    }
  };
  ws.onerror = () => {
    // PCH-01-026 #9: surface the error and unblock any parked receive() by
    // closing. Browsers usually follow error with close, but not always; an
    // explicit close guarantees the handshake fails fast instead of hanging.
    console.warn("[Peach] sync WebSocket error; closing transport");
    try { ws.close(); } catch { /* ignore */ }
  };
  ws.onclose = (event: CloseEvent) => {
    closed = true;
    if (event.code === 1009) {
      closeError = new SyncFrameTooLargeError();
    }
    while (readers.length > 0) {
      const r = readers.shift()!;
      r(new Uint8Array(0));
    }
  };

  function send(bytes: Uint8Array): Promise<void> {
    return new Promise((resolve, reject) => {
      if (bytes.byteLength > SYNC_LIMITS.MAX_FRAME_BYTES) {
        reject(
          new SyncFrameTooLargeError(
            `sync frame exceeds ${SYNC_LIMITS.MAX_FRAME_BYTES} bytes`,
          ),
        );
        return;
      }
      try {
        ws.send(bytes as Uint8Array<ArrayBuffer>);
        resolve();
      } catch (err) {
        reject(err as Error);
      }
    });
  }
  function receive(): Promise<Uint8Array> {
    if (closed) return Promise.reject(closeError);
    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(closeError);
        else resolve(b);
      });
    });
  }

  let session: NoiseNKSession;
  try {
    // NK's responder static key is a handshake INPUT, not something learned
    // mid-handshake — a wrong key simply fails msg2 decryption below, so
    // there is no separate mismatch check to run first.
    session = await noiseNKInitiator(hexToPubkey(expectedRemotePubkeyHex), send, receive);
  } catch (err) {
    try { ws.close(); } catch { /* ignore */ }
    throw err;
  }

  try {
    await send(session.encrypt(encodePairIdentity(ourIdentity)));
    const theirFrame = JSON.parse(new TextDecoder().decode(session.decrypt(await receive())));
    if (theirFrame?.type !== "PAIR_IDENTITY") {
      throw new Error(`expected PAIR_IDENTITY, got ${theirFrame?.type as string}`);
    }
  } catch (err) {
    try { ws.close(); } catch { /* ignore */ }
    throw err;
  }

  return {
    send(plaintext: Uint8Array) {
      const ciphertext = session.encrypt(plaintext);
      if (ciphertext.byteLength > SYNC_LIMITS.MAX_FRAME_BYTES) {
        throw new SyncFrameTooLargeError(
          `sync frame exceeds ${SYNC_LIMITS.MAX_FRAME_BYTES} bytes`,
        );
      }
      ws.send(ciphertext as Uint8Array<ArrayBuffer>);
    },
    async receive() {
      const ct = await receive();
      return session.decrypt(ct);
    },
    async close() {
      try { ws.close(); } catch { /* ignore */ }
    },
    remoteStaticPublicKey: hexToPubkey(expectedRemotePubkeyHex),
  };
}

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 NK
 * 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 NK 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,
  ourIdentity: SyncPeerIdentity,
  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 },
      ourIdentity,
      expectedRemotePubkeyHex,
      timeoutMs,
    )));
    const hit = results.find((r): r is SyncEndpoint => r !== null);
    if (hit) return hit;
  }
  return null;
}

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