import type { MutableRefObject, PropsWithChildren } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { v4 as uuidv4 } from "uuid";

import * as api from "../api";
import { getTopicStorageKey, localStorageSessionIdKey } from "../constants";
import { chatEventBus } from "../event-buses";
import type {
	ChatMessage,
	ChatMessageText,
	ChatOptions,
	ChatStore,
	ChatTopic,
	FeedbackVote,
	SendMessageResponse,
} from "../types";
import { ChatContext } from "../context/ChatContext";
import { ChatOptionsContext } from "../context/ChatOptionsContext";
import { StreamingMessageManager, createBotMessage } from "../utils/streaming";
import {
	handleNodeComplete,
	handleNodeStart,
	handleStreamingChunk,
} from "../utils/streamingHandlers";

interface ChatProviderProps extends PropsWithChildren {
	options: ChatOptions;
}

function emitScrollToBottom() {
	if (typeof queueMicrotask === "function") {
		queueMicrotask(() => chatEventBus.emit("scrollToBottom"));
		return;
	}

	setTimeout(() => chatEventBus.emit("scrollToBottom"), 0);
}

function createUserMessage(text: string, files: File[] = []): ChatMessage {
	return {
		id: uuidv4(),
		text,
		sender: "user",
		files,
	};
}

function handleEmptyStreamResponse(
	receivedMessage: MutableRefObject<ChatMessageText | null>,
	appendMessage: (message: ChatMessage) => void,
	replaceMessage: (messageId: string, message: ChatMessageText) => void,
	messagesRef: React.MutableRefObject<ChatMessage[]>,
) {
	if (!receivedMessage.current) {
		receivedMessage.current = createBotMessage();
		appendMessage(receivedMessage.current);
	} else {
		const hasContent = messagesRef.current.some(
			(msg) =>
				msg.sender === "bot" &&
				"text" in msg &&
				(msg as ChatMessageText).text?.trim().length > 0,
		);

		if (!hasContent) {
			receivedMessage.current = createBotMessage();
			appendMessage(receivedMessage.current);
		}
	}

	const updatedMessage: ChatMessageText = {
		...receivedMessage.current,
		text: "[No response received. This could happen if streaming is enabled in the trigger but disabled in agent node(s)]",
		canReceiveFeedback: true,
	};

	replaceMessage(updatedMessage.id, updatedMessage);
	receivedMessage.current = updatedMessage;
}

function handleMessageError(
	error: unknown,
	receivedMessage: MutableRefObject<ChatMessageText | null>,
	appendMessage: (message: ChatMessage) => void,
	replaceMessage: (messageId: string, message: ChatMessageText) => void,
) {
	const fallbackMessage = receivedMessage.current ?? createBotMessage();
	const updatedMessage: ChatMessageText = {
		...fallbackMessage,
		text: "Error: Failed to receive response",
		canReceiveFeedback: false,
	};

	if (!receivedMessage.current) {
		appendMessage(updatedMessage);
	} else {
		replaceMessage(updatedMessage.id, updatedMessage);
	}

	receivedMessage.current = updatedMessage;
	console.error("Chat API error:", error);
}

async function handleNonStreamingMessage(
	text: string,
	files: File[],
	sessionId: string,
	options: ChatOptions,
): Promise<{ response?: SendMessageResponse; botMessage?: ChatMessageText }> {
	const sendMessageResponse = await api.sendMessage(
		text,
		files,
		sessionId,
		options,
	);

	if (sendMessageResponse?.executionStarted) {
		return { response: sendMessageResponse };
	}

	const receivedMessage = createBotMessage();
	const textMessage =
		sendMessageResponse.output ??
		sendMessageResponse.text ??
		sendMessageResponse.message ??
		"";

	if (textMessage) {
		receivedMessage.text = textMessage;
		receivedMessage.canReceiveFeedback = true;
	} else if (Object.keys(sendMessageResponse).length > 0) {
		try {
			receivedMessage.text = JSON.stringify(sendMessageResponse, null, 2);
			receivedMessage.canReceiveFeedback = true;
		} catch {
			receivedMessage.text = "";
		}
	}

	return { botMessage: receivedMessage };
}

export function ChatProvider({ options, children }: ChatProviderProps) {
	const [messages, setMessages] = useState<ChatMessage[]>([]);
	const [waitingForResponse, setWaitingForResponse] = useState(false);
	const [currentSessionId, setCurrentSessionId] = useState<string | null>(
		null,
	);
	const [selectedTopic, setSelectedTopic] = useState<ChatTopic | null>(null);
	const availableTopics = options.topics ?? [];
	const websocketRef = useRef<WebSocket | null>(null);
	const messagesRef = useRef<ChatMessage[]>([]);
	const topicSelectionLockedRef = useRef(false);

	const persistTopicSlug = useCallback(
		(sessionId: string, slug: string | null) => {
			if (!sessionId) {
				return;
			}

			try {
				const storageKey = getTopicStorageKey(sessionId);
				if (!slug) {
					localStorage.removeItem(storageKey);
					return;
				}

				localStorage.setItem(storageKey, slug);
			} catch (error) {
				// LocalStorage is best-effort; ignore failures.
			}
		},
		[],
	);

	const resolveTopicBySlug = useCallback(
		(slug?: string | null): ChatTopic | null => {
			if (!slug) {
				return null;
			}

			return availableTopics.find((topic) => topic.slug === slug) ?? null;
		},
		[availableTopics],
	);

	const readTopicFromStorage = useCallback(
		(
			sessionId: string,
		): { topic: ChatTopic | null; slug: string | null } => {
			if (!sessionId) {
				return { topic: null, slug: null };
			}

			try {
				const storedSlug = localStorage.getItem(
					getTopicStorageKey(sessionId),
				);
				return {
					topic: resolveTopicBySlug(storedSlug),
					slug: storedSlug,
				};
			} catch (error) {
				return { topic: null, slug: null };
			}
		},
		[resolveTopicBySlug],
	);

	const buildOptionsWithTopic = useCallback(
		(topicOverride?: ChatTopic | null): ChatOptions => {
			const effectiveTopic =
				topicOverride === undefined ? selectedTopic : topicOverride;

			if (!effectiveTopic) {
				return options;
			}

			return {
				...options,
				metadata: {
					...(options.metadata ?? {}),
					topic: effectiveTopic.slug,
				},
			};
		},
		[options, selectedTopic],
	);

	useEffect(() => {
		messagesRef.current = messages;
	}, [messages]);

	useEffect(() => {
		return () => {
			if (websocketRef.current) {
				websocketRef.current.close();
				websocketRef.current = null;
			}
		};
	}, []);

	useEffect(() => {
		if (!selectedTopic) {
			topicSelectionLockedRef.current = false;
		}
	}, [selectedTopic]);

	useEffect(() => {
		if (!selectedTopic) {
			return;
		}

		const stillAvailable = availableTopics.some(
			(topic) => topic.slug === selectedTopic.slug,
		);

		if (stillAvailable) {
			return;
		}

		setSelectedTopic(null);
		if (currentSessionId) {
			persistTopicSlug(currentSessionId, null);
		}
	}, [availableTopics, currentSessionId, persistTopicSlug, selectedTopic]);

	useEffect(() => {
		if (!currentSessionId || !selectedTopic) {
			return;
		}

		persistTopicSlug(currentSessionId, selectedTopic.slug);
	}, [currentSessionId, persistTopicSlug, selectedTopic]);

	const initialMessages = useMemo<ChatMessage[]>(() => {
		return (options.initialMessages ?? []).map((text) => ({
			id: uuidv4(),
			text,
			sender: "bot" as const,
			canReceiveFeedback: false,
		}));
	}, [options.initialMessages]);

	const hasConversationMessages = messages.length > 0;
	const isTopicSelectionRequired = availableTopics.length > 0;
	const isTopicSelectionPending =
		isTopicSelectionRequired && !selectedTopic && !hasConversationMessages;

	const appendMessage = useCallback((message: ChatMessage) => {
		setMessages((prev) => [...prev, message]);
	}, []);

	const selectTopic = useCallback(
		(topic: ChatTopic, acknowledgement?: string) => {
			if (!isTopicSelectionRequired) {
				return;
			}

			if (!topic) {
				return;
			}

			if (selectedTopic?.slug === topic.slug) {
				return;
			}

			if (topicSelectionLockedRef.current) {
				return;
			}

			topicSelectionLockedRef.current = true;

			setSelectedTopic(topic);

			if (currentSessionId) {
				persistTopicSlug(currentSessionId, topic.slug);
			}

			if (acknowledgement) {
				appendMessage({
					id: uuidv4(),
					sender: "bot",
					text: acknowledgement,
					transparent: true,
					canReceiveFeedback: false,
				});
			}
		},
		[
			appendMessage,
			currentSessionId,
			isTopicSelectionRequired,
			persistTopicSlug,
			selectedTopic?.slug,
		],
	);

	const replaceMessage = useCallback(
		(messageId: string, updatedMessage: ChatMessageText) => {
			setMessages((prev) =>
				prev.map((msg) =>
					msg.id === messageId ? { ...msg, ...updatedMessage } : msg,
				),
			);
		},
		[],
	);

	const markMessageFeedbackReady = useCallback((messageId: string) => {
		setMessages((prev) =>
			prev.map((message) =>
				message.id === messageId
					? { ...message, canReceiveFeedback: true }
					: message,
			),
		);
	}, []);

	const sendMessage = useCallback(
		async (text: string, files: File[] = []) => {
			if (!currentSessionId) {
				throw new Error("Chat session has not been initialized.");
			}

			if (isTopicSelectionPending) {
				throw new Error(
					"Topic selection is required before sending messages.",
				);
			}

			const requestOptions = buildOptionsWithTopic();

			const sentMessage = createUserMessage(text, files);
			appendMessage(sentMessage);
			emitScrollToBottom();
			setWaitingForResponse(true);

			const receivedMessage: MutableRefObject<ChatMessageText | null> = {
				current: null,
			};
			const streamingManager = new StreamingMessageManager();

			try {
				if (options.enableStreaming) {
					const handlers = {
						onChunk: (
							chunk: string,
							nodeId?: string,
							runIndex?: number,
						) => {
							handleStreamingChunk(
								chunk,
								nodeId,
								streamingManager,
								receivedMessage,
								appendMessage,
								replaceMessage,
								runIndex,
							);
						},
						onBeginMessage: (nodeId: string, runIndex?: number) => {
							handleNodeStart(nodeId, streamingManager, runIndex);
						},
						onEndMessage: (nodeId: string, runIndex?: number) => {
							const completedMessage = handleNodeComplete(
								nodeId,
								streamingManager,
								runIndex,
							);

							if (completedMessage) {
								markMessageFeedbackReady(completedMessage.id);
							}
						},
					};

					const { hasReceivedChunks } =
						await api.sendMessageStreaming(
							text,
							files,
							currentSessionId,
							requestOptions,
							handlers,
						);

					if (!hasReceivedChunks) {
						handleEmptyStreamResponse(
							receivedMessage,
							appendMessage,
							replaceMessage,
							messagesRef,
						);
					}

					if (receivedMessage.current) {
						markMessageFeedbackReady(receivedMessage.current.id);
					}
				} else {
					const result = await handleNonStreamingMessage(
						text,
						files,
						currentSessionId,
						requestOptions,
					);

					if (result.response?.executionStarted) {
						setWaitingForResponse(false);
						return result.response;
					}

					if (result.botMessage) {
						appendMessage(result.botMessage);
					}
				}
			} catch (error) {
				handleMessageError(
					error,
					receivedMessage,
					appendMessage,
					replaceMessage,
				);
			} finally {
				setWaitingForResponse(false);
			}

			emitScrollToBottom();
			return null;
		},
		[
			appendMessage,
			buildOptionsWithTopic,
			isTopicSelectionPending,
			markMessageFeedbackReady,
			replaceMessage,
			currentSessionId,
		],
	);

	const sendFeedback = useCallback(
		async (messageId: string, vote: FeedbackVote) => {
			if (!options.enableFeedback) {
				return;
			}

			if (!currentSessionId) {
				return;
			}

			const targetMessage = messagesRef.current.find(
				(message) =>
					message.id === messageId &&
					message.sender === "bot" &&
					"text" in message,
			) as ChatMessageText | undefined;

			if (!targetMessage || !targetMessage.canReceiveFeedback) {
				return;
			}

			setMessages((prev) =>
				prev.map((message) =>
					message.id === messageId
						? { ...message, feedbackSubmitting: true }
						: message,
				),
			);

			try {
				await api.sendFeedback(
					targetMessage.text,
					currentSessionId,
					vote,
					buildOptionsWithTopic(),
				);

				setMessages((prev) =>
					prev.map((message) =>
						message.id === messageId
							? {
									...message,
									feedback: vote,
									feedbackSubmitting: false,
								}
							: message,
					),
				);
			} catch (error) {
				console.error("Chat feedback error:", error);
				setMessages((prev) =>
					prev.map((message) =>
						message.id === messageId
							? { ...message, feedbackSubmitting: false }
							: message,
					),
				);
			}
		},
		[buildOptionsWithTopic, currentSessionId, options.enableFeedback],
	);

	const loadPreviousSession = useCallback(async () => {
		if (!options.loadPreviousSession) {
			return false;
		}

		const sessionId =
			localStorage.getItem(localStorageSessionIdKey) ?? uuidv4();
		const { topic: storedTopic, slug: storedSlug } =
			readTopicFromStorage(sessionId);

		if (storedTopic) {
			setSelectedTopic(storedTopic);
		} else if (storedSlug) {
			persistTopicSlug(sessionId, null);
		}

		const previousMessagesResponse = await api.loadPreviousSession(
			sessionId,
			buildOptionsWithTopic(storedTopic),
		);

		const previousMessages = (previousMessagesResponse?.data || []).map(
			(message, index) => ({
				id: `${index}`,
				text: message.kwargs.content,
				sender: message.id.includes("HumanMessage")
					? ("user" as const)
					: ("bot" as const),
				canReceiveFeedback: message.id.includes("HumanMessage")
					? false
					: true,
			}),
		);

		setMessages(previousMessages);

		if (previousMessages.length) {
			setCurrentSessionId(sessionId);
			return true;
		}

		return false;
	}, [
		buildOptionsWithTopic,
		options,
		persistTopicSlug,
		readTopicFromStorage,
		setCurrentSessionId,
		setSelectedTopic,
	]);

	const startNewSession = useCallback(async () => {
		const sessionId = uuidv4();

		setMessages([]);
		messagesRef.current = [];
		setWaitingForResponse(false);

		if (websocketRef.current) {
			websocketRef.current.close();
			websocketRef.current = null;
		}

		topicSelectionLockedRef.current = false;
		setSelectedTopic(null);
		setCurrentSessionId(sessionId);
		localStorage.setItem(localStorageSessionIdKey, sessionId);
		persistTopicSlug(sessionId, null);
	}, [
		persistTopicSlug,
		setCurrentSessionId,
		setMessages,
		setWaitingForResponse,
		websocketRef,
	]);

	const value = useMemo<ChatStore>(
		() => ({
			initialMessages,
			messages,
			currentSessionId,
			waitingForResponse,
			loadPreviousSession,
			startNewSession,
			sendMessage,
			sendFeedback,
			appendMessage,
			replaceMessage,
			setWaitingForResponse,
			setCurrentSessionId,
			websocketRef,
			selectTopic,
			selectedTopic,
			isTopicSelectionPending,
			isTopicSelectionRequired,
		}),
		[
			initialMessages,
			messages,
			currentSessionId,
			waitingForResponse,
			loadPreviousSession,
			startNewSession,
			sendMessage,
			sendFeedback,
			appendMessage,
			replaceMessage,
			setWaitingForResponse,
			setCurrentSessionId,
			websocketRef,
			selectTopic,
			selectedTopic,
			isTopicSelectionPending,
			isTopicSelectionRequired,
		],
	);

	return (
		<ChatOptionsContext.Provider value={options}>
			<ChatContext.Provider value={value}>
				{children}
			</ChatContext.Provider>
		</ChatOptionsContext.Provider>
	);
}