import { createCipheriv, createDecipheriv, createHash, randomBytes } from "node:crypto";
import { deflateRawSync, inflateRawSync } from "node:zlib";
import { writeFileSync } from "node:fs";
import { resolve, dirname } from "node:path";
import BN from "bn.js";
import { ec as EllipticEC } from "elliptic";

const DATA_TYPE = {
  PROTOCOL_LIST: 1,
  SELECTED_PROTOCOL: 2,
  SESSION_KEY_EXCHANGE: 3,
  ERROR: 4,
  DATA: 5,
  ECJPAKE_ROUND_ONE: 32,
  ECJPAKE_ROUND_TWO: 33,
} as const;

function usage(exitCode = 0) {
  const script = "bun run fetch-devices";
  const text = [
    `Usage: ${script} <ip> <passcode>`,
    "",
    "Examples:",
    `  ${script} 192.168.31.166 048889`,
    "",
    "Notes:",
    "  - The passcode is read only from argv and is not written to disk.",
    "  - Output is written to src/devices.json.",
  ].join("\n");
  (exitCode === 0 ? console.log : console.error)(text);
  process.exit(exitCode);
}

function normalizeUrl(input: string): string {
  let url: URL;
  if (/^wss?:\/\//i.test(input)) {
    url = new URL(input);
  } else if (/^https?:\/\//i.test(input)) {
    url = new URL(input);
    url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
  } else {
    url = new URL(`ws://${input.replace(/\/+$/, "")}`);
  }

  if (!url.pathname.endsWith("/centrallinkws/")) {
    url.pathname = `${url.pathname.replace(/\/+$/, "")}/centrallinkws/`;
  }
  return url.toString();
}

function concatBytes(...parts: Array<Buffer | Uint8Array>): Buffer {
  return Buffer.concat(parts.map((part) => Buffer.from(part)));
}

function bnBytes(bn: BN, length = 32): Buffer {
  return Buffer.from(bn.toArray("be", length));
}

function sha256Bytes(buf: Buffer | Uint8Array): Buffer {
  return createHash("sha256").update(Buffer.from(buf)).digest();
}

class ECJPAKE {
  role: string;
  peerRole: string;
  secret: BN;
  curveName: string;
  serverRoundTwoPrefix = 22;
  failed = false;
  wroteRoundOne = false;
  readRoundOneDone = false;
  wroteRoundTwo = false;
  readRoundTwoDone = false;
  x1: ReturnType<InstanceType<typeof EllipticEC>["genKeyPair"]> | null = null;
  x2: ReturnType<InstanceType<typeof EllipticEC>["genKeyPair"]> | null = null;
  peerX1: ReturnType<InstanceType<typeof EllipticEC>["keyFromPublic"]>["pub"] | null = null;
  peerX2: ReturnType<InstanceType<typeof EllipticEC>["keyFromPublic"]>["pub"] | null = null;

  constructor({ role, secret }: { role: string; secret: string }) {
    if (role !== "server" && role !== "client") {
      throw new TypeError('role must be "client" or "server"');
    }
    if (typeof secret !== "string") {
      throw new TypeError("secret must be a string");
    }
    this.role = role;
    this.peerRole = role === "client" ? "server" : "client";
    this.secret = new BN(Buffer.from(secret, "utf8"));
    this.curveName = "secp256k1";
  }

  writeRoundOne(): Buffer {
    this.assertUsable();
    if (this.wroteRoundOne || this.wroteRoundTwo || this.readRoundTwoDone) {
      throw new Error("Wrong step");
    }
    this.wroteRoundOne = true;
    const ec = new EllipticEC(this.curveName);
    this.x1 = ec.genKeyPair();
    this.x2 = ec.genKeyPair();
    const zkp1 = this.createZkp(ec, ec.g, this.x1.getPublic(), this.x1.getPrivate(), this.role);
    const zkp2 = this.createZkp(ec, ec.g, this.x2.getPublic(), this.x2.getPrivate(), this.role);
    return concatBytes(
      this.encodePublic(this.x1.getPublic()),
      this.encodeZkp(zkp1),
      this.encodePublic(this.x2.getPublic()),
      this.encodeZkp(zkp2),
    );
  }

