"use client";

import type { ExportedMessageRepository } from "@assistant-ui/core";
import { AssistantRuntimeProvider, useAui } from "@assistant-ui/react";
import {
	AssistantChatTransport,
	useChatRuntime,
} from "@assistant-ui/react-ai-sdk";
import type { UIMessage } from "ai";
import { PanelRightOpenIcon, XIcon } from "lucide-react";
import {
	type MutableRefObject,
	useCallback,
	useEffect,
	useMemo,
	useRef,
	useState,
} from "react";
import {
	ActivityPanel,
	ActivityProvider,
	useActivityPanel,
} from "@/components/assistant-ui/activity-panel";
import { Thread } from "@/components/assistant-ui/thread";
import {
	type fetchLatestRunStatus,
	toReplayRepository,
} from "@/components/assistant-ui/thread-list";
import { ThreadListSidebar } from "@/components/assistant-ui/threadlist-sidebar";
import { TrainingPipelinePanel } from "@/components/training/training-pipeline-panel";
import { Button } from "@/components/ui/button";
import {
	SidebarInset,
	SidebarProvider,
	SidebarTrigger,
} from "@/components/ui/sidebar";
import {
	ACTIVE_SESSION_CHANGED_EVENT,
	consumeFreshLocalId,
	createLocalSessionId,
	readActiveSessionId,
	readPendingWorkspaceSessionId,
	writeActiveSessionId,
} from "@/lib/claw-active-session";
import {
	fetchSessionReplaySnapshot,
	readCachedSessionReplay,
	type SessionReplaySnapshot,
} from "@/lib/claw-session-cache";
import { ClawSessionReplayProvider } from "@/lib/claw-session-replay";
import { pushSessionUrl, readSessionIdFromUrl } from "@/lib/claw-session-url";
import {
	SKILL_TOGGLED_EVENT,
	type SkillToggledDetail,
	useAppliedTraining,
	useModelIterationEnabled,
} from "@/lib/use-training-mode";
import { cn } from "@/lib/utils";
import { ClawAccountGate, useClawAccount } from "./claw-account-gate";

type AssistantProps = {
	initialSessionId?: string;
};

const OPTIMISTIC_REPLAY_GUARD_MS = 30_000;

