// Encrypted session on top of a Relay.
//
// Mirrors `secure_connection` + login-request handling in
// /Users/sn0/Desktop/rustdesk/src/client.rs:758-834. Wire shape:
//
//   Peer  ->  Us:   Message{signed_id: <signed IdPk{id, curve25519_box_pk}>}
//             - signed by peer's Ed25519 SIGN sk
//             - we verify with peer_sign_pk (extracted from rendezvous's signed_id_pk)
//             - inner pk is the peer's CURVE25519 BOX pk
//   Us    ->  Peer: Message{public_key: {asymmetric_value, symmetric_value}}
//             - asymmetric_value = our ephemeral Curve25519 box pk
//             - symmetric_value = box_seal(secretbox_key, peer_box_pk, our_box_sk)
//             - this message is unencrypted (the LAST one before secretbox stream)
//   [from now on every Message is secretbox(seqN, secretbox_key)]
//   Peer  ->  Us:   Message{hash: {salt, challenge}}
//   Us    ->  Peer: Message{login_request: {username, password, my_id, ...}}
//             - password = SHA256(SHA256(pwd_bytes || base64dec(salt)) || base64dec(challenge))
//             - empty password = SHA256(SHA256("" || salt) || challenge) on the desktop
//   Peer  ->  Us:   Message{login_response: {peer_info or error}}
//
// Sequence-counter rule (the bit that bites every reimplementation —
// see libs/hbb_common/src/tcp.rs:296-344): SEND and RECV each maintain
// their OWN counter starting at 0, only the secretbox-encrypted frames
// are counted, both directions advance independently.

import { Relay } from "./relay.js";
import { hbb } from "../proto/generated.js";
import {
  signOpen,
  genBoxKeypair,
  genSecretboxKey,
  boxSeal,
  secretboxSeal,
  secretboxOpen,
  sha256,
  concat,
} from "../crypto.js";

/** Parameters for the secure-then-login handshake. */
export interface SessionOptions {
  relay: Relay;
  peerId: string;
  /** The peer's signed Ed25519 SIGN pk extracted from the rendezvous IdPk. */
  peerSignPk: Uint8Array;
  /** Display name shown on the host's "incoming connection" notice. */
  myName: string;
  /** Peer password — empty Uint8Array for passwordless / accept-all peers. */
  password: Uint8Array;
  /** Client version string we advertise (e.g. "1.4.0"). */
  clientVersion: string;
  /** Random session id (uint64). Number value (≤ 53 bits) — opaque to peer. */
  sessionId: number;
}

export interface SessionReady {
  session: Session;
  peerInfo: hbb.IPeerInfo;
  /**
   * Messages that arrived during the login phase but weren't a Hash /
   * LoginResponse / TestDelay. The most important one in practice is
   * `Misc{audio_format}` — the host's audio_service runs its first
   * snapshot between add_connection and login_response, so AudioFormat
   * lands BEFORE the caller has a chance to set up its receive loop.
   * The caller must replay these into its post-login dispatch.
   */
  preloginExtras: hbb.Message[];
}

/** A live encrypted session. After login, use `send`/`recv` to swap Messages. */
export class Session {
  private constructor(
    private relay: Relay,
    private secretboxKey: Uint8Array,
    private sendSeq: bigint,
    private recvSeq: bigint,
  ) {}