  readRoundOne(payload: Uint8Array): void {
    this.assertUsable();
    if (this.readRoundOneDone || this.wroteRoundTwo || this.readRoundTwoDone) {
      throw new Error("Wrong step");
    }
    this.readRoundOneDone = true;
    const ec = new EllipticEC(this.curveName);
    let offset = 0;

    const peerX1Len = payload[offset++]!;
    const peerX1Wire = payload.slice(offset, offset + peerX1Len);
    offset += peerX1Len;
    this.peerX1 = ec.keyFromPublic(Buffer.from(peerX1Wire)).getPublic();

    const zkp1Length = 1 + payload[offset]! + 1 + payload[offset + 1 + payload[offset]!]!;
    const zkp1 = this.decodeZkp(payload.slice(offset, offset + zkp1Length));
    offset += zkp1Length;
    if (!this.verifyZkp(ec, ec.g, this.peerX1!, zkp1.V, zkp1.r, this.peerRole)) {
      this.failed = true;
      throw new Error("ECJPAKE round one failed");
    }

    const peerX2Len = payload[offset++]!;
    const peerX2Wire = payload.slice(offset, offset + peerX2Len);
    offset += peerX2Len;
    this.peerX2 = ec.keyFromPublic(Buffer.from(peerX2Wire)).getPublic();

    const zkp2Length = 1 + payload[offset]! + 1 + payload[offset + 1 + payload[offset]!]!;
    const zkp2 = this.decodeZkp(payload.slice(offset, offset + zkp2Length));
    if (!this.verifyZkp(ec, ec.g, this.peerX2!, zkp2.V, zkp2.r, this.peerRole)) {
      this.failed = true;
      throw new Error("ECJPAKE round one failed");
    }
  }

  writeRoundTwo(): Buffer {
    this.assertUsable();
    if (this.wroteRoundTwo || !this.wroteRoundOne || !this.readRoundOneDone) {
      throw new Error("Wrong step");
    }
    this.wroteRoundTwo = true;
    const ec = new EllipticEC(this.curveName);
    const generator = this.x1!.getPublic().add(this.peerX1!).add(this.peerX2!);
    ec.g = generator;
    const n = new BN(randomBytes(16)).mul(ec.n).add(this.secret);
    const s = this.x2!.getPrivate().mul(n).umod(ec.n);
    const publicKey = generator.mul(s);
    const zkp = this.createZkp(ec, generator, publicKey, s, this.role);
    const encodedPublic = this.encodePublic(publicKey);
    const encodedZkp = this.encodeZkp(zkp);

    if (this.role === "server") {
      return concatBytes(this.encodeServerRoundTwoPrefix(), encodedPublic, encodedZkp);
    }
    return concatBytes(encodedPublic, encodedZkp);
  }

  readRoundTwo(payload: Uint8Array): Buffer {
    this.assertUsable();
    if (this.readRoundTwoDone || !this.wroteRoundOne || !this.readRoundOneDone) {
      throw new Error("Wrong step");
    }
    this.readRoundTwoDone = true;
    const ec = new EllipticEC(this.curveName);
    const generator = this.x1!.getPublic().add(this.x2!.getPublic()).add(this.peerX1!);
    let offset = this.role === "client" ? 3 : 0;

    const publicLen = payload[offset++]!;
    const publicWire = payload.slice(offset, offset + publicLen);
    offset += publicLen;
    const peerPublic = ec.keyFromPublic(Buffer.from(publicWire)).getPublic();

    const zkpLength = 1 + payload[offset]! + 1 + payload[offset + 1 + payload[offset]!]!;
    const zkp = this.decodeZkp(payload.slice(offset, offset + zkpLength));
    if (!this.verifyZkp(ec, generator, peerPublic, zkp.V, zkp.r, this.peerRole)) {
      this.failed = true;
      throw new Error("ECJPAKE round two failed");
    }

    const n = new BN(randomBytes(16)).mul(ec.n).add(this.secret);
    const m = this.x2!.getPrivate().mul(n).umod(ec.n);
    const sharedPoint = peerPublic.add(this.peerX2!.mul(m).neg()).mul(this.x2!.getPrivate());
    return sha256Bytes(bnBytes(sharedPoint.getX(), 32));
  }

