/**
 * `@mizan/ws-transport` — RPC, context bundles and channel subscriptions over one socket.
 *
 * The same connection carries every direction of traffic, so an app opens one and needs
 * no HTTP fallback for reads:
 *
 *     browser -> server
 *       { action: 'rpc', id, fn, args }
 *       { action: 'ctx', id, context, params }
 *       { action: 'subscribe' | 'unsubscribe', channel, params }
 *       { action: 'message', channel, params, data }
 *
 *     server -> browser
 *       { id, ok: true,  data }                      // correlated reply
 *       { id, ok: false, error: { code, message } }
 *       { channel, params, type, data }              // a channel push, uncorrelated
 *
 * `call` resolves to the `{result, invalidate, merge}` envelope, which is what `mizanCall`
 * reads to apply server-driven invalidation — so a mutation sent over the socket
 * invalidates exactly as one sent over HTTP.
 */

import type { MizanTransport } from '@mizan/base'

type Pending = { resolve: (v: any) => void; reject: (e: any) => void }

export class MizanSocketError extends Error {
    code: string
    details?: unknown

    constructor(code: string, message: string, details?: unknown) {
        super(message)
        this.name = 'MizanSocketError'
        this.code = code
        this.details = details
    }
}

/** A push from a channel, delivered to every handler subscribed to it. */
export type ChannelMessage = {
    channel: string
    params: Record<string, any>
    type: string
    data: any
}

export type ChannelHandler = (msg: ChannelMessage) => void

export type WsTransportOptions = {
    /** Milliseconds between reconnect attempts. 0 disables reconnection. */
    retryMs?: number
    /** A reply that never arrives rejects after this many milliseconds. 0 disables. */
    timeoutMs?: number
}

export class MizanSocket {
    private url: string
    private socket: WebSocket | null = null
    private pending = new Map<string, Pending>()
    private handlers = new Map<string, Set<ChannelHandler>>()
    // re-sent on every open, so a subscription survives a reconnect
    private subscriptions = new Map<string, { channel: string; params?: Record<string, any> }>()
    private opening: Promise<void> | null = null
    private seq = 0
    private retryMs: number
    private timeoutMs: number

    constructor(url: string, opts: WsTransportOptions = {}) {
        this.url = url
        this.retryMs = opts.retryMs ?? 1000
        this.timeoutMs = opts.timeoutMs ?? 30000
    }

    private key(channel: string, params?: Record<string, any>): string {
        const parts = Object.entries(params ?? {})
            .map(([k, v]) => `${k}=${v}`)
            .sort()
        return parts.length ? `${channel}:${parts.join(':')}` : channel
    }

    /** The open socket, connecting if needed. Concurrent callers share one attempt. */
    private connect(): Promise<void> {
        if (this.socket && this.socket.readyState === WebSocket.OPEN) return Promise.resolve()
        if (this.opening) return this.opening

        this.opening = new Promise<void>((resolve, reject) => {
            const socket = new WebSocket(this.url)
            this.socket = socket

            socket.onopen = () => {
                this.opening = null
                for (const sub of this.subscriptions.values()) {
                    socket.send(JSON.stringify({ action: 'subscribe', ...sub }))
                }
                resolve()
            }
            socket.onmessage = (ev) => this.receive(ev.data)
            socket.onerror = () => {
                this.opening = null
                reject(new MizanSocketError('TRANSPORT', `websocket error on ${this.url}`))
            }
            socket.onclose = () => {
                this.opening = null
                this.socket = null
                // an in-flight call cannot be answered by a socket that is gone
                for (const [, p] of this.pending) {
                    p.reject(new MizanSocketError('TRANSPORT', 'socket closed before reply'))
                }
                this.pending.clear()
                if (this.retryMs > 0 && this.handlers.size > 0) this.retry()
            }
        })
        return this.opening
    }

    /** Reconnect for the subscriptions still held. A failed attempt is reported and retried. */
    private retry(): void {
        setTimeout(() => {
            this.connect().catch((e) => {
                console.error(`mizan socket: reconnect to ${this.url} failed`, e)
                if (this.retryMs > 0 && this.handlers.size > 0) this.retry()
            })
        }, this.retryMs)
    }

    private receive(raw: string): void {
        const msg = JSON.parse(raw)

        if (msg.id !== undefined && this.pending.has(msg.id)) {
            const p = this.pending.get(msg.id)!
            this.pending.delete(msg.id)
            if (msg.ok) p.resolve(msg.data)
            else
                p.reject(
                    new MizanSocketError(
                        msg.error?.code ?? 'INTERNAL_ERROR',
                        msg.error?.message ?? 'call failed',
                        msg.error?.details,
                    ),
                )
            return
        }

        if (msg.channel !== undefined) {
            const handlers = this.handlers.get(this.key(msg.channel, msg.params))
            if (handlers) for (const h of handlers) h(msg as ChannelMessage)
            return
        }

        // an error with no id belongs to no call — surface it rather than drop it
        if (msg.error !== undefined) console.error('mizan socket:', msg.error)
    }

    /** Send an action expecting a correlated reply. */
    private async request(action: string, body: Record<string, any>): Promise<any> {
        await this.connect()
        const id = String(++this.seq)
        return new Promise((resolve, reject) => {
            this.pending.set(id, { resolve, reject })
            if (this.timeoutMs > 0) {
                setTimeout(() => {
                    if (this.pending.delete(id)) {
                        reject(new MizanSocketError('TIMEOUT', `${action} ${id} had no reply`))
                    }
                }, this.timeoutMs)
            }
            this.socket!.send(JSON.stringify({ action, id, ...body }))
        })
    }

    /** Join a channel. The returned function leaves it. */
    async subscribe(
        channel: string,
        params: Record<string, any> | undefined,
        handler: ChannelHandler,
    ): Promise<() => void> {
        const key = this.key(channel, params)
        const first = !this.handlers.has(key)
        if (first) this.handlers.set(key, new Set())
        this.handlers.get(key)!.add(handler)

        await this.connect()
        if (first) {
            this.subscriptions.set(key, { channel, params })
            this.socket!.send(JSON.stringify({ action: 'subscribe', channel, params }))
        }

        return () => {
            const set = this.handlers.get(key)
            if (!set) return
            set.delete(handler)
            if (set.size > 0) return
            this.handlers.delete(key)
            this.subscriptions.delete(key)
            if (this.socket?.readyState === WebSocket.OPEN) {
                this.socket.send(JSON.stringify({ action: 'unsubscribe', channel, params }))
            }
        }
    }

    /** Send into a channel. Whether it reaches the group is the channel's decision. */
    async send(channel: string, params: Record<string, any> | undefined, data: any): Promise<void> {
        await this.connect()
        this.socket!.send(JSON.stringify({ action: 'message', channel, params, data }))
    }

    close(): void {
        this.retryMs = 0
        this.socket?.close()
    }

    transport(): MizanTransport {
        return {
            call: (fn, args) => this.request('rpc', { fn, args }),
            fetch: (context, params) => this.request('ctx', { context, params }),
        }
    }
}

/** The transport, plus the socket it rides, so channels are reachable from the same object. */
export function wsTransport(
    url: string,
    opts: WsTransportOptions = {},
): { transport: MizanTransport; socket: MizanSocket } {
    const socket = new MizanSocket(url, opts)
    return { transport: socket.transport(), socket }
}