  static async open(opts: SessionOptions): Promise<SessionReady> {
    // -------- 1. Receive Message{signed_id} --------
    const sig1Buf = await opts.relay.next();
    if (!sig1Buf) throw new Error("session: relay closed before SignedId");
    const sig1 = hbb.Message.decode(sig1Buf);
    if (!sig1.signed_id) throw new Error(`session: expected signed_id, got ${describeMessage(sig1)}`);
    const signedIdBytes = sig1.signed_id.id as Uint8Array;
    if (!signedIdBytes || signedIdBytes.length === 0) {
      throw new Error("session: SignedId.id is empty");
    }

    // -------- 2. Verify with peer's Ed25519 sign pk; extract Curve25519 box pk --------
    let peerBoxPk: Uint8Array;
    let verifiedId: string;
    try {
      const idPkBytes = signOpen(signedIdBytes, opts.peerSignPk);
      const idPk = hbb.IdPk.decode(idPkBytes);
      verifiedId = idPk.id || "";
      peerBoxPk = (idPk.pk as Uint8Array) || new Uint8Array(0);
    } catch (e) {
      throw new Error(`session: SignedId verify failed (peer's Ed25519 sign pk does not match the rendezvous-signed pk): ${e}`);
    }
    if (verifiedId !== opts.peerId) {
      throw new Error(`session: SignedId.id mismatch (expected ${opts.peerId}, got ${verifiedId})`);
    }
    if (peerBoxPk.length !== 32) {
      throw new Error(`session: peer Curve25519 box pk has wrong length ${peerBoxPk.length}`);
    }

    // -------- 3. Generate ephemeral Curve25519 keypair + symmetric key --------
    const ourBox = genBoxKeypair();
    const sbKey = genSecretboxKey();
    const sealedKey = boxSeal(sbKey, peerBoxPk, ourBox.secretKey);

    // -------- 4. Send Message{public_key} unencrypted --------
    const pkMsg = hbb.Message.create({
      public_key: hbb.PublicKey.create({
        asymmetric_value: ourBox.publicKey,
        symmetric_value: sealedKey,
      }),
    });
    opts.relay.send(hbb.Message.encode(pkMsg).finish());

    // -------- 5. From now on, all messages are secretbox-encrypted --------
    const session = new Session(opts.relay, sbKey, /*sendSeq=*/ 0n, /*recvSeq=*/ 0n);

    // -------- 6. Receive Message{hash} --------
    const hashMsg = await session.recv();
    if (!hashMsg.hash) {
      throw new Error(`session: expected hash, got ${describeMessage(hashMsg)}`);
    }
    const salt = hashMsg.hash.salt || "";
    const challenge = hashMsg.hash.challenge || "";

    // -------- 7. Compute password challenge response --------
    // The desktop client at client.rs:3414 / 3477 uses the salt and
    // challenge as RAW UTF-8 BYTES of the proto `string` fields (not
    // base64-decoded). protobufjs gives us JS strings; encode them back to
    // UTF-8 bytes the way the Rust side reads them.
    //   pwd_hash         = SHA256(password_text || salt_utf8_bytes)
    //   password_response = SHA256(pwd_hash || challenge_utf8_bytes)
    const enc = new TextEncoder();
    const saltBytes = enc.encode(salt);
    const challengeBytes = enc.encode(challenge);
    const pwdHash = sha256(concat(opts.password, saltBytes));
    const passwordResp = sha256(concat(pwdHash, challengeBytes));

    // -------- 8. Send Message{login_request} --------
    const loginReq = hbb.Message.create({
      login_request: hbb.LoginRequest.create({
        username: opts.peerId,
        password: passwordResp,
        my_id: "web-client",
        my_name: opts.myName,
        version: opts.clientVersion,
        my_platform: "Web",
        session_id: opts.sessionId,
        video_ack_required: false,
        option: hbb.OptionMessage.create({
          supported_decoding: hbb.SupportedDecoding.create({
            ability_vp9: 1,
            ability_h264: 1,
            ability_vp8: 1,
            // AV1 / H.265 left at 0 for v1.
            // Prefer VP8 — it's the cheapest software encoder of the
            // lot. Hosts without hwcodec/nvenc can keep up at 15-30 fps
            // for screen sharing, where VP9 caps at single digits and
            // H.264 falls back to "auto" (= VP9 in practice). VP8 is
            // universally supported in WebCodecs.
            prefer: hbb.SupportedDecoding.PreferCodec.VP8,
          }),
          // Explicitly opt INTO audio. The host's update_options() at
          // server/connection.rs:3974 only re-subscribes audio_service
          // when this is a definite Yes/No — `NotSet` (the default) is a
          // silent no-op and the host never starts streaming audio.
          disable_audio: hbb.OptionMessage.BoolOption.No,
          // Also explicitly enable clipboard while we're here — same gate
          // pattern, future-proofs M6h+ when we add clipboard sync.
          disable_clipboard: hbb.OptionMessage.BoolOption.No,
        }),
      }),
    });
    await session.send(loginReq);

    // -------- 9. Wait for login_response, echoing test_delay along the way --------
    // Peers send periodic TestDelay frames for latency measurement; we
    // echo them with from_client=true so the peer can compute RTT.
    // CRITICAL: the host's audio_service first-snapshot fires between
    // add_connection and login_response — Misc{audio_format} arrives here
    // and would be dropped if we just ignored non-login_response. Capture
    // every non-trivial message into preloginExtras so the caller can
    // replay them into its main receive dispatch.
    const preloginExtras: hbb.Message[] = [];
    while (true) {
      const respMsg = await session.recv();
      if (respMsg.test_delay) {
        const td = respMsg.test_delay;
        await session.send(hbb.Message.create({
          test_delay: hbb.TestDelay.create({
            time: td.time,
            from_client: true,
            last_delay: td.last_delay,
            target_bitrate: td.target_bitrate,
          }),
        }));
        continue;
      }
      if (respMsg.login_response) {
        const lr = respMsg.login_response;
        if (lr.error) throw new Error(`session: login refused: ${lr.error}`);
        if (!lr.peer_info) throw new Error("session: login_response missing peer_info");
        console.log(`[rustdesk-web] session: login OK, peer=${lr.peer_info.hostname}/${lr.peer_info.platform} v${lr.peer_info.version}`);
        return { session, peerInfo: lr.peer_info, preloginExtras };
      }
      // Stash the message for the caller to replay post-login.
      preloginExtras.push(respMsg);
    }
  }

  /** Encrypt + send a Message. The desktop's `Encrypt::enc` PRE-increments
   *  its counter — the first encrypted message uses nonce sequence = 1,
   *  not 0. See libs/hbb_common/src/tcp.rs:317-320. */
  async send(msg: hbb.IMessage): Promise<void> {
    const plain = hbb.Message.encode(msg as hbb.Message).finish();
    this.sendSeq += 1n;
    const cipher = secretboxSeal(plain, this.sendSeq, this.secretboxKey);
    this.relay.send(cipher);
  }

  /** Receive + decrypt the next Message. Same pre-increment as `send`. */
  async recv(): Promise<hbb.Message> {
    const buf = await this.relay.next();
    if (!buf) throw new Error("session: relay closed");
    this.recvSeq += 1n;
    const plain = secretboxOpen(buf, this.recvSeq, this.secretboxKey);
    return hbb.Message.decode(plain);
  }

  close(): void {
    this.relay.close();
  }
}

function describeMessage(m: hbb.IMessage): string {
  // Find the first set field for a useful diagnostic.
  for (const k of [
    "signed_id", "public_key", "test_delay", "video_frame", "login_request",
    "login_response", "hash", "mouse_event", "audio_frame", "key_event",
    "clipboard", "misc", "peer_info",
  ] as const) {
    if ((m as Record<string, unknown>)[k]) return k;
  }
  return "unknown";
}