  assertUsable(): void {
    if (this.failed) {
      throw new Error("Reusing failed ECJPAKE context is insecure.");
    }
  }

  encodePointForHash(point: InstanceType<typeof EllipticEC>["point"]): Buffer {
    const out = Buffer.alloc(69);
    out.writeUInt32BE(65, 0);
    out[4] = 4;
    bnBytes(point.getX(), 32).copy(out, 5);
    bnBytes(point.getY(), 32).copy(out, 37);
    return out;
  }

  encodePublic(point: InstanceType<typeof EllipticEC>["point"]): Buffer {
    const out = Buffer.alloc(66);
    out[0] = 65;
    out[1] = 4;
    bnBytes(point.getX(), 32).copy(out, 2);
    bnBytes(point.getY(), 32).copy(out, 34);
    return out;
  }

  encodeServerRoundTwoPrefix(): Buffer {
    const out = Buffer.alloc(3);
    out[0] = 3;
    out.writeUInt16BE(this.serverRoundTwoPrefix, 1);
    return out;
  }

  zkpChallenge(
    generator: InstanceType<typeof EllipticEC>["point"],
    v: InstanceType<typeof EllipticEC>["point"],
    x: InstanceType<typeof EllipticEC>["point"],
    signerId: string,
    order: BN,
  ): BN {
    const id = Buffer.from(signerId, "utf8");
    const len = Buffer.alloc(4);
    len.writeUInt32BE(id.length, 0);
    const hashInput = concatBytes(
      this.encodePointForHash(generator),
      this.encodePointForHash(v),
      this.encodePointForHash(x),
      len,
      id,
    );
    return new BN(sha256Bytes(hashInput).toString("hex"), "hex", "be").umod(order);
  }

  createZkp(
    ec: InstanceType<typeof EllipticEC>,
    generator: InstanceType<typeof EllipticEC>["point"],
    publicKey: InstanceType<typeof EllipticEC>["point"],
    privateKey: BN,
    signerId: string,
  ): { V: InstanceType<typeof EllipticEC>["point"]; r: BN } {
    const zkpEc = new EllipticEC(this.curveName);
    zkpEc.g = generator;
    const vKey = zkpEc.genKeyPair();
    const challenge = this.zkpChallenge(zkpEc.g, vKey.getPublic(), publicKey, signerId, zkpEc.n);
    const r = vKey.getPrivate().sub(challenge.mul(privateKey)).umod(zkpEc.n);
    return { V: vKey.getPublic(), r };
  }

  encodeZkp(zkp: { V: InstanceType<typeof EllipticEC>["point"]; r: BN }): Buffer {
    const out = Buffer.alloc(99);
    this.encodePublic(zkp.V).copy(out, 0);
    out[66] = 32;
    bnBytes(zkp.r, 32).copy(out, 67);
    return out;
  }

  decodeZkp(payload: Uint8Array): { V: InstanceType<typeof EllipticEC>["point"]; r: BN } {
    const ec = new EllipticEC(this.curveName);
    const publicLen = payload[0]!;
    const publicWire = payload.slice(1, 1 + publicLen);
    const rLenOffset = 1 + publicLen;
    const rLen = payload[rLenOffset]!;
    const rBytes = payload.slice(rLenOffset + 1, rLenOffset + 1 + rLen);
    return {
      V: ec.keyFromPublic(Buffer.from(publicWire)).getPublic(),
      r: new BN(Buffer.from(rBytes)),
    };
  }

  verifyZkp(
    ec: InstanceType<typeof EllipticEC>,
    generator: InstanceType<typeof EllipticEC>["point"],
    publicKey: InstanceType<typeof EllipticEC>["point"],
    v: InstanceType<typeof EllipticEC>["point"],
    r: BN,
    signerId: string,
  ): boolean {
    const verifyEc = new EllipticEC(this.curveName);
    verifyEc.g = generator;
    const challenge = this.zkpChallenge(generator, v, publicKey, signerId, verifyEc.n);
    return publicKey.mul(challenge).add(generator.mul(r)).eq(v);
  }
}

