import { useChat, type Message } from 'ai/react';
import { useAtom } from 'jotai';
import { toast } from 'react-hot-toast';
import { datasetAtom } from '../../state';
import { useEffect, useState } from 'react';
import { MessageWithSelectedDataset } from '../types';

const useChatWithDataset = () => {
	const [dataset] = useAtom(datasetAtom);
	const {
		messages,
		append,
		reload,
		stop,
		isLoading,
		input,
		setInput,
		setMessages,
	} = useChat({
		sendExtraMessageFields: true,
		onResponse(response) {
			if (response.status !== 200) {
				toast.error(response.statusText);
			}
		},
	});

	const [loadingDots, setLoadingDots] = useState('');

	useEffect(() => {
		let loadingInterval: NodeJS.Timeout;

		if (isLoading) {
			loadingInterval = setInterval(() => {
				setLoadingDots(prevMessage => {
					switch (prevMessage) {
						case '':
							return '.';
						case '.':
							return '..';
						case '..':
							return '...';
						case '...':
							return '';
						default:
							return '';
					}
				});
			}, 500);
		}

		return () => {
			clearInterval(loadingInterval);
		};
	}, [isLoading]);

	const assistantLoadingMessage = {
		id: 'loading',
		content: loadingDots,
		role: 'assistant',
	};

	const messageWithLoading =
		isLoading &&
		messages.length &&
		messages[messages.length - 1].role !== 'assistant'
			? [...messages, assistantLoadingMessage]
			: messages;

	const selectedDataset = dataset.find(entity => entity.selected)
		? dataset.filter(entity => entity.selected)
		: // If there is no selected dataset, use the entire dataset
			dataset;

	const appendWithDataset: typeof append = message => {
		const newSystemMessage: Message = {
			id: 'fake-id',
			content:
				'For the next prompt, here are names of images provided by user, please use these name if you need reference: ' +
				selectedDataset.map(entity => entity.name).join(', '),
			role: 'system',
		};
		setMessages([...messages, newSystemMessage]);
		return append({
			...message,
			// @ts-ignore this is extra fields
			dataset: selectedDataset,
		} satisfies MessageWithSelectedDataset);
	};

	return {
		messages: messageWithLoading as MessageWithSelectedDataset[],
		append: appendWithDataset,
		reload,
		stop,
		isLoading,
		input,
		setInput,
	};
};

export default useChatWithDataset;