// A MOQT draft-22 client session over a transport of ./transport.js, the codec being
// joule-moq-web's wasm build.
//
// - SETUP on this end's control stream (its first unidirectional stream), with joule-moq's
//   receipt extension, and either a room grant with its proof of possession or a token by
//   value. A room grant is public, so it is presented only with a proof (WAI live-rooms
//   §10.5.2): the session waits for the relay's SETUP, which carries the session's
//   challenge, and sends its own after it, the grant as its authorization token and a proof
//   by the principal's key (signed in the wasm build) over the challenge and the
//   `//authority/path` of the URL the transport opened. The relay's SETUP is read from its
//   control stream, data streams that come first held until it.
// - Requests on bidirectional streams, Request IDs even from 0 (§6.4.2): PUBLISH,
//   PUBLISH_NAMESPACE, SUBSCRIBE (with priorities, filters and FILL_PARAMETERS), FETCH,
//   SUBSCRIBE_NAMESPACE, TRACK_STATUS.
// - Every data stream read concurrently, each by a task of its own as it arrives: a
//   subscriber that reads its streams one at a time, in the order they arrive, can stall
//   itself, since the stream it waits on may have its bytes held back for flow-control
//   credit that the unread streams behind it hold (joule-moq-relay's README, "A subscriber
//   that does not read").
// - GOAWAY (§9.2): no new request is opened once one arrives; the caller is told the new
//   session URI and the timeout, and moves.

import { StreamReset } from "./transport.js";

const STEP_MS = 10000;
const ALIAS_WAIT_MS = 3000;
// A Track Alias a subgroup stream names before this end has read the SUBSCRIBE_OK that
// gives it waits up to ALIAS_WAIT_MS.

export function withTimeout(p, ms, what) {
  let t;
  return Promise.race([p, new Promise((_, no) => { t = setTimeout(() => no(new Error(`${what}: no answer in ${ms} ms`)), ms); })]).finally(() => clearTimeout(t));
}

const hexToBytes = (h) => new Uint8Array((h.match(/../g) || []).map((x) => parseInt(x, 16)));
const utf8 = (h) => new TextDecoder().decode(hexToBytes(h));

export class RequestRefused extends Error {
  constructor(message, first) {
    super(message);
    this.first = first;
  }
}

class Queue {
  constructor() { this.items = []; this.waiters = []; this.done = false; }
  push(x) { const w = this.waiters.shift(); if (w) w(x); else this.items.push(x); }
  end() { this.done = true; for (const w of this.waiters.splice(0)) w(null); }
  next() {
    if (this.items.length) return Promise.resolve(this.items.shift());
    if (this.done) return Promise.resolve(null);
    return new Promise((r) => this.waiters.push(r));
  }
}

export class MoqSession {
  constructor(transport, moq, { token, grant, implementation = "joule-moq-web", log = () => {}, maxPayload = 4 << 20, properties = false } = {}) {
    this.t = transport;
    // Each object's property list, encoded, as its fourth argument to a sink: what a
    // member of an end-to-end encrypted room opens an object with.
    this.properties = properties;
    this.moq = moq;
    this.token = token;
    // { record, seed }: a room grant's record and its principal's 32-byte seed.
    this.grant = grant;
    this.implementation = implementation;
    this.log = log;
    this.maxPayload = maxPayload;
    this.nextRequestId = 0;
    this.tracks = new Map(); // downstream Track Alias -> handler { object, end }
    this.aliasWaiters = new Map();
    this.fetches = new Map(); // request id -> fetch handler
    this.goaway = null;
    this.onGoaway = null;
    this.closedInfo = null;
    this.counters = { uni_streams: 0, subgroups: 0, fetch_streams: 0, objects: 0, bytes: 0, resets: 0, unrouted: 0, decode_errors: 0, streams_open_max: 0 };
    this.openStreams = 0;
    this.peerSetup = null;
    this.closed = transport.closed.then((i) => { this.closedInfo = i; return i; });
  }

  static async connect(transport, moq, opts) {
    const s = new MoqSession(transport, moq, opts);
    await s.start();
    return s;
  }

  async start() {
    this.setupReceived = new Promise((ok, no) => { this._setupOk = ok; this._setupNo = no; });
    this.acceptUniLoop();
    this.acceptBidiLoop();
    this.control = await this.t.openUni();
    if (!this.grant) {
      await this.control.write(this.moq.clientSetup(this.token || undefined, this.implementation, true, 0n));
      this.peerSetup = await withTimeout(this.setupReceived, STEP_MS, "the relay's SETUP");
      this.log(`session: SETUP exchanged over ${this.t.kind}`);
      return;
    }
    this.peerSetup = await withTimeout(this.setupReceived, STEP_MS, "the relay's SETUP");
    const type = Number(this.moq.grantChallengeOption());
    const opt = (this.peerSetup.options || []).find((o) => o.type === type && o.bytes);
    if (!opt) throw new Error("the relay's SETUP carries no challenge for a grant's proof");
    const u = new URL(this.t.url);
    this.target = `//${u.host}${u.pathname}`;
    const challenge = hexToBytes(opt.bytes);
    // `proofBy`: a test's proof by another key than the principal's, which the relay refuses.
    const setup = this.grant.proofBy
      ? this.moq.clientSetupWithProof(this.grant.record, this.moq.grantProofBy(this.grant.record, this.grant.proofBy, challenge, this.target), this.implementation, true, 0n)
      : this.moq.clientSetupWithGrant(this.grant.record, this.grant.seed, challenge, this.target, this.implementation, true, 0n);
    await this.control.write(setup);
    this.log(`session: SETUP exchanged over ${this.t.kind}, the grant presented with its proof for ${this.target}`);
  }

  // ------------------------------------------------------------------- incoming streams

  async acceptUniLoop() {
    for (;;) {
      const r = await this.t.acceptUni();
      if (!r) return;
      this.counters.uni_streams++;
      // Each stream is read by a task of its own, at once.
      this.readUni(r);
    }
  }

  async acceptBidiLoop() {
    for (;;) {
      const s = await this.t.acceptBidi();
      if (!s) return;
      // This client publishes by PUBLISH, so the relay opens no request to it; one it
      // opens anyway is refused (REQUEST_ERROR NOT_SUPPORTED).
      this.refuseIncoming(s);
    }
  }

  async refuseIncoming({ writer, reader }) {
    const ms = new this.moq.MessageStream();
    try {
      for (;;) {
        const { value, done } = await reader.read();
        if (done) break;
        ms.push(value);
        const m = ms.next();
        if (m) {
          await writer.write(this.moq.encodeControl(JSON.stringify({ message: "REQUEST_ERROR", code: 3, reason: "this client takes no requests" })));
          await writer.close();
          break;
        }
      }
    } catch { /* the stream ended */ }
    ms.free();
  }

  async readUni(reader) {
    const dec = new this.moq.UniDecoder(BigInt(this.maxPayload));
    let sink = null;
    let kind = null;
    this.openStreams++;
    this.counters.streams_open_max = Math.max(this.counters.streams_open_max, this.openStreams);
    try {
      for (;;) {
        let chunk;
        try {
          chunk = await reader.read();
        } catch (e) {
          this.counters.resets++;
          if (sink && sink.reset) sink.reset(e instanceof StreamReset ? e.code : -1);
          return;
        }
        if (chunk.done) {
          try { dec.finish(); } catch (e) { this.counters.decode_errors++; this.log(`a stream ended inside an item: ${e.message || e}`); }
          if (sink && sink.end) sink.end();
          return;
        }
        this.counters.bytes += chunk.value.length;
        dec.push(chunk.value);
        for (;;) {
          let item;
          try { item = dec.next(); } catch (e) {
            this.counters.decode_errors++;
            this.log(`a data stream does not decode: ${e.message || e}`);
            reader.cancel(0);
            if (sink && sink.reset) sink.reset(-2);
            return;
          }
          if (!item) break;
          const k = item.kind;
          const j = JSON.parse(item.json);
          if (k === "header") {
            kind = j.header;
            sink = await this.route(j, dec);
            if (!sink) { this.counters.unrouted++; reader.cancel(0); item.free(); return; }
          } else if (k === "control") {
            this.onControl(j.control);
          } else if (k === "object" || k === "fetch_object") {
            this.counters.objects++;
            const payload = item.payload();
            const props = this.properties ? item.properties() : null;
            sink.object && sink.object(j, payload, k, props);
          } else if (k === "end_of_range") {
            sink.endOfRange && sink.endOfRange(j.end_of_range);
          }
          item.free();
        }
      }
    } finally {
      this.openStreams--;
      dec.free();
      void kind;
    }
  }

  async route(h, dec) {
    if (h.header === "control") {
      return { object() {}, end: () => { if (!this.closedInfo) this.log("the relay's control stream ended"); } };
    }
    if (h.header === "subgroup") {
      this.counters.subgroups++;
      let t = this.tracks.get(h.track_alias);
      if (!t) t = await this.waitAlias(h.track_alias);
      if (!t) return null;
      return t.stream(h);
    }
    if (h.header === "fetch") {
      this.counters.fetch_streams++;
      const f = this.fetches.get(h.request_id);
      if (!f) return null;
      dec.setGroupOrder(f.order || 1);
      return f.stream(h);
    }
    return null;
  }

  waitAlias(alias) {
    return new Promise((r) => {
      const t = setTimeout(() => { done(null); }, ALIAS_WAIT_MS);
      const done = (x) => {
        clearTimeout(t);
        const ws = (this.aliasWaiters.get(alias) || []).filter((w) => w !== done);
        if (ws.length) this.aliasWaiters.set(alias, ws); else this.aliasWaiters.delete(alias);
        r(x);
      };
      this.aliasWaiters.set(alias, [...(this.aliasWaiters.get(alias) || []), done]);
    });
  }

  setTrack(alias, handler) {
    this.tracks.set(alias, handler);
    for (const w of this.aliasWaiters.get(alias) || []) w(handler);
  }

  onControl(m) {
    if (m.message === "SETUP") {
      this._setupOk(m);
    } else if (m.message === "GOAWAY") {
      this.goaway = { uri: utf8(m.new_session_uri), timeout_ms: m.timeout_ms };
      this.log(`GOAWAY: new session URI "${this.goaway.uri}", timeout ${this.goaway.timeout_ms} ms`);
      if (this.onGoaway) this.onGoaway(this.goaway);
    }
  }

  // ---------------------------------------------------------------------- requests

  async request(desc) {
    if (this.goaway) throw new Error(`GOAWAY received: no new ${desc.message}`);
    if (this.closedInfo) throw new Error("the session is closed");
    const id = this.nextRequestId;
    this.nextRequestId += 2;
    // The authorization token rides in SETUP, where the relay's grant rules read it for
    // every request of the session.
    const d = { ...desc, request_id: id };
    const { writer, reader } = await this.t.openBidi();
    const bytes = this.moq.encodeControl(JSON.stringify(d));
    await writer.write(bytes);
    const q = new Queue();
    const ms = new this.moq.MessageStream();
    (async () => {
      try {
        for (;;) {
          const { value, done } = await reader.read();
          if (done) break;
          ms.push(value);
          for (let m; (m = ms.next()); ) q.push(JSON.parse(m));
        }
      } catch (e) {
        q.push({ message: "STREAM_RESET", code: e.code });
      } finally {
        ms.free();
        q.end();
      }
    })();
    const first = await withTimeout(q.next(), STEP_MS, desc.message);
    if (!first) throw new RequestRefused(`${desc.message}: the stream ended before its answer`, null);
    return { id, first, messages: q, writer, reader };
  }

  static refusal(what, first) {
    if (first.message === "REQUEST_ERROR") return new RequestRefused(`${what}: REQUEST_ERROR ${first.code} ${utf8(first.reason || "")}`, first);
    return new RequestRefused(`${what}: an unexpected ${first.message}`, first);
  }

  async publish(namespace, name, alias, trackProperties) {
    const r = await this.request({ message: "PUBLISH", namespace, name, track_alias: alias, track_properties: trackProperties });
    if (r.first.message !== "REQUEST_OK") throw MoqSession.refusal(`PUBLISH ${namespace.join("/")}--${name}`, r.first);
    return r;
  }

  // SUBSCRIBE: `handler` gets each subgroup stream of the track (`stream(header)` gives
  // its object sink) and, with `params.fill`, the fill fetch stream named by the request.
  async subscribe(namespace, name, params, handler) {
    const id = this.nextRequestId;
    if (params && params.fill) this.fetches.set(id, { order: params.group_order || 1, stream: (h) => handler.fill(h) });
    let r;
    try {
      r = await this.request({ message: "SUBSCRIBE", namespace, name, params });
    } catch (e) {
      this.fetches.delete(id);
      throw e;
    }
    if (r.first.message !== "SUBSCRIBE_OK") {
      this.fetches.delete(id);
      throw MoqSession.refusal(`SUBSCRIBE ${namespace.join("/")}--${name}`, r.first);
    }
    const alias = r.first.track_alias;
    this.setTrack(alias, handler);
    (async () => {
      for (let m; (m = await r.messages.next()); ) {
        if (m.message === "PUBLISH_DONE") { handler.done && handler.done(m); }
      }
    })();
    return { ...r, alias };
  }

  // FETCH: every entry of the fetch stream, until its FIN.
  async fetch(namespace, name, params, onObject) {
    const id = this.nextRequestId;
    let finished;
    const ended = new Promise((r) => { finished = r; });
    const entries = { objects: 0, ranges: [] };
    this.fetches.set(id, {
      order: (params && params.group_order) || 1,
      stream: () => ({
        object: (j, payload, k, props) => { entries.objects++; onObject(j, payload, props); },
        endOfRange: (e) => entries.ranges.push(e),
        end: () => finished("fin"),
        reset: (code) => finished(`reset ${code}`),
      }),
    });
    let r;
    try {
      r = await this.request({ message: "FETCH", namespace, name, params });
    } catch (e) {
      this.fetches.delete(id);
      throw e;
    }
    if (r.first.message !== "FETCH_OK") {
      this.fetches.delete(id);
      throw MoqSession.refusal(`FETCH ${namespace.join("/")}--${name}`, r.first);
    }
    const how = await withTimeout(ended, STEP_MS * 2, "the fetch stream");
    this.fetches.delete(id);
    return { ok: r.first, how, ...entries };
  }

  // One subgroup stream of `bytes`, at `sendOrder` where the transport has it.
  async sendStream(bytes, sendOrder) {
    const w = await this.t.openUni({ sendOrder });
    await w.write(bytes);
    await w.close();
  }

  async openStream(sendOrder) {
    return this.t.openUni({ sendOrder });
  }

  close(code = 0, reason = "") {
    this.t.close(code, reason);
  }
}