import type { ChannelId, ServerId, UserId } from "~/lib/api/types";
import { createPrefixedLogger } from "~/lib/utils";
import {
    type ClientMessage,
    ClientMessageType,
    ConnectionState,
    ErrorCode,
    type EventData,
    EventType,
    type ServerMessage,
    ServerMessageType,
} from "./types";

export type GatewayEvents = {
    [K in EventType]: (data: Extract<EventData, { type: K }>["data"]) => void;
};

export type ControlEvents = {
    stateChange: (state: ConnectionState) => void;
    error: (error: Error, code?: ErrorCode) => void;
    authenticated: (userId: UserId, sessionKey: string) => void;
};

export interface GatewayClientOptions {
    reconnect?: boolean;
    reconnectDelay?: number;
    maxReconnectAttempts?: number;
}

export class GatewayClient {
    private socket: WebSocket | null = null;
    private state: ConnectionState = ConnectionState.DISCONNECTED;
    private url: string;
    private token: string | null = null;
    private sessionKey: string | null = null;
    private userId: string | null = null;
    private reconnectAttempts = 0;
    private reconnectTimeout: NodeJS.Timeout | null = null;
    private eventHandlers: Partial<ControlEvents> = {};
    private serverEventHandlers: Partial<GatewayEvents> = {};
    private options: Required<GatewayClientOptions>;
    private closeInitiatedByClient = false;

    private connectionLock = false;

    constructor(url: string, options: GatewayClientOptions = {}) {
        this.url = url;
        this.options = {
            reconnect: options.reconnect ?? true,
            reconnectDelay: options.reconnectDelay ?? 5000,
            maxReconnectAttempts: options.maxReconnectAttempts ?? 10,
        };
    }

    // Public methods
    public connect(token: string): void {
        logger.log("Connecting to %s", this.url);

        if (this.connectionLock) {
            logger.warn("Connection already in progress");
            return;
        }

        if (this.token === token) {
            logger.warn("Token is the same as the current token");
            return;
        }

        this.connectionLock = true;

        if (this.state !== ConnectionState.DISCONNECTED) {
            this.disconnect();
        }

        this.token = token;
        this.closeInitiatedByClient = false;
        this.reconnectAttempts = 0;
        this.connectToWebSocket();
    }

    public disconnect(): void {
        logger.log("Disconnecting");

        this.closeInitiatedByClient = true;
        this.cleanupSocket();
        this.setState(ConnectionState.DISCONNECTED);

        this.connectionLock = false;
    }

    public updateVoiceState(serverId: ServerId, channelId: ChannelId): void {
        this.sendMessage({
            type: ClientMessageType.VOICE_STATE_UPDATE,
            data: { serverId, channelId },
        });
    }

    public requestVoiceStates(serverId: ServerId): void {
        this.sendMessage({
            type: ClientMessageType.REQUEST_VOICE_STATES,
            data: { serverId },
        });
    }

    public onEvent<K extends keyof GatewayEvents>(event: K | string, handler: GatewayEvents[K]): void {
        this.serverEventHandlers[event as K] = handler;
    }

    public offEvent<K extends keyof GatewayEvents>(event: K): void {
        delete this.serverEventHandlers[event];
    }

    public onControl<K extends keyof ControlEvents>(event: K, handler: ControlEvents[K]): void {
        this.eventHandlers[event] = handler;
    }

    public offControl<K extends keyof ControlEvents>(event: K): void {
        delete this.eventHandlers[event];
    }

    public get connectionState(): ConnectionState {
        return this.state;
    }

    public get currentUserId(): UserId | null {
        return this.userId;
    }

    public get currentSessionKey(): string | null {
        return this.sessionKey;
    }

    // Private methods
    private connectToWebSocket(): void {
        this.setState(ConnectionState.CONNECTING);

        try {
            this.socket = new WebSocket(this.url);

            this.socket.onopen = this.onSocketOpen.bind(this);
            this.socket.onmessage = this.onSocketMessage.bind(this);
            this.socket.onerror = this.onSocketError.bind(this);
            this.socket.onclose = this.onSocketClose.bind(this);
        } catch (error) {
            this.emitError(new Error("Failed to create WebSocket connection"));
            this.setState(ConnectionState.ERROR);
        }
    }

    private onSocketOpen(): void {
        this.connectionLock = false;

        logger.log("Socket opened");
        this.setState(ConnectionState.AUTHENTICATING);

        if (this.token) {
            this.sendMessage({
                type: ClientMessageType.AUTHENTICATE,
                data: { token: this.token },
            });
        } else {
            this.emitError(new Error("No authentication token provided"));
            this.disconnect();
        }
    }

    private onSocketMessage(event: MessageEvent): void {
        try {
            const message = JSON.parse(event.data) as ServerMessage;
            this.handleServerMessage(message);
        } catch (error) {
            this.emitError(
                new Error("Failed to parse WebSocket message", {
                    cause: error,
                }),
            );
        }
    }

    private onSocketError(event: Event): void {
        this.connectionLock = false;
        logger.log("Socket error: %s", event);

        this.emitError(new Error("WebSocket error occurred"));
    }

    private onSocketClose(event: CloseEvent): void {
        logger.log("Socket closed: %s", event);

        this.connectionLock = false;

        if (
            this.options.reconnect &&
            !this.closeInitiatedByClient &&
            this.reconnectAttempts < this.options.maxReconnectAttempts
        ) {
            logger.log(
                "Reconnecting in %d seconds (%d/%d)",
                this.options.reconnectDelay / 1000,
                this.reconnectAttempts + 1,
                this.options.maxReconnectAttempts,
            );
            this.reconnectAttempts++;

            this.reconnectTimeout = setTimeout(() => {
                if (this.token) {
                    this.connectToWebSocket();
                }
            }, this.options.reconnectDelay);
        } else {
            this.setState(ConnectionState.DISCONNECTED);
        }
    }

    private handleServerMessage(message: ServerMessage): void {
        logger.log("Received message: ", message);

        switch (message.type) {
            case ServerMessageType.AUTHENTICATE_ACCEPTED:
                this.userId = message.data.userId;
                this.sessionKey = message.data.sessionKey;
                this.setState(ConnectionState.CONNECTED);
                this.emitControl("authenticated", message.data.userId, message.data.sessionKey);
                break;

            case ServerMessageType.AUTHENTICATE_DENIED:
                this.emitError(new Error("Authentication denied"));
                this.disconnect();
                break;

            case ServerMessageType.ERROR:
                this.emitError(new Error(`Server error: ${message.data.code}`), message.data.code);
                break;

            case ServerMessageType.EVENT:
                this.handleEventMessage(message.data.event);
                break;

            default:
                console.warn("Unhandled server message type:", message);
        }
    }

    private handleEventMessage(event: EventData): void {
        this.emitEvent(event.type, event.data as any);
    }

    private sendMessage(message: ClientMessage): void {
        logger.log("Sending message: %o", message);

        if (this.socket && this.socket.readyState === WebSocket.OPEN) {
            this.socket.send(JSON.stringify(message));
        } else {
            this.emitError(new Error("Cannot send message: socket not connected"));
        }
    }

    private setState(state: ConnectionState): void {
        if (this.state !== state) {
            logger.log("State changed to %s", state);

            this.state = state;
            this.emitControl("stateChange", state);
        }
    }

    private emitError(error: Error, code?: ErrorCode): void {
        logger.error("Error: %s", error, error.cause);

        this.setState(ConnectionState.ERROR);
        this.emitControl("error", error, code);
    }

    private emitControl<K extends keyof ControlEvents>(event: K, ...args: Parameters<ControlEvents[K]>): void {
        const handler = this.eventHandlers[event];
        if (handler) {
            (handler as Function)(...args);
        }
    }

    private emitEvent<K extends keyof GatewayEvents>(event: K, ...args: Parameters<GatewayEvents[K]>): void {
        const handler = this.serverEventHandlers[event];
        if (handler) {
            (handler as Function)(...args);
        }
    }

    private cleanupSocket(): void {
        logger.log("Cleaning up socket");

        if (this.reconnectTimeout) {
            clearTimeout(this.reconnectTimeout);
            this.reconnectTimeout = null;
        }

        if (this.socket) {
            // Remove all event listeners
            this.socket.onopen = null;
            this.socket.onmessage = null;
            this.socket.onerror = null;
            this.socket.onclose = null;

            // Close the connection if it's still open
            if (this.socket.readyState === WebSocket.OPEN || this.socket.readyState === WebSocket.CONNECTING) {
                this.socket.close();
            }

            this.socket = null;
        }

        this.sessionKey = null;
        this.userId = null;
    }
}

const logger = createPrefixedLogger("%cGateway WS%c:", ["color: red; font-weight: bold;", ""]);