"use client";

import { AssistantRuntimeProvider, useAuiState } from "@assistant-ui/react";
import { useAISDKRuntime } from "@assistant-ui/react-ai-sdk";
import { useChat } from "@ai-sdk/react";
import { DefaultChatTransport } from "ai";
import { AlertCircle, Box, Loader2, Moon, Sun } from "lucide-react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { activeCheckpointPreview, latestSuccessfulResult } from "@/lib/cad-artifacts";
import { normalizeCadMessages } from "@/lib/cad-messages";
import type { ViewerSelectionContext } from "@/lib/viewer-selection";
import type {
  BackendConfig,
  CadError,
  CadAttachment,
  CadProgress,
  CadResult,
  CadUIMessage,
  ConversationRecord,
  TaskRecord,
} from "@/lib/cad-types";
import { AgentThread } from "./agent-thread";
import { CadViewerPreview } from "./cad-viewer-preview";
import type { AssistantRuntime } from "@assistant-ui/react";

type LoadState = "loading" | "ready" | "error";

export function AgentStudio() {
  const [loadState, setLoadState] = useState<LoadState>("loading");
  const [conversationId, setConversationId] = useState("");
  const [selectedTaskId, setSelectedTaskId] = useState("");
  const [initialMessages, setInitialMessages] = useState<CadUIMessage[]>([]);
  const [config, setConfig] = useState<BackendConfig | null>(null);
  const [cadResult, setCadResult] = useState<CadResult | null>(null);
  const [lastError, setLastError] = useState("");
  const [attachments, setAttachments] = useState<CadAttachment[]>([]);
  const [uploading, setUploading] = useState(false);
  const [uploadError, setUploadError] = useState("");
  const [providerId, setProviderId] = useState("");
  const [modelId, setModelId] = useState("");
  const [theme, setTheme] = useState<"light" | "dark">("light");
  const [viewerSelection, setViewerSelection] = useState<ViewerSelectionContext | null>(null);
  const [taskRunning, setTaskRunning] = useState(false);
  const [taskRecord, setTaskRecord] = useState<TaskRecord | null>(null);

  const syncUrl = useCallback((conversation: string, task: string) => {
    const params = new URLSearchParams(window.location.search);
    if (conversation) params.set("conversationId", conversation);
    if (task) params.set("taskId", task);
    else params.delete("taskId");
    window.history.replaceState(null, "", `${window.location.pathname}?${params.toString()}`);
  }, []);

  useEffect(() => {
    let cancelled = false;
    async function boot() {
      try {
        const params = new URLSearchParams(window.location.search);
        let nextConversationId = params.get("conversationId") || "";
        const urlTaskId = params.get("taskId") || "";
        const [configResponse] = await Promise.all([fetch("/api/config", { cache: "no-store" })]);
        if (!configResponse.ok) throw new Error(await configResponse.text());
        const nextConfig = (await configResponse.json()) as BackendConfig;

        let conversation: ConversationRecord;
        if (nextConversationId) {
          const response = await fetch(`/api/conversations/${encodeURIComponent(nextConversationId)}`, { cache: "no-store" });
          if (!response.ok) throw new Error(await response.text());
          conversation = (await response.json()) as ConversationRecord;
        } else {
          const response = await fetch("/api/conversations", { method: "POST" });
          if (!response.ok) throw new Error(await response.text());
          conversation = (await response.json()) as ConversationRecord;
          nextConversationId = conversation.conversation_id;
        }

        const taskId = urlTaskId || conversation.current_task_id || "";
        let restored: CadResult | null = null;
        let restoredTask: TaskRecord | null = null;
        if (taskId) {
          const taskResponse = await fetch(`/api/tasks/${encodeURIComponent(taskId)}`, { cache: "no-store" });
          if (taskResponse.ok) {
            restoredTask = (await taskResponse.json()) as TaskRecord;
            restored = activeCheckpointPreview(restoredTask) ?? latestSuccessfulResult(restoredTask);
          }
        }

        if (cancelled) return;
        setConfig(nextConfig);
        setConversationId(nextConversationId);
        setSelectedTaskId(taskId);
        setInitialMessages(normalizeCadMessages(conversation.messages));
        setAttachments(conversation.attachments || []);
        const defaultProvider = nextConfig.providers.find((provider) => provider.id === nextConfig.default_provider) ?? nextConfig.providers[0];
        setProviderId(defaultProvider?.id || "");
        setModelId(defaultProvider?.models.find((model) => model.id === nextConfig.default_model)?.id || defaultProvider?.models[0]?.id || "");
        setCadResult(restored);
        setTaskRunning(restoredTask?.lifecycle === "running");
        setTaskRecord(restoredTask);
        setLoadState("ready");
        syncUrl(nextConversationId, taskId);
      } catch (error) {
        if (cancelled) return;
        setLastError(error instanceof Error ? error.message : "启动失败");
        setLoadState("error");
      }
    }
    void boot();
    return () => {
      cancelled = true;
    };
  }, [syncUrl]);

  useEffect(() => {
    const stored = window.localStorage.getItem("cdsl-cad.ui-theme");
    const next = stored === "dark" || stored === "light"
      ? stored
      : window.matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
    setTheme(next);
    document.documentElement.dataset.uiTheme = next;
  }, []);

  const toggleTheme = useCallback(() => {
    setTheme((current) => {
      const next = current === "light" ? "dark" : "light";
      document.documentElement.dataset.uiTheme = next;
      window.localStorage.setItem("cdsl-cad.ui-theme", next);
      return next;
    });
  }, []);

  const handleResult = useCallback((result: CadResult) => {
    setCadResult(result);
    setViewerSelection(null);
    setSelectedTaskId(result.taskId);
    setLastError("");
    if (result.lifecycle) setTaskRunning(result.lifecycle === "running");
    if (conversationId) {
      syncUrl(conversationId, result.taskId);
      void fetch(`/api/conversations/${encodeURIComponent(conversationId)}`, {
        method: "PATCH",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ currentTaskId: result.taskId }),
      });
    }
  }, [conversationId, syncUrl]);

  const handleTaskState = useCallback((progress: CadProgress) => {
    const taskId = String(progress.taskId || "");
    if (taskId) {
      setSelectedTaskId(taskId);
      if (conversationId) {
        syncUrl(conversationId, taskId);
        void fetch(`/api/conversations/${encodeURIComponent(conversationId)}`, {
          method: "PATCH", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ currentTaskId: taskId }),
        });
      }
    }
    if (progress.step === "requirements_document" && progress.status === "running") setTaskRunning(true);
    if (progress.step === "task_terminal") setTaskRunning(progress.lifecycle === "running");
  }, [conversationId, syncUrl]);

  const handleError = useCallback((error: CadError) => {
    setLastError(error.message);
    if (error.stage === "generation") setTaskRunning(false);
  }, []);

  const handleCancel = useCallback(() => {
    setTaskRunning(false);
    setLastError("正在停止 CAD 任务...");
    void (async () => {
      let taskId = selectedTaskId;
      if (!taskId && conversationId) {
        const conversationResponse = await fetch(`/api/conversations/${encodeURIComponent(conversationId)}`, { cache: "no-store" });
        if (conversationResponse.ok) {
          const conversation = await conversationResponse.json() as ConversationRecord;
          taskId = conversation.current_task_id || "";
        }
      }
      if (!taskId) throw new Error("尚未取得运行中的任务编号，请稍后重试");
      const response = await fetch(`/api/tasks/${encodeURIComponent(taskId)}`, { method: "DELETE" });
      if (!response.ok) {
        const payload = await response.json().catch(() => ({})) as { error?: string };
        throw new Error(payload.error || "停止 CAD 任务失败");
      }
      setLastError("CAD 任务已停止");
    })().catch((error) => {
      setLastError(error instanceof Error ? error.message : "停止 CAD 任务失败");
    });
  }, [conversationId, selectedTaskId]);

  const handleUpload = useCallback(async (files: FileList | null) => {
    const selectedFiles = Array.from(files || []);
    if (!selectedFiles.length) return;
    setUploadError("");
    setUploading(true);
    try {
      const uploaded: CadAttachment[] = [];
      for (const file of selectedFiles) {
        const form = new FormData();
        form.set("file", file);
        const response = await fetch(`/api/conversations/${encodeURIComponent(conversationId)}/attachments`, { method: "POST", body: form });
        const payload = await response.json() as CadAttachment & { error?: string };
        if (!response.ok) throw new Error(payload.error || `${file.name} 上传失败`);
        uploaded.push(payload);
      }
      setAttachments((current) => [...current, ...uploaded]);
      setLastError("");
    } catch (error) {
      const message = error instanceof Error ? error.message : "附件上传失败";
      setUploadError(message);
      setLastError(message);
    } finally {
      setUploading(false);
    }
  }, [conversationId]);

  useEffect(() => {
    if (!taskRunning || !selectedTaskId) return;
    let cancelled = false;
    const refresh = async () => {
      try {
        const response = await fetch(`/api/tasks/${encodeURIComponent(selectedTaskId)}`, { cache: "no-store" });
        if (!response.ok) return;
        const next = await response.json() as TaskRecord;
        if (cancelled) return;
        setTaskRecord(next);
        const running = next.lifecycle === "running";
        setTaskRunning(running);
        if (running) {
          const preview = activeCheckpointPreview(next);
          if (preview) setCadResult(preview);
        } else {
          const restored = latestSuccessfulResult(next);
          setCadResult(restored);
        }
      } catch {
        // Keep the persisted lock until a later poll can prove a terminal state.
      }
    };
    void refresh();
    const timer = window.setInterval(() => void refresh(), 2000);
    return () => { cancelled = true; window.clearInterval(timer); };
  }, [selectedTaskId, taskRunning]);

  if (loadState === "loading") {
    return <StudioLoading />;
  }
  if (loadState === "error") {
    return <StudioError message={lastError} />;
  }

  return (
    <AgentRuntime
      key={conversationId}
      conversationId={conversationId}
      selectedTaskId={selectedTaskId}
      providerId={providerId}
      modelId={modelId}
      viewerSelection={viewerSelection}
      initialMessages={initialMessages}
      onCadResult={handleResult}
      onCadError={handleError}
      onCadProgress={handleTaskState}
    >
      <StudioShell
        config={config}
        cadResult={cadResult}
        lastError={lastError}
        attachments={attachments}
        uploading={uploading}
        uploadError={uploadError}
        onUpload={handleUpload}
        onCancel={handleCancel}
        theme={theme}
        onToggleTheme={toggleTheme}
        providerId={providerId}
        modelId={modelId}
        onProviderChange={setProviderId}
        onModelChange={setModelId}
        onCadResult={handleResult}
        onCadError={handleError}
        onSelectionChange={setViewerSelection}
        taskRunning={taskRunning}
      />
    </AgentRuntime>
  );
}

