"use client";

import type { ExportedMessageRepository } from "@assistant-ui/core";
import { AssistantRuntimeProvider } from "@assistant-ui/react";
import {
	AssistantChatTransport,
	useChatRuntime,
} from "@assistant-ui/react-ai-sdk";
import type { UIMessage } from "ai";
import { useCallback, useMemo } from "react";
import {
	ActivityPanel,
	ActivityProvider,
} from "@/components/assistant-ui/activity-panel";
import { Thread } from "@/components/assistant-ui/thread";
import { ThreadListSidebar } from "@/components/assistant-ui/threadlist-sidebar";
import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
import { ClawSessionReplayProvider } from "@/lib/claw-session-replay";
import { ClawAccountGate, useClawAccount } from "./claw-account-gate";

export const Assistant = () => {
	const { account, isLoading, setAccount } = useClawAccount();
	const transport = useMemo(
		() =>
			new AssistantChatTransport({
				api: "/api/chat",
				prepareSendMessagesRequest: async (options) => {
					const body = options.body as Record<string, unknown>;
					const messages = options.messages as UIMessage[];
					const selectedSessionId =
						typeof window !== "undefined"
							? window.localStorage.getItem("claw.activeSessionId")
							: null;
					const lastSessionId = getLastSessionId(messages);
					const selectedResumeSessionId =
						messages.length > 1 ? selectedSessionId : null;

					return {
						body: {
							...body,
							id: options.id,
							messages,
							resumeSessionId:
								lastSessionId ?? selectedResumeSessionId ?? undefined,
						},
					};
				},
			}),
		[],
	);
	const runtime = useChatRuntime({
		transport,
	});
	const replaySession = useCallback(
		(sessionId: string, repository: ExportedMessageRepository) => {
			window.localStorage.setItem("claw.activeSessionId", sessionId);
			runtime.thread.import(repository);
		},
		[runtime],
	);

	if (isLoading) {
		return (
			<div className="flex min-h-dvh items-center justify-center bg-background text-muted-foreground text-sm">
				正在加载账号...
			</div>
		);
	}

	if (!account) {
		return <ClawAccountGate onAccount={setAccount} />;
	}

	return (
		<AssistantRuntimeProvider runtime={runtime}>
			<ClawSessionReplayProvider value={{ replaySession }}>
				<ActivityProvider>
					<SidebarProvider>
						<div className="flex h-dvh w-full">
							<ThreadListSidebar
								account={account}
								onLogout={() => setAccount(null)}
							/>
							<SidebarInset>
								<div className="flex min-h-0 flex-1 overflow-hidden">
									<div className="min-w-0 flex-1 overflow-hidden">
										<Thread />
									</div>
									<ActivityPanel />
								</div>
							</SidebarInset>
						</div>
					</SidebarProvider>
				</ActivityProvider>
			</ClawSessionReplayProvider>
		</AssistantRuntimeProvider>
	);
};

function getLastSessionId(messages: UIMessage[]) {
	for (const message of [...messages].reverse()) {
		if (message.role !== "assistant") continue;
		const metadata = message.metadata;
		if (metadata && typeof metadata === "object" && "sessionId" in metadata) {
			const sessionId = (metadata as { sessionId?: unknown }).sessionId;
			if (typeof sessionId === "string" && sessionId) return sessionId;
		}
	}
	return undefined;
}