class CounterGcm {
  key: Buffer;
  salt: Buffer;
  selfCounter = 1;
  peerCounter = 0;

  constructor(key: Buffer, salt: Buffer) {
    this.key = Buffer.from(key);
    this.salt = Buffer.from(salt);
    if (this.key.length !== 16) throw new Error("AES key must be 16 bytes");
    if (this.salt.length !== 8) throw new Error("AES salt must be 8 bytes");
  }

  iv(counter: number): Buffer {
    const iv = Buffer.alloc(12);
    this.salt.copy(iv, 0);
    iv.writeUInt32LE(counter, 8);
    return iv;
  }

  encrypt(plain: Buffer | Uint8Array): Buffer {
    const counter = this.selfCounter++;
    const cipher = createCipheriv("aes-128-gcm", this.key, this.iv(counter), {
      authTagLength: 16,
    });
    const ciphertext = Buffer.concat([cipher.update(Buffer.from(plain)), cipher.final()]);
    const tag = cipher.getAuthTag();
    const out = Buffer.alloc(4 + ciphertext.length + tag.length);
    out.writeUInt32LE(counter, 0);
    ciphertext.copy(out, 4);
    tag.copy(out, 4 + ciphertext.length);
    return out;
  }

  decrypt(frame: Buffer | Uint8Array): Buffer {
    const input = Buffer.from(frame);
    const counter = input.readUInt32LE(0);
    if (counter <= this.peerCounter) {
      throw new Error(`Replay or out-of-order frame: ${counter}`);
    }
    this.peerCounter = counter;
    const tag = input.subarray(input.length - 16);
    const ciphertext = input.subarray(4, input.length - 16);
    const decipher = createDecipheriv("aes-128-gcm", this.key, this.iv(counter), {
      authTagLength: 16,
    });
    decipher.setAuthTag(tag);
    return Buffer.concat([decipher.update(ciphertext), decipher.final()]);
  }
}

function pack(type: number, payload: Buffer = Buffer.alloc(0)): Buffer {
  return Buffer.concat([Buffer.from([type]), Buffer.from(payload)]);
}

function compressJson(obj: unknown): Buffer {
  const raw = Buffer.from(JSON.stringify(obj), "utf8");
  const deflated = deflateRawSync(raw);
  const out = Buffer.alloc(4 + deflated.length);
  out.writeUInt32LE(raw.length, 0);
  deflated.copy(out, 4);
  return out;
}

function decompressJson(payload: Buffer | Uint8Array): unknown {
  const input = Buffer.from(payload);
  const expectedLength = input.readUInt32LE(0);
  const raw = inflateRawSync(input.subarray(4));
  if (raw.length !== expectedLength) {
    throw new Error(`Inflated JSON length mismatch: ${raw.length} != ${expectedLength}`);
  }
  return JSON.parse(raw.toString("utf8"));
}

async function wsDataToBuffer(data: unknown): Promise<Buffer> {
  if (typeof data === "string") {
    throw new Error("Unexpected websocket text frame");
  }
  if (data instanceof ArrayBuffer) return Buffer.from(data);
  if (ArrayBuffer.isView(data)) {
    return Buffer.from(data.buffer, data.byteOffset, data.byteLength);
  }
  if (data && typeof (data as Blob).arrayBuffer === "function") {
    return Buffer.from(await (data as Blob).arrayBuffer());
  }
  return Buffer.from(data as ArrayBufferView);
}

class GatewayClient {
  url: string;
  passcode: string;
  stage = "init";
  ecjpake: ECJPAKE | null = null;
  sessionCipher: CounterGcm | null = null;
  outCipher: CounterGcm | null = null;
  inCipher: CounterGcm | null = null;
  ready = false;
  nextId = 0;
  pending = new Map<string, { resolve: (value: unknown) => void; reject: (reason: Error) => void; timeout: ReturnType<typeof setTimeout> }>();
  ws!: WebSocket;

  constructor(url: string, passcode: string) {
    this.url = url;
    this.passcode = passcode;
  }