function AgentRuntime({
  conversationId,
  selectedTaskId,
  providerId,
  modelId,
  viewerSelection,
  initialMessages,
  onCadResult,
  onCadError,
  onCadProgress,
  children,
}: {
  conversationId: string;
  selectedTaskId: string;
  providerId: string;
  modelId: string;
  viewerSelection: ViewerSelectionContext | null;
  initialMessages: CadUIMessage[];
  onCadResult: (result: CadResult) => void;
  onCadError: (error: CadError) => void;
  onCadProgress: (progress: CadProgress) => void;
  children: React.ReactNode;
}) {
  const conversationRef = useRef(conversationId);
  const taskRef = useRef(selectedTaskId);
  const providerRef = useRef(providerId);
  const modelRef = useRef(modelId);
  const viewerSelectionRef = useRef(viewerSelection);
  const onCadResultRef = useRef(onCadResult);
  const onCadErrorRef = useRef(onCadError);
  const onCadProgressRef = useRef(onCadProgress);
  conversationRef.current = conversationId;
  taskRef.current = selectedTaskId;
  providerRef.current = providerId;
  modelRef.current = modelId;
  viewerSelectionRef.current = viewerSelection;
  onCadResultRef.current = onCadResult;
  onCadErrorRef.current = onCadError;
  onCadProgressRef.current = onCadProgress;

  const transport = useMemo(() => new DefaultChatTransport<CadUIMessage>({
    api: "/api/chat",
    prepareSendMessagesRequest: (options) => ({
      body: {
        ...options.body,
        id: options.id,
        messages: options.messages,
        conversationId: conversationRef.current,
        selectedTaskId: taskRef.current || null,
        providerId: providerRef.current || null,
        modelId: modelRef.current || null,
        viewerContext: viewerSelectionRef.current ? [viewerSelectionRef.current] : [],
        trigger: options.trigger,
        messageId: options.messageId,
      },
    }),
  }), []);

  const chat = useChat<CadUIMessage>({
    id: conversationId,
    messages: initialMessages,
    transport,
    onData: (part) => {
      if (part.type === "data-cad-progress") {
        onCadProgressRef.current(part.data as CadProgress);
      }
      if (part.type === "data-cad-result") {
        onCadResultRef.current(part.data as CadResult);
        taskRef.current = (part.data as CadResult).taskId;
      }
      if (part.type === "data-cad-error") {
        onCadErrorRef.current(part.data as CadError);
      }
    },
    onError: (error) => {
      onCadErrorRef.current({ stage: "chat", message: error.message });
    },
  });
  const runtime = useAISDKRuntime<CadUIMessage>(chat, {
    joinStrategy: "none",
  });
  const stableRuntime = useMemo(() => stabilizeAssistantRuntimeSnapshots(runtime), [runtime]);

  return <AssistantRuntimeProvider runtime={stableRuntime}>{children}</AssistantRuntimeProvider>;
}

