import { useEcho } from '@laravel/echo-vue';
import { onBeforeUnmount, ref, shallowRef } from 'vue';
import type { Iteration } from '@/types/perceptron';

type TrainingEvent = {
    trainingId: string;
};

type IterationEvent = TrainingEvent & {
    iterations: Iteration[];
};

type InitializationEvent = TrainingEvent & {
    activationFunction: string;
};

type TrainingEndedEvent = TrainingEvent & {
    reason: string;
};

export function usePerceptronTraining(sessionId: string) {
    const trainingId = ref('');
    const iterations = shallowRef<Iteration[]>([]);
    const trainingEnded = ref(false);
    const trainingEndReason = ref('');
    const activationFunction = ref('');

    let pendingIterations: Iteration[] = [];
    let renderFrame: number | null = null;

    function isCurrentTraining(event: TrainingEvent): boolean {
        return event.trainingId === trainingId.value;
    }

    function flushPendingIterations(): void {
        renderFrame = null;
        if (pendingIterations.length === 0) {
            return;
        }

        iterations.value = [...iterations.value, ...pendingIterations];
        pendingIterations = [];
    }

    function handleIterations(event: IterationEvent): void {
        if (!isCurrentTraining(event)) {
            return;
        }

        pendingIterations.push(...event.iterations);
        if (renderFrame === null) {
            renderFrame = requestAnimationFrame(flushPendingIterations);
        }
    }

    function handleInitialization(event: InitializationEvent): void {
        if (isCurrentTraining(event)) {
            activationFunction.value = event.activationFunction;
        }
    }

    function handleTrainingEnded(event: TrainingEndedEvent): void {
        if (!isCurrentTraining(event)) {
            return;
        }

        flushPendingIterations();
        trainingEnded.value = true;
        trainingEndReason.value = event.reason;
    }

    function reset(): void {
        if (renderFrame !== null) {
            cancelAnimationFrame(renderFrame);
            renderFrame = null;
        }

        pendingIterations = [];
        iterations.value = [];
        trainingEnded.value = false;
        trainingEndReason.value = '';
        activationFunction.value = '';
    }

    function setTrainingId(newTrainingId: string): void {
        reset();
        trainingId.value = newTrainingId;
    }

    const channel = `${sessionId}-perceptron-training`;
    useEcho(channel, 'PerceptronTrainingIteration', handleIterations, [{}], 'public');
    useEcho(channel, 'PerceptronTrainingEnded', handleTrainingEnded, [{}], 'public');
    useEcho(channel, 'PerceptronInitialization', handleInitialization, [{}], 'public');

    onBeforeUnmount(reset);

    return {
        activationFunction,
        iterations,
        setTrainingId,
        trainingEnded,
        trainingEndReason,
    };
}