export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
	const { account, isLoading, setAccount } = useClawAccount();
	const loadedInitialSessionRef = useRef<string | null>(null);
	const optimisticSendRef = useRef<Map<string, number>>(new Map());
	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 =
						readSessionIdFromUrl() ?? readActiveSessionId() ?? null;
					const pendingWorkspaceSessionId =
						!selectedSessionId && messages.length <= 1
							? readPendingWorkspaceSessionId()
							: null;
					const selectedResumeSessionId = selectedSessionId;
					const outgoingSessionId =
						selectedResumeSessionId ??
						pendingWorkspaceSessionId ??
						createLocalSessionId();
					const lastUserMessage = [...messages]
						.reverse()
						.find((message) => message.role === "user");
					if (lastUserMessage) {
						optimisticSendRef.current.set(outgoingSessionId, Date.now());
					}
					writeActiveSessionId(outgoingSessionId);
					pushSessionUrl(outgoingSessionId);

					return {
						body: {
							...body,
							id: outgoingSessionId,
							messages: lastUserMessage ? [lastUserMessage] : [],
							resumeSessionId: selectedResumeSessionId ?? undefined,
						},
					};
				},
			}),
		[],
	);
	const runtime = useChatRuntime({
		transport,
	});
	const replaySession = useCallback(
		(sessionId: string, repository: ExportedMessageRepository) => {
			console.log("[replay-session] called", { sessionId });
			writeActiveSessionId(sessionId);
			pushSessionUrl(sessionId);
			runtime.thread.import(repository);
		},
		[runtime],
	);
	const clearSession = useCallback(() => {
		runtime.thread.import({ headId: null, messages: [] });
	}, [runtime]);

	useEffect(() => {
		const sessionId = initialSessionId?.trim();
		if (!account || !sessionId) return;
		if (loadedInitialSessionRef.current === sessionId) return;
		loadedInitialSessionRef.current = sessionId;
		const targetSessionId = sessionId;
		let cancelled = false;
		const abortController = new AbortController();

		async function loadSession() {
			try {
				writeActiveSessionId(targetSessionId);
				const cached = readCachedSessionReplay<
					Parameters<typeof toReplayRepository>[0],
					Awaited<ReturnType<typeof fetchLatestRunStatus>>
				>(targetSessionId);
				if (cached && !cancelled) {
					if (
						!shouldSkipOptimisticEmptyReplay(
							optimisticSendRef,
							targetSessionId,
							cached,
						)
					) {
						replaySession(
							targetSessionId,
							toReplayRepository(
								cached.session,
								targetSessionId,
								cached.runStatus,
							),
						);
					}
				}
				const snapshot = await fetchSessionReplaySnapshot<
					Parameters<typeof toReplayRepository>[0],
					Awaited<ReturnType<typeof fetchLatestRunStatus>>
				>(targetSessionId, { signal: abortController.signal });
				if (cancelled) return;
				if (!snapshot) return;
				if (
					shouldSkipOptimisticEmptyReplay(
						optimisticSendRef,
						targetSessionId,
						snapshot,
					)
				) {
					return;
				}
				if (!cached || snapshot.signature !== cached.signature) {
					replaySession(
						targetSessionId,
						toReplayRepository(
							snapshot.session,
							targetSessionId,
							snapshot.runStatus,
						),
					);
				}
				pushSessionUrl(targetSessionId);
			} catch (error) {
				if (error instanceof DOMException && error.name === "AbortError") {
					return;
				}
				if (!cancelled) {
					loadedInitialSessionRef.current = null;
				}
			}
		}

		void loadSession();
		return () => {
			cancelled = true;
			abortController.abort();
		};
	}, [account, initialSessionId, replaySession]);

	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, clearSession }}>
				<ActivityProvider>
					<SidebarProvider>
						<div className="flex h-dvh w-full">
							<SidebarTrigger className="fixed top-3 left-3 z-40 size-9 rounded-lg bg-background/90 shadow-sm ring-1 ring-border backdrop-blur sm:hidden" />
							<ThreadListSidebar
								account={account}
								onLogout={() => setAccount(null)}
							/>
							<SidebarInset>
								<AssistantWorkspace />
							</SidebarInset>
						</div>
					</SidebarProvider>
				</ActivityProvider>
			</ClawSessionReplayProvider>
		</AssistantRuntimeProvider>
	);
};

function shouldSkipOptimisticEmptyReplay(
	optimisticSendRef: MutableRefObject<Map<string, number>>,
	sessionId: string,
	snapshot: SessionReplaySnapshot | null,
) {
	if (hasReplayMessages(snapshot)) {
		optimisticSendRef.current.delete(sessionId);
		return false;
	}
	const sentAt = optimisticSendRef.current.get(sessionId);
	if (sentAt === undefined) return false;
	if (Date.now() - sentAt > OPTIMISTIC_REPLAY_GUARD_MS) {
		optimisticSendRef.current.delete(sessionId);
		return false;
	}
	return true;
}

function hasReplayMessages(snapshot: SessionReplaySnapshot | null) {
	const session = snapshot?.session;
	if (!session || typeof session !== "object" || Array.isArray(session)) {
		return false;
	}
	const messages = (session as { messages?: unknown }).messages;
	return Array.isArray(messages) && messages.length > 0;
}

// 共享 composer 草稿缓存里 newTask（无 sessionId）使用的 key。
const NEWTASK_KEY = "__newtask__";

// AssistantWorkspace cache effect 通知 ImeComposerInput 直接更新 textarea
// 内容（绕过 store→local 间接路径，避免 switchToNewThread 时序问题）。
export const COMPOSER_RESTORE_EVENT = "claw-composer-restore-draft";

// __LOCALID_xxx 是侧栏在 newTask 状态下点工作区时临时分配的 placeholder
// session id（让 jupyter bind 有 id 可传）。每个 LOCALID 拥有独立的 draft
// key，避免多个 LOCALID 会话在 cache 里互相覆盖。newTask → 首次生成 LOCALID
// 这一瞬的 textarea 闪空，由下面 save/restore effect 里的一次性继承处理。
function getDraftKey(activeSessionId: string | null): string {
	if (!activeSessionId) return NEWTASK_KEY;
	return activeSessionId;
}

function AssistantWorkspace() {
	const aui = useAui();
	const [activeSessionId, setActiveSessionId] = useState<string | null>(
		() => readSessionIdFromUrl() ?? readActiveSessionId(),
	);
	const skill = useModelIterationEnabled();
	const { applied, setApplied } = useAppliedTraining(activeSessionId);
	const { close: closeActivityPanel } = useActivityPanel();
	const showPipeline = skill.enabled && applied;
	const pipelineContainerRef = useRef<HTMLDivElement>(null);
	const lastPipelineYRef = useRef<number | null>(null);
	const [exitButtonVisible, setExitButtonVisible] = useState(false);
	// Composer 是 per-thread 的，但本应用所有 backend session 都共用同一个
	// assistant-ui thread（从不调 switchToThread），因此 composer.text 跨 session
	// 共享。在 activeSessionId 切换时手动 save+restore，给每个 session 一份
	// 独立的 draft。null sessionId（newTask）用 NEWTASK_KEY 落盘。
	const composerCacheRef = useRef<Map<string, string>>(new Map());
	const prevSessionKeyRef = useRef<string>(getDraftKey(activeSessionId));

	const handlePipelineMouseMove = useCallback(
		(event: React.MouseEvent<HTMLDivElement>) => {
			const node = pipelineContainerRef.current;
			if (!node) return;
			const rect = node.getBoundingClientRect();
			const relativeY = event.clientY - rect.top;
			const inTopThird = relativeY >= 0 && relativeY <= rect.height / 3;
			const lastY = lastPipelineYRef.current;
			const movingUp = lastY !== null && event.clientY < lastY;
			const movingDown = lastY !== null && event.clientY > lastY;
			lastPipelineYRef.current = event.clientY;
			setExitButtonVisible((prev) => {
				if (!inTopThird) return false;
				if (movingUp) return true;
				if (movingDown) return false;
				return prev;
			});
		},
		[],
	);

	const handlePipelineMouseLeave = useCallback(() => {
		lastPipelineYRef.current = null;
		setExitButtonVisible(false);
	}, []);

	useEffect(() => {
		if (!showPipeline) {
			lastPipelineYRef.current = null;
			setExitButtonVisible(false);
		}
	}, [showPipeline]);

	useEffect(() => {
		const update = () => {
			const next = readSessionIdFromUrl() ?? readActiveSessionId();
			console.log("[active-session] focus-update", { next });
			setActiveSessionId(next);
		};
		const handleChanged = (event: Event) => {
			const customEvent = event as CustomEvent<{ sessionId?: string | null }>;
			const hasExplicitSession =
				customEvent.detail && Object.hasOwn(customEvent.detail, "sessionId");
			const next = hasExplicitSession
				? (customEvent.detail.sessionId ?? null)
				: (readSessionIdFromUrl() ?? readActiveSessionId());
			console.log("[active-session] event-changed", {
				detailSessionId: customEvent.detail?.sessionId,
				next,
			});
			setActiveSessionId(next);
		};
		window.addEventListener(ACTIVE_SESSION_CHANGED_EVENT, handleChanged);
		window.addEventListener("focus", update);
		window.addEventListener("popstate", update);
		update();
		return () => {
			window.removeEventListener(ACTIVE_SESSION_CHANGED_EVENT, handleChanged);
			window.removeEventListener("focus", update);
			window.removeEventListener("popstate", update);
		};
	}, []);

	// 切 session 时保存当前 composer 草稿到旧 key、恢复新 key 的草稿。
	useEffect(() => {
		const cache = composerCacheRef.current;
		const nextKey = getDraftKey(activeSessionId);
		const prevKey = prevSessionKeyRef.current;
		console.log("[draft-cache] effect-fire", {
			prevKey,
			nextKey,
			activeSessionId,
			willSkip: prevKey === nextKey,
		});
		if (prevKey === nextKey) return;
		// 不能信任 aui store 的 composer.text——assistant-ui 内部 reducer 在
		// 切换 / isEditing 翻转时会让 store 与 textarea 脱钩（典型现象：textarea
		// 显示着 "hello" 但 store 已经是 ""），后果是这里把空串存进 cache、
		// 用户的草稿被静默吞掉。直接读 DOM 拿 textarea 当前真实值。
		// EditComposer 用的是 aui-edit-composer-input，不会撞 selector。
		const composerEl = document.querySelector<HTMLTextAreaElement>(
			"textarea.aui-composer-input",
		);
		const composerState = aui.composer().getState();
		const savingText = composerEl?.value ?? composerState.text ?? "";
		cache.set(prevKey, savingText);
		// newTask 里点工作区会**新生成**一个 LOCALID 并切过去，用户视角里
		// 这是同一份草稿的延续——把 newTask 的文本继承到 LOCALID 名下，
		// 避免 textarea 闪空。consumeFreshLocalId 只对前端刚生成的 LOCALID
		// 命中，点击侧栏老 LOCALID 不命中，确保历史会话之间互不串。
		const inheritFromNewTask =
			prevKey === NEWTASK_KEY && consumeFreshLocalId(activeSessionId);
		if (inheritFromNewTask) {
			cache.set(nextKey, savingText);
		}
		const restored = cache.get(nextKey) ?? "";
		console.log("[draft-cache] do-switch", {
			prevKey,
			nextKey,
			savingText,
			savingTextSource: composerEl ? "dom" : "store",
			storeText: composerState.text,
			restored,
			inheritFromNewTask,
		});
		aui.composer().setText(restored);
		window.dispatchEvent(
			new CustomEvent(COMPOSER_RESTORE_EVENT, {
				detail: { text: restored },
			}),
		);
		console.log("[draft-cache] dispatched", { restored });
		prevSessionKeyRef.current = nextKey;
	}, [activeSessionId, aui]);

	// biome-ignore lint/correctness/useExhaustiveDependencies: close the activity drawer whenever training mode toggles.
	useEffect(() => {
		closeActivityPanel();
	}, [showPipeline, closeActivityPanel]);

	// 面板只通过用户在 SkillInsertDialog 里显式点击开/关 model-iteration 来控制。
	// 监听 SkillInsertDialog 发出的自定义事件，不再从 poll 状态变化推断。
	useEffect(() => {
		const handler = (e: Event) => {
			const detail = (e as CustomEvent<SkillToggledDetail>).detail;
			if (detail.skill !== "model-iteration") return;
			setApplied(detail.enabled);
		};
		window.addEventListener(SKILL_TOGGLED_EVENT, handler);
		return () => window.removeEventListener(SKILL_TOGGLED_EVENT, handler);
	}, [setApplied]);

	// 注意：用同一棵树 + 条件渲染，保证 <Thread /> 和 <ActivityPanel /> 在
	// pipeline 切换前后都保持挂载，否则它们会 remount，
	// useJupyterWorkspaceStatus 的内部 state 会被清空，UI 上 Jupyter 工作区
	// indicator 会闪一下变成"切换工作区"。
	return (
		<div className="flex min-h-0 flex-1 overflow-hidden">
			{showPipeline ? (
				<section
					aria-label="训练模式面板"
					ref={pipelineContainerRef}
					onMouseMove={handlePipelineMouseMove}
					onMouseLeave={handlePipelineMouseLeave}
					className="relative flex min-w-0 flex-1 flex-col"
				>
					<button
						type="button"
						onClick={() => setApplied(false)}
						title="退出训练模式"
						aria-label="退出训练模式"
						className={cn(
							"-translate-x-1/2 absolute top-0 left-1/2 z-20 inline-flex items-center gap-1 rounded-b-md border-x border-b border-border bg-background px-3 py-1 text-muted-foreground text-xs shadow-md ring-offset-background transition-transform duration-200 hover:bg-muted hover:text-foreground focus:outline-hidden focus:ring-2 focus:ring-ring focus:ring-offset-2",
							exitButtonVisible ? "translate-y-0" : "-translate-y-full",
						)}
					>
						<XIcon className="size-3.5" />
						<span>退出训练模式</span>
					</button>
					<TrainingPipelinePanel sessionId={activeSessionId} />
				</section>
			) : null}
			<div
				className={cn(
					"flex flex-col overflow-hidden",
					showPipeline ? "w-[28rem] shrink-0 border-l" : "min-w-0 flex-1",
				)}
			>
				<div className="min-h-0 flex-1 overflow-hidden">
					<Thread />
				</div>
			</div>
			<ActivityPanel asDrawer={showPipeline} />
			{showPipeline ? <ActivityDrawerTrigger /> : null}
		</div>
	);
}

function ActivityDrawerTrigger() {
	const { open, openActivity } = useActivityPanel();
	if (open) return null;
	return (
		<Button
			type="button"
			variant="outline"
			size="icon"
			onClick={openActivity}
			className="fixed top-3 right-3 z-30 size-9 rounded-lg bg-background/90 shadow-sm backdrop-blur"
			aria-label="打开活动面板"
		>
			<PanelRightOpenIcon className="size-4" />
		</Button>
	);
}