const stableSnapshotSymbol = Symbol.for("cdsl-cad.stableRuntimeSnapshot");
const stableChildRuntimeSymbol = Symbol.for("cdsl-cad.stableChildRuntimeMethods");

type SnapshotRuntime = {
  getState: () => unknown;
  [stableSnapshotSymbol]?: true;
  [stableChildRuntimeSymbol]?: true;
};

function stabilizeAssistantRuntimeSnapshots(runtime: AssistantRuntime) {
  stabilizeThreadListRuntime(runtime.threads);
  stabilizeThreadRuntime(runtime.thread);
  return runtime;
}

function stabilizeThreadListRuntime(runtime: AssistantRuntime["threads"]) {
  const threadList = runtime as AssistantRuntime["threads"] & SnapshotRuntime;
  stabilizeSnapshot(threadList);
  stabilizeThreadRuntime(threadList.main);
  stabilizeSnapshotIfRuntime(threadList.mainItem);
  if (threadList[stableChildRuntimeSymbol]) return;

  wrapRuntimeFactory(threadList, "getById", stabilizeThreadRuntime);
  wrapRuntimeFactory(threadList, "getItemById", stabilizeSnapshotIfRuntime);
  wrapRuntimeFactory(threadList, "getItemByIndex", stabilizeSnapshotIfRuntime);
  wrapRuntimeFactory(threadList, "getArchivedItemByIndex", stabilizeSnapshotIfRuntime);
  threadList[stableChildRuntimeSymbol] = true;
}

function stabilizeThreadRuntime(runtime: unknown) {
  if (!isSnapshotRuntime(runtime)) return runtime;
  stabilizeSnapshot(runtime);
  const thread = runtime as SnapshotRuntime & {
    composer?: unknown;
    getMessageById?: (...args: unknown[]) => unknown;
    getMessageByIndex?: (...args: unknown[]) => unknown;
  };
  stabilizeComposerRuntime(thread.composer);
  if (thread[stableChildRuntimeSymbol]) return thread;

  wrapRuntimeFactory(thread, "getMessageById", stabilizeMessageRuntime);
  wrapRuntimeFactory(thread, "getMessageByIndex", stabilizeMessageRuntime);
  thread[stableChildRuntimeSymbol] = true;
  return thread;
}

function stabilizeMessageRuntime(runtime: unknown) {
  if (!isSnapshotRuntime(runtime)) return runtime;
  stabilizeSnapshot(runtime);
  const message = runtime as SnapshotRuntime & {
    composer?: unknown;
    getAttachmentByIndex?: (...args: unknown[]) => unknown;
    getMessagePartByIndex?: (...args: unknown[]) => unknown;
    getMessagePartByToolCallId?: (...args: unknown[]) => unknown;
  };
  stabilizeComposerRuntime(message.composer);
  if (message[stableChildRuntimeSymbol]) return message;

  wrapRuntimeFactory(message, "getAttachmentByIndex", stabilizeSnapshotIfRuntime);
  wrapRuntimeFactory(message, "getMessagePartByIndex", stabilizeSnapshotIfRuntime);
  wrapRuntimeFactory(message, "getMessagePartByToolCallId", stabilizeSnapshotIfRuntime);
  message[stableChildRuntimeSymbol] = true;
  return message;
}

function stabilizeComposerRuntime(runtime: unknown) {
  if (!isSnapshotRuntime(runtime)) return runtime;
  stabilizeSnapshot(runtime);
  const composer = runtime as SnapshotRuntime & {
    getAttachmentByIndex?: (...args: unknown[]) => unknown;
  };
  if (composer[stableChildRuntimeSymbol]) return composer;

  wrapRuntimeFactory(composer, "getAttachmentByIndex", stabilizeSnapshotIfRuntime);
  composer[stableChildRuntimeSymbol] = true;
  return composer;
}

function stabilizeSnapshot(runtime: SnapshotRuntime) {
  if (runtime[stableSnapshotSymbol]) return;

  const getState = runtime.getState.bind(runtime);
  let previous: unknown;
  runtime.getState = () => {
    const next = getState();
    if (isShallowSameSnapshot(previous, next)) return previous;
    previous = next;
    return next;
  };
  runtime[stableSnapshotSymbol] = true;
}

function stabilizeSnapshotIfRuntime(value: unknown) {
  if (isSnapshotRuntime(value)) stabilizeSnapshot(value);
  return value;
}

function wrapRuntimeFactory(
  runtime: Record<PropertyKey, unknown>,
  method: string,
  stabilize: (value: unknown) => unknown,
) {
  const original = runtime[method];
  if (typeof original !== "function") return;
  runtime[method] = (...args: unknown[]) => stabilize(original.apply(runtime, args));
}

function isSnapshotRuntime(value: unknown): value is SnapshotRuntime {
  return Boolean(value && typeof value === "object" && typeof (value as SnapshotRuntime).getState === "function");
}

function isShallowSameSnapshot(left: unknown, right: unknown) {
  if (Object.is(left, right)) return true;
  if (!left || !right || typeof left !== "object" || typeof right !== "object") return false;
  const leftRecord = left as Record<string, unknown>;
  const rightRecord = right as Record<string, unknown>;
  const leftKeys = Object.keys(leftRecord);
  if (leftKeys.length !== Object.keys(rightRecord).length) return false;
  return leftKeys.every((key) => Object.is(leftRecord[key], rightRecord[key]));
}

function StudioShell({
  config,
  cadResult,
  lastError,
  attachments,
  uploading,
  uploadError,
  onUpload,
  onCancel,
  theme,
  onToggleTheme,
  providerId,
  modelId,
  onProviderChange,
  onModelChange,
  onCadResult,
  onCadError,
  onSelectionChange,
  taskRunning,
}: {
  config: BackendConfig | null;
  cadResult: CadResult | null;
  lastError: string;
  attachments: CadAttachment[];
  uploading: boolean;
  uploadError: string;
  onUpload: (files: FileList | null) => void;
  onCancel: () => void;
  theme: "light" | "dark";
  onToggleTheme: () => void;
  providerId: string;
  modelId: string;
  onProviderChange: (id: string) => void;
  onModelChange: (id: string) => void;
  onCadResult: (result: CadResult) => void;
  onCadError: (error: CadError) => void;
  onSelectionChange: (selection: ViewerSelectionContext | null) => void;
  taskRunning: boolean;
}) {
  const running = useAuiState((state) => state.thread.isRunning) || taskRunning;
  const provider = config?.providers.find((item) => item.id === providerId);
  const handleViewerError = useCallback((message: string) => {
    onCadError({ stage: "viewer", message });
  }, [onCadError]);
  return (
    <main className="studio-app">
      <header className="app-header">
        <div className="app-brand"><Box size={16} /><strong>CDSL CAD Studio</strong>{cadResult ? <span className="task-badge">{cadResult.taskId}</span> : null}</div>
        <div className="app-controls">
          <select aria-label="模型提供商" value={providerId} disabled={running} onChange={(event) => { const id = event.target.value; const models = config?.providers.find((item) => item.id === id)?.models || []; onProviderChange(id); onModelChange(models[0]?.id || ""); }}>
            {config?.providers.map((item) => <option key={item.id} value={item.id}>{item.label}</option>)}
          </select>
          <select aria-label="模型" value={modelId} disabled={running} onChange={(event) => onModelChange(event.target.value)}>
            {provider?.models.map((model) => <option key={model.id} value={model.id}>{model.id}{model.vision ? " · Vision" : ""}</option>)}
          </select>
          <button className="theme-button" type="button" title="切换亮暗主题" disabled={running} onClick={onToggleTheme}>{theme === "light" ? <Moon size={16} /> : <Sun size={16} />}</button>
        </div>
      </header>
      {!config?.configured ? <div className="config-warning"><AlertCircle size={16} /><span>未配置模型环境变量，聊天会保留诊断但不会生成虚假模型。</span></div> : null}
      {config?.autonomous_generation && !config.review_configured ? <div className="config-warning"><AlertCircle size={16} /><span>最终视觉复核未配置，任务在最终发布前会停止：{config.review_error || "请配置独立视觉模型。"}</span></div> : null}
      <div className="studio-main">
        <aside className="agent-pane"><AgentThread attachments={attachments} uploading={uploading} uploadError={uploadError} taskRunning={taskRunning} onUpload={onUpload} onCancel={onCancel} /></aside>
        <section className="preview-pane">
          <CadViewerPreview result={cadResult} isGenerating={running} lastError={lastError} theme={theme} onError={handleViewerError} onSelectionChange={onSelectionChange} />
        </section>
      </div>
    </main>
  );
}

function StudioLoading() {
  return (
    <main className="boot-screen">
      <Loader2 className="spin" size={22} />
      <span>启动 CDSL CAD Studio</span>
    </main>
  );
}

function StudioError({ message }: { message: string }) {
  return (
    <main className="boot-screen boot-error">
      <AlertCircle size={22} />
      <span>{message}</span>
    </main>
  );
}