  async connect(): Promise<void> {
    await new Promise<void>((resolve, reject) => {
      const timeout = setTimeout(() => reject(new Error("WebSocket connect timeout")), 10000);
      this.ws = new WebSocket(this.url);
      this.ws.binaryType = "arraybuffer";

      this.ws.onopen = () => {
        clearTimeout(timeout);
        this.sendProtocolList();
      };
      this.ws.onerror = () => {
        clearTimeout(timeout);
        reject(new Error(`WebSocket error while connecting to ${this.url}`));
      };
      this.ws.onclose = (event) => {
        if (!this.ready) {
          clearTimeout(timeout);
          reject(new Error(`WebSocket closed before ready: ${event.code} ${event.reason}`));
        }
      };
      this.ws.onmessage = async (event) => {
        try {
          this.handleMessage(await wsDataToBuffer(event.data));
          if (this.ready) resolve();
        } catch (err) {
          reject(err);
        }
      };
    });
  }

  close(): void {
    try {
      this.ws?.close();
    } catch {}
  }

  send(buf: Buffer): void {
    this.ws.send(Buffer.from(buf));
  }

  log(message: string): void {
    if (process.env.MIOT_TRACE === "1") console.error(`[trace] ${message}`);
  }

  sendProtocolList(): void {
    this.stage = "protocol-list";
    this.log("send PROTOCOL_LIST");
    this.send(pack(DATA_TYPE.PROTOCOL_LIST, Buffer.from(JSON.stringify(["passcode"]))));
  }

  handleMessage(frame: Buffer): void {
    const type = frame[0]!;
    const payload = frame.subarray(1);
    this.log(`recv type=${type} len=${frame.length}`);

    switch (type) {
      case DATA_TYPE.SELECTED_PROTOCOL:
        return this.handleSelectedProtocol(payload);
      case DATA_TYPE.ECJPAKE_ROUND_ONE:
        return this.handleRoundOne(payload);
      case DATA_TYPE.ECJPAKE_ROUND_TWO:
        return this.handleRoundTwo(payload);
      case DATA_TYPE.SESSION_KEY_EXCHANGE:
        return this.handleSessionKeyExchange(payload);
      case DATA_TYPE.DATA:
        return this.handleData(payload);
      case DATA_TYPE.ERROR:
        throw new Error(`Gateway returned ERROR frame during ${this.stage}`);
      default:
        throw new Error(`Unexpected frame type: ${type}`);
    }
  }

  handleSelectedProtocol(payload: Buffer): void {
    this.stage = "selected-protocol";
    const selected = JSON.parse(payload.toString("utf8")) as { protocol: string };
    if (selected.protocol !== "passcode") {
      throw new Error(`Gateway selected unsupported protocol: ${selected.protocol}`);
    }
    this.ecjpake = new ECJPAKE({ role: "client", secret: this.passcode });
    this.stage = "ecjpake-round-one";
    this.log("send ECJPAKE_ROUND_ONE");
    this.send(pack(DATA_TYPE.ECJPAKE_ROUND_ONE, this.ecjpake.writeRoundOne()));
  }

  handleRoundOne(payload: Buffer): void {
    this.stage = "ecjpake-round-two";
    this.ecjpake!.readRoundOne(new Uint8Array(payload));
    this.log("send ECJPAKE_ROUND_TWO");
    this.send(pack(DATA_TYPE.ECJPAKE_ROUND_TWO, this.ecjpake!.writeRoundTwo()));
  }

  handleRoundTwo(payload: Buffer): void {
    this.stage = "session-key-exchange";
    const shared = Buffer.from(this.ecjpake!.readRoundTwo(new Uint8Array(payload)));
    this.sessionCipher = new CounterGcm(shared.subarray(0, 16), shared.subarray(16, 24));

    const localMaterial = randomBytes(24);
    this.outCipher = new CounterGcm(localMaterial.subarray(0, 16), localMaterial.subarray(16, 24));
    this.log("send SESSION_KEY_EXCHANGE");
    this.send(pack(DATA_TYPE.SESSION_KEY_EXCHANGE, this.sessionCipher.encrypt(localMaterial)));
  }

  handleSessionKeyExchange(payload: Buffer): void {
    this.stage = "secure-session";
    const remoteMaterial = this.sessionCipher!.decrypt(payload);
    this.inCipher = new CounterGcm(remoteMaterial.subarray(0, 16), remoteMaterial.subarray(16, 24));
    this.ready = true;
    this.log("secure session ready");
  }

  handleData(payload: Buffer): void {
    const rpc = decompressJson(this.inCipher!.decrypt(payload)) as { id: number; error?: { message: string }; result?: unknown };
    const key = String(rpc.id);
    const pending = this.pending.get(key);
    if (!pending) return;

    clearTimeout(pending.timeout);
    this.pending.delete(key);
    if ("error" in rpc && rpc.error) {
      pending.reject(new Error(rpc.error?.message || JSON.stringify(rpc.error)));
    } else {
      pending.resolve(rpc.result);
    }
  }

  async callAPI(method: string, params: Record<string, unknown> = {}, timeoutMs = 10000): Promise<unknown> {
    if (!this.ready) throw new Error("Secure session is not ready");
    this.stage = `api:${method}`;
    const id = this.nextId++;
    const rpc = {
      jsonrpc: "2.0",
      id,
      method: `/api/${method}`,
      params,
    };
    const encrypted = this.outCipher!.encrypt(compressJson(rpc));
    this.send(pack(DATA_TYPE.DATA, encrypted));

    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.pending.delete(String(id));
        reject(new Error(`API timeout: ${method}`));
      }, timeoutMs);
      this.pending.set(String(id), { resolve, reject, timeout });
    });
  }
}

function getOutputPath(): string {
  return resolve(dirname(import.meta.file), "..", "devices.json");
}

async function main(): Promise<void> {
  const args = process.argv.slice(2);
  if (args.includes("-h") || args.includes("--help")) usage(0);
  if (args[0] === "--self-test") {
    const client = new ECJPAKE({ role: "client", secret: "000000" });
    const server = new ECJPAKE({ role: "server", secret: "000000" });
    const clientRoundOne = client.writeRoundOne();
    const serverRoundOne = server.writeRoundOne();
    if (clientRoundOne.length !== 330 || serverRoundOne.length !== 330) {
      throw new Error(
        `Unexpected ECJPAKE round-one lengths: ${clientRoundOne.length}, ${serverRoundOne.length}`,
      );
    }
    client.readRoundOne(serverRoundOne);
    server.readRoundOne(clientRoundOne);
    const clientRoundTwo = client.writeRoundTwo();
    const serverRoundTwo = server.writeRoundTwo();
    if (clientRoundTwo.length !== 165 || serverRoundTwo.length !== 168) {
      throw new Error(
        `Unexpected ECJPAKE round-two lengths: ${clientRoundTwo.length}, ${serverRoundTwo.length}`,
      );
    }
    const clientKey = Buffer.from(client.readRoundTwo(serverRoundTwo));
    const serverKey = Buffer.from(server.readRoundTwo(clientRoundTwo));
    if (clientKey.length !== 32 || !clientKey.equals(serverKey)) {
      throw new Error("ECJPAKE self-test shared key mismatch");
    }
    console.error("Self-test OK");
    return;
  }
  if (args.length !== 2) usage(1);

  const url = normalizeUrl(args[0]!);
  const passcode = args[1]!;
  if (!/^\d{6}$/.test(passcode)) {
    throw new Error("Passcode must be a 6-digit string");
  }

  const outputPath = getOutputPath();
  console.error(`Connecting to ${url}`);
  const client = new GatewayClient(url, passcode);
  try {
    await client.connect();
    console.error("Secure session established");
    const result = (await client.callAPI("getDevList", {}, 10000)) as { devList?: Record<string, unknown> };
    const devList = result.devList || {};
    const payload = {
      fetchedAt: new Date().toISOString(),
      url,
      count: Object.keys(devList).length,
      devList,
    };
    writeFileSync(outputPath, `${JSON.stringify(payload, null, 2)}\n`);
    console.error(`已更新设备清单: ${outputPath} (${Object.keys(devList).length} 台设备)`);
  } finally {
    client.close();
  }
}

main().catch((err: unknown) => {
  console.error(`Error: ${err instanceof Error ? err.message : String(err)}`);
  process.exitCode = 1;
});