export type SocketMessage = {
    type: string;
    data: any;
}
export type CompatEvent = {
    request: {
        headers: Headers;
    };
    context: any;
} | {
    headers: Headers;
    context: any;
};

export class Socket
{
    private _frequency?: number;
    private _timeout?: number;
    private _heartbeat: boolean = false;
    private _heartbeatWaiting = false;
    private _timeoutTimer?: NodeJS.Timeout;

    private _ws: WebSocket;
    private _handlers: Map<string, (data: any) => void> = new Map();

    constructor(url: string, heartbeat?: true | { timeout?: number, frequency?: number })
    {
        this._frequency = heartbeat === true ? 10000 : heartbeat?.frequency ?? 10000;
        this._timeout = heartbeat === true ? 100000 : heartbeat?.timeout ?? 100000;
        this._heartbeat = heartbeat !== undefined;

        this._ws = new WebSocket(`${location.protocol.endsWith('s:') ? 'wss' : 'ws'}://${location.host}${url.startsWith('/') ? url : '/' + url}`);

        this._ws.addEventListener('open', (e) => {
            console.log(`[ws] Connected to ${this._ws.url}`);
            this._heartbeat && setTimeout(() => this.heartbeat(), this._frequency);
        });
        this._ws.addEventListener('close', (e) => {
            console.log(`[ws] Disconnected from ${this._ws.url} (code: ${e.code}, reason: ${e.reason}, ${e.wasClean ? 'clean close' : 'dirty close'})`)
            this._heartbeatWaiting = false;
            this._timeoutTimer && clearTimeout(this._timeoutTimer);
        });
        this._ws.addEventListener('message', (e) => {
            const data = JSON.parse(e.data) as SocketMessage;

            switch(data.type)
            {
                case 'PONG':
                    if(this._heartbeatWaiting)
                    {
                        this._heartbeatWaiting = false;
                        this._timeoutTimer && clearTimeout(this._timeoutTimer);
                        this._heartbeat && setTimeout(() => this.heartbeat(), this._frequency);
                    }
                    return;
                default: return this._handlers.has(data.type) && queueMicrotask(() => this._handlers.get(data.type)!(data.data));
            }
        });
    }
    public handleMessage<T>(type: string, callback: (data: T) => void)
    {
        this._handlers.set(type, callback);
    }
    public send(type: string, data: any)
    {
        this._ws.readyState === WebSocket.OPEN && this._ws.send(JSON.stringify({ type, data }));
    }
    public close()
    {
        this._ws.close(1000);
    }
    private heartbeat()
    {
        if(this._heartbeat && this._ws.readyState === WebSocket.OPEN)
        {
            this._ws.send(JSON.stringify({ type: 'PING' }));
            this._heartbeatWaiting = true;
            this._timeoutTimer = setTimeout(() => this._ws.close(3000, 'Timeout'), this._timeout);
        }
    }
}