import { LOCAL_API_BASE } from 'shared/config/api';
import { TaskInfo } from 'shared/domain/librarian/task';

export class TaskWebSocket {
    private ws: WebSocket | null = null;
    private taskId: string;
    private taskType: string;
    private onUpdate: (taskInfo: TaskInfo) => void;
    private onError: (error: any) => void;

    constructor(
        taskId: string,
        taskType: string,
        onUpdate: (taskInfo: TaskInfo) => void,
        onError: (error: any) => void
    ) {
        this.taskId = taskId;
        this.taskType = taskType;
        this.onUpdate = onUpdate;
        this.onError = onError;
    }

    connect() {
        this.ws = new WebSocket(
            `ws://${LOCAL_API_BASE.replace(/^http/, 'ws')}/${this.taskType}/ws/${this.taskId}`
        );

        this.ws.onmessage = (event) => {
            try {
                const taskInfo = JSON.parse(event.data);
                this.onUpdate(taskInfo);

                if (
                    taskInfo.state === 'success' ||
                    taskInfo.state === 'error'
                ) {
                    this.disconnect();
                }
            } catch (e) {
                this.onError(new Error('Failed to parse WebSocket message'));
                this.disconnect();
            }
        };

        this.ws.onerror = (errorEvent) => {
            const error = new Error('WebSocket error');
            (error as any).event = errorEvent;
            this.onError(error);
            this.disconnect();
        };

        this.ws.onclose = (event) => {
            if (event.wasClean === false) {
                this.onError(
                    new Error(
                        `WebSocket closed unexpectedly: ${event.code} ${event.reason || ''}`.trim()
                    )
                );
            }
            this.ws = null;
        };
    }

    disconnect() {
        if (this.ws) {
            this.ws.close();
            this.ws = null;
        }
    }
}