/**
 * Stream router — PLAN_RELAY_AGENT T8. Maps §4.1 streamId ⇄ a loopback socket. Each OPEN gets a
 * FRESH per-stream allocation (socket + transform state); there are NO global mutable buffers, so
 * cross-tenant/cross-stream buffer bleed is structurally impossible (EXPLORE §4b failure #3).
 *
 * MANDATORY INV1 defense-in-depth: the FIRST thing handleOpen does — before any allocation or
 * dial — is compare MuxOpen.subdomain (§4.1) against this agent's enrolled subdomain. A mismatch
 * is RST and NEVER dialed (metadata-only audit log, INV10). Belt-and-suspenders to relay authz.
 */
import type { MuxOpen } from 'relay-contracts'
import { MuxOpenSchema } from 'relay-contracts'
import type { AgentConfig } from '../config/agentConfig.js'
import type { Logger } from '../log/logger.js'
import type { WsLike } from './seams.js'
import type { DialLoopback } from './loopback.js'
import { closeHeader, dataHeader, type Tunnel } from './tunnel.js'

/**
 * Per-stream cipher transform. Identity in v0.9 (plaintext passthrough); replaced by the E2E
 * codec in v0.10 (T15). `takeControlFrames` lets an E2E transform emit host→client control frames
 * (e.g. HostHello) that the router forwards upstream — no-op for the identity transform.
 */
export interface FrameTransform {
  inbound(streamId: number, cipher: Uint8Array): Uint8Array | null
  outbound(streamId: number, plain: Uint8Array): Uint8Array
  openStream(streamId: number): void
  closeStream(streamId: number): void
  takeControlFrames?(streamId: number): Uint8Array[]
}

export const identityTransform: FrameTransform = {
  inbound: (_s, cipher) => cipher,
  outbound: (_s, plain) => plain,
  openStream: () => {},
  closeStream: () => {},
}

export interface StreamRouter {
  handleOpen(open: MuxOpen): void
  handleData(streamId: number, payload: Uint8Array): void
  handleClose(streamId: number, rst: boolean): void
  activeStreamCount(): number
}

interface StreamState {
  socket: WsLike | null
  readonly pending: Uint8Array[] // inbound bytes buffered until the loopback socket is open
  closed: boolean
}

export function createStreamRouter(
  cfg: AgentConfig,
  tunnel: Tunnel,
  dial: DialLoopback,
  transform: FrameTransform,
  logger: Logger,
): StreamRouter {
  const streams = new Map<number, StreamState>()

  function flushControlFrames(streamId: number): void {
    const frames = transform.takeControlFrames?.(streamId) ?? []
    for (const frame of frames) {
      tunnel.send(dataHeader(streamId, frame.length), frame)
    }
  }

  function teardown(streamId: number, rst: boolean): void {
    const state = streams.get(streamId)
    if (state === undefined) return
    // Delete FIRST so a synchronous socket 'close' event can't re-enter this teardown.
    streams.delete(streamId)
    state.closed = true
    transform.closeStream(streamId)
    tunnel.send(closeHeader(streamId, rst), new Uint8Array(0))
    state.socket?.close()
  }

  return {
    handleOpen(open: MuxOpen): void {
      // INV1 defense-in-depth — FIRST statement, before any allocation or dial.
      if (open.subdomain !== cfg.subdomain) {
        tunnel.sendRst(open.streamId)
        logger.log('error', 'open.subdomain mismatch — refusing to dial', { streamId: open.streamId })
        return
      }
      if (!MuxOpenSchema.safeParse(open).success) {
        tunnel.sendRst(open.streamId)
        return
      }
      if (streams.has(open.streamId)) {
        tunnel.sendRst(open.streamId) // duplicate OPEN for a live stream
        return
      }

      const state: StreamState = { socket: null, pending: [], closed: false }
      streams.set(open.streamId, state)
      transform.openStream(open.streamId)

      dial(open.requestPath, open.originHeader)
        .then((socket) => {
          if (state.closed) {
            socket.close()
            return
          }
          state.socket = socket
          // loopback output → transform.outbound → tunnel DATA
          socket.on('message', (data: unknown) => {
            if (!(data instanceof Uint8Array)) return
            const cipher = transform.outbound(open.streamId, data)
            tunnel.send(dataHeader(open.streamId, cipher.length), cipher)
          })
          socket.on('close', () => teardown(open.streamId, false))
          // flush anything buffered before the socket opened
          for (const buffered of state.pending) socket.send(buffered)
          state.pending.length = 0
        })
        .catch((err: unknown) => {
          logger.log('error', 'loopback dial failed', { streamId: open.streamId })
          void err
          teardown(open.streamId, true)
        })
    },

    handleData(streamId: number, payload: Uint8Array): void {
      const state = streams.get(streamId)
      if (state === undefined) {
        tunnel.sendRst(streamId) // DATA before OPEN / after CLOSE / unknown stream → RST
        return
      }
      const plain = transform.inbound(streamId, payload)
      flushControlFrames(streamId) // E2E: emit HostHello etc. (no-op in v0.9)
      if (plain === null) return // consumed (handshake), nothing to forward
      if (state.socket === null) {
        state.pending.push(plain)
        return
      }
      state.socket.send(plain)
    },

    handleClose(streamId: number, rst: boolean): void {
      teardown(streamId, rst)
    },

    activeStreamCount(): number {
      return streams.size
    },
  }
}
