"use client";

import type { DragEvent } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import {
  AssistantRuntimeProvider,
  ComposerPrimitive,
  MessagePartPrimitive,
  MessagePrimitive,
  ThreadPrimitive,
  useMessage,
  type DataMessagePartProps,
  type ToolCallMessagePartProps,
} from "@assistant-ui/react";
import { AssistantChatTransport, useChatRuntime } from "@assistant-ui/react-ai-sdk";
import type { UIMessage } from "ai";
import { CadViewerPreview } from "@/components/cad-viewer-preview";
import { AiSelectionMode } from "@/lib/cad-edit-tools";
import {
  Bot,
  Box,
  Check,
  CircleAlert,
  FileImage,
  FileUp,
  Loader2,
  MessageSquare,
  Paperclip,
  Send,
  Settings2,
  Sparkles,
  Square,
  Wrench,
} from "lucide-react";

type Attachment = {
  id: string;
  taskId: string;
  name: string;
  kind: "image" | "step" | "source" | "manifest" | "other";
  path: string;
  size: number;
  sha256: string;
  mime: string;
  artifactUrl?: string;
  previewPath?: string;
  previewUrl?: string;
  viewerAssetPath?: string;
  viewerAssetUrl?: string;
  previewError?: string;
};

type ViewerContext = {
  id: string;
  type: string;
  payload: Record<string, unknown>;
  createdAt: string;
};

type CadResult = {
  taskId: string;
  sourcePath?: string;
  artifactPath?: string;
  artifactUrl?: string;
  previewPath?: string;
  previewUrl?: string;
  viewerAssetPath?: string;
  viewerAssetUrl?: string;
  summary?: string;
};

type CadProgress = {
  step: string;
  label: string;
  status: "pending" | "running" | "success" | "error";
  message?: string;
};

type CadError = {
  stage?: string;
  message: string;
};

type PublicConfig = {
  defaultProvider: string;
  defaultModel: string;
  providers: Record<string, {
    type: string;
    baseURL: string;
    configured: boolean;
    models: Record<string, string>;
  }>;
  uploads: {
    maxImageMB: number;
    maxStepMB: number;
  };
};

const initialMessages: UIMessage[] = [{
  id: "welcome",
  role: "assistant",
  parts: [{
    type: "text",
    text: "CAD Agent Studio 已就绪。上传图片、STEP 或其他文件，或直接描述要生成/修改的模型。",
  }],
}];

const DEEPSEEK_PROVIDER_ID = "deepseek";

function uniqueId(prefix: string) {
  return `${prefix}_${Math.random().toString(36).slice(2, 10)}`;
}

function formatBytes(size: number) {
  if (size < 1024) return `${size} B`;
  if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KB`;
  return `${(size / 1024 / 1024).toFixed(1)} MB`;
}

function isStepFile(file: File) {
  return /\.(step|stp)$/i.test(file.name);
}

function extractCadResult(text: string) {
  const match = /::cad-result\{base64="([^"]+)"\}/.exec(text);
  if (!match) return null;
  try {
    return JSON.parse(atob(match[1])) as CadResult;
  } catch {
    return null;
  }
}

function setIfChanged(setter: (value: string | ((current: string) => string)) => void, nextValue = "") {
  setter((current) => (current === nextValue ? current : nextValue));
}

function UploadedAttachmentCard({ attachment }: { attachment: Attachment }) {
  return (
    <div className="flex min-w-0 items-center gap-2 px-1 py-1 text-[11px] text-[#a9b0b6]">
      {attachment.kind === "image" ? <FileImage className="size-3.5 text-[#d8b66c]" /> : <FileUp className="size-3.5 text-[#7dc8cf]" />}
      <span className="truncate">{attachment.name}</span>
      {attachment.kind === "step" ? (
        <span className="shrink-0 text-[10px] text-[#9fbfc3]">重建参考</span>
      ) : null}
      <span className="ml-auto shrink-0 text-[#969b9f]">{formatBytes(attachment.size)}</span>
    </div>
  );
}

function CadProgressCard({ data }: DataMessagePartProps<CadProgress>) {
  const icon = data.status === "running"
    ? <Loader2 className="size-3 animate-spin text-[#7dc8cf]" />
    : data.status === "success"
      ? <Check className="size-3 text-[#9cd67a]" />
      : data.status === "error"
        ? <CircleAlert className="size-3 text-[#f2b0b2]" />
        : <Loader2 className="size-3 animate-spin text-[#969b9f]" />;
  return (
    <div className="my-1.5 pl-1 text-[12px]">
      <div className="flex items-center gap-2 text-[#d7dcdd]">
        {icon}
        <span className="font-medium">{data.label}</span>
        <span className="text-[10px] uppercase tracking-wide text-[#777e84]">{data.status}</span>
      </div>
      {data.message ? <div className="ml-5 mt-0.5 text-[11px] leading-4 text-[#969b9f]">{data.message}</div> : null}
    </div>
  );
}

function CadArtifactCard({ data }: DataMessagePartProps<CadResult>) {
  return (
    <div className="my-2 pl-1 text-[12px]">
      <div className="flex items-center gap-2 font-medium text-[#dff5f6]">
        <Box className="size-3 text-[#7dc8cf]" />
        <span>CAD 结果</span>
      </div>
      <div className="ml-5 mt-1 grid gap-0.5 text-[11px] text-[#a9b0b6]">
        <div>{data.taskId}</div>
        {data.sourcePath ? <div>{data.sourcePath}</div> : null}
        {data.artifactPath ? <div>{data.artifactPath}</div> : null}
      </div>
      {data.artifactUrl ? (
        <a className="ml-5 mt-1 inline-flex text-[11px] text-[#8bd5dc] underline-offset-2 hover:underline" href={data.artifactUrl} download>
          下载 STEP
        </a>
      ) : null}
    </div>
  );
}

function CadErrorCard({ data }: DataMessagePartProps<CadError>) {
  return (
    <div className="my-1.5 pl-1 text-[12px] text-[#f2b0b2]">
      <div className="flex items-center gap-2 font-medium">
        <CircleAlert className="size-3" />
        <span>{data.stage ? `${data.stage} 失败` : "CAD 操作失败"}</span>
      </div>
      <div className="ml-5 mt-0.5 whitespace-pre-wrap text-[11px] leading-4">{data.message}</div>
    </div>
  );
}

function CadToolCallCard(props: ToolCallMessagePartProps) {
  const status = String(props.status?.type || props.status || "");
  return (
    <div className="my-1.5 pl-1 text-[12px]">
      <div className="flex items-center gap-2 text-[#d7dcdd]">
        {status.includes("running") || status.includes("input-streaming") ? <Loader2 className="size-3 animate-spin text-[#7dc8cf]" /> : <Wrench className="size-3 text-[#7dc8cf]" />}
        <span>{props.toolName}</span>
        <span className="text-[10px] uppercase tracking-wide text-[#777e84]">{status || "tool call"}</span>
      </div>
    </div>
  );
}

function DataFallback({ name, data }: DataMessagePartProps) {
  return (
    <div className="my-1.5 pl-1 text-[11px] text-[#969b9f]">
      {name}: {JSON.stringify(data).slice(0, 500)}
    </div>
  );
}

function MessageText() {
  return (
    <MessagePrimitive.Parts
      components={{
        Text: () => (
          <p className="whitespace-pre-wrap">
            <MessagePartPrimitive.Text />
          </p>
        ),
        data: {
          by_name: {
            "cad-progress": CadProgressCard,
            "cad-result": CadArtifactCard,
            "cad-error": CadErrorCard,
          },
          Fallback: DataFallback,
        },
        tools: {
          Fallback: CadToolCallCard,
        },
      }}
    />
  );
}

function UserMessage() {
  return (
    <MessagePrimitive.Root className="py-3 pl-8 pr-1 text-[12px] leading-5">
      <div className="mb-2 flex items-center gap-2 text-[11px] font-semibold text-[#969b9f]">
        <Sparkles className="size-3.5" />
        You
      </div>
      <MessageText />
    </MessagePrimitive.Root>
  );
}

function AssistantMessage() {
  return (
    <MessagePrimitive.Root className="py-3 pr-1 text-[12px] leading-5">
      <div className="mb-2 flex items-center gap-2 text-[11px] font-semibold text-[#969b9f]">
        <Bot className="size-3.5 text-[#7dc8cf]" />
        Agent
      </div>
      <MessageText />
      <MessagePrimitive.Error>
        <div className="mt-2 text-[11px] text-[#f2b0b2]">
          Agent 请求失败，请检查服务端日志和模型配置。
        </div>
      </MessagePrimitive.Error>
    </MessagePrimitive.Root>
  );
}

function AnyMessage() {
  const message = useMessage();
  if (message.role === "user") return <UserMessage />;
  if (message.role === "assistant") return <AssistantMessage />;
  return null;
}

export function AgentStudio() {
  const [config, setConfig] = useState<PublicConfig | null>(null);
  const [provider, setProvider] = useState(DEEPSEEK_PROVIDER_ID);
  const [model, setModel] = useState("");
  const [attachments, setAttachments] = useState<Attachment[]>([]);
  const [pendingViewerContexts, setPendingViewerContexts] = useState<ViewerContext[]>([]);
  const [currentTaskId, setCurrentTaskId] = useState("");
  const [previewUrl, setPreviewUrl] = useState("");
  const [artifactUrl, setArtifactUrl] = useState("");
  const [artifactPath, setArtifactPath] = useState("");
  const [viewerAssetUrl, setViewerAssetUrl] = useState("");
  const [activeEditToolId, setActiveEditToolId] = useState("");
  const [aiSelectionMode, setAiSelectionMode] = useState<AiSelectionMode>("point");
  const [uploading, setUploading] = useState(false);
  const [draggingUpload, setDraggingUpload] = useState(false);
  const [error, setError] = useState("");
  const dragDepthRef = useRef(0);
  const fileInputRef = useRef<HTMLInputElement | null>(null);

  const applyCadResult = useCallback((result: CadResult) => {
    if (!result.taskId) return;
    setIfChanged(setCurrentTaskId, result.taskId);
    setIfChanged(setArtifactPath, result.artifactPath || "");
    setIfChanged(setArtifactUrl, result.artifactUrl || "");
    setIfChanged(setPreviewUrl, result.previewUrl || "");
    setIfChanged(setViewerAssetUrl, result.viewerAssetUrl || "");
    setActiveEditToolId("");
    setAiSelectionMode("point");
  }, []);

  const handleData = useCallback((dataPart: { type?: string; data?: unknown }) => {
    if (dataPart.type === "data-cad-result") {
      applyCadResult(dataPart.data as CadResult);
      setPendingViewerContexts([]);
      return;
    }
    if (dataPart.type === "data-cad-error") {
      const data = dataPart.data as CadError;
      setError(data.message || "CAD 生成失败");
    }
  }, [applyCadResult]);

  const transport = useMemo(() => new AssistantChatTransport<UIMessage>({
    api: "/api/chat",
    body: {
      provider: DEEPSEEK_PROVIDER_ID,
      model,
      attachments,
      viewerContext: pendingViewerContexts,
      selectedTaskId: currentTaskId,
    },
  }), [attachments, currentTaskId, model, pendingViewerContexts]);

  const runtime = useChatRuntime({
    messages: initialMessages,
    transport,
    onData: handleData,
    onError: (nextError) => {
      setError(nextError instanceof Error ? nextError.message : "Agent 回复失败");
    },
  });

  useEffect(() => {
    void fetch("/api/config")
      .then((response) => response.json())
      .then((nextConfig: PublicConfig) => {
        setConfig(nextConfig);
        setProvider(DEEPSEEK_PROVIDER_ID);
        const deepseekProvider = nextConfig.providers[DEEPSEEK_PROVIDER_ID];
        setModel(deepseekProvider?.models?.default || "");
      })
      .catch((nextError) => {
        setError(nextError instanceof Error ? nextError.message : "配置加载失败");
      });
  }, []);

  const providerModels = useMemo(() => {
    const models = config?.providers?.[DEEPSEEK_PROVIDER_ID]?.models || {};
    return Object.entries(models);
  }, [config]);

  const pushViewerContext = useCallback((type: string, payload: Record<string, unknown>) => {
    if (!currentTaskId) {
      setError("请先生成或上传一个模型，再使用右侧工具栏。");
      return;
    }
    const context: ViewerContext = {
      id: uniqueId("viewer"),
      type,
      createdAt: new Date().toISOString(),
      payload,
    };
    setPendingViewerContexts((current) => {
      const nextCurrent = type === "cadViewer.selection"
        ? current.filter((item) => item.type !== "cadViewer.selection")
        : current;
      return [context, ...nextCurrent].slice(0, 8);
    });
  }, [currentTaskId]);

  const handleEditIntent = useCallback((intent: Record<string, unknown>) => {
    if (!currentTaskId) return;
    pushViewerContext("cadViewer.editIntent", {
      ...intent,
      viewerState: {
        viewerAssetUrl,
        renderer: "cad-viewer.ported",
      },
    });
  }, [currentTaskId, pushViewerContext, viewerAssetUrl]);

  const handleGeometrySelection = useCallback((selectionIntent: Record<string, unknown>) => {
    if (!currentTaskId) return;
    pushViewerContext("cadViewer.selection", {
      ...selectionIntent,
      viewerState: {
        viewerAssetUrl,
        renderer: "cad-viewer.ported",
      },
    });
  }, [currentTaskId, pushViewerContext, viewerAssetUrl]);

  const handleUpload = useCallback(async (files: FileList | null) => {
    if (!files?.length) return;
    setUploading(true);
    setError("");
    try {
      for (const file of Array.from(files)) {
        const form = new FormData();
        form.append("file", file);
        const stepUpload = isStepFile(file);
        if (currentTaskId && !stepUpload) {
          form.append("taskId", currentTaskId);
        }
        const response = await fetch("/api/upload", {
          method: "POST",
          body: form,
        });
        const attachment = await response.json();
        if (!response.ok) {
          throw new Error(attachment.error || `${file.name} 上传失败`);
        }
        setAttachments((current) => [attachment, ...current]);
        if (attachment.kind === "step") {
          continue;
        }
        setCurrentTaskId(attachment.taskId);
        if (attachment.artifactUrl) {
          setArtifactUrl(attachment.artifactUrl);
          setArtifactPath(attachment.path);
        }
        if (attachment.previewUrl) {
          setPreviewUrl(attachment.previewUrl);
        }
        if (attachment.viewerAssetUrl) {
          setViewerAssetUrl(attachment.viewerAssetUrl);
        }
        if (attachment.previewError) {
          setError(`STEP 已上传，但预览网格生成失败：${attachment.previewError}`);
        }
      }
    } catch (nextError) {
      setError(nextError instanceof Error ? nextError.message : "上传失败");
    } finally {
      setUploading(false);
      if (fileInputRef.current) {
        fileInputRef.current.value = "";
      }
    }
  }, [currentTaskId]);

  const handleDragEnter = useCallback((event: DragEvent<HTMLElement>) => {
    event.preventDefault();
    event.stopPropagation();
    dragDepthRef.current += 1;
    if (event.dataTransfer.types.includes("Files")) {
      setDraggingUpload(true);
    }
  }, []);

  const handleDragOver = useCallback((event: DragEvent<HTMLElement>) => {
    event.preventDefault();
    event.stopPropagation();
    event.dataTransfer.dropEffect = "copy";
    if (event.dataTransfer.types.includes("Files")) {
      setDraggingUpload(true);
    }
  }, []);

  const handleDragLeave = useCallback((event: DragEvent<HTMLElement>) => {
    event.preventDefault();
    event.stopPropagation();
    dragDepthRef.current = Math.max(0, dragDepthRef.current - 1);
    if (dragDepthRef.current === 0) {
      setDraggingUpload(false);
    }
  }, []);

  const handleDrop = useCallback((event: DragEvent<HTMLElement>) => {
    event.preventDefault();
    event.stopPropagation();
    dragDepthRef.current = 0;
    setDraggingUpload(false);
    void handleUpload(event.dataTransfer.files);
  }, [handleUpload]);

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <main className="flex h-screen min-h-0 flex-col bg-[#111315] text-[#e8e8e3]">
        <header className="flex h-12 shrink-0 items-center justify-between border-b border-[#2d3339] bg-[#15181b] px-3">
          <div className="flex min-w-0 items-center gap-2">
            <Box className="size-4 text-[#7dc8cf]" aria-hidden="true" />
            <div className="truncate text-sm font-semibold">CAD Agent Studio</div>
            {currentTaskId ? (
              <div className="hidden truncate rounded border border-[#2d3339] px-2 py-1 text-[11px] text-[#969b9f] sm:block">
                {currentTaskId}
              </div>
            ) : null}
          </div>
          <div className="flex items-center gap-2">
            <div className="flex h-8 items-center gap-1.5 rounded border border-[#2d3339] bg-[#1e2226] px-2 text-[11px] text-[#c9cccf]">
              <Settings2 className="size-3.5" aria-hidden="true" />
              <span>DeepSeek</span>
            </div>
            <label className="hidden h-8 items-center gap-1.5 rounded border border-[#2d3339] bg-[#1e2226] px-2 text-[11px] text-[#c9cccf] md:flex">
              <Sparkles className="size-3.5" aria-hidden="true" />
              <select
                className="max-w-40 bg-transparent outline-none"
                value={model}
                onChange={(event) => setModel(event.target.value)}
              >
                {providerModels.map(([label, value]) => (
                  <option key={`${label}:${value}`} value={value}>{label}: {value}</option>
                ))}
              </select>
            </label>
          </div>
        </header>

        <div className="grid min-h-0 flex-1 grid-cols-1 md:grid-cols-[420px_minmax(0,1fr)]">
          <section
            className="relative flex min-h-0 flex-col border-r border-[#2d3339] bg-[#171a1d]"
            onDragEnter={handleDragEnter}
            onDragOver={handleDragOver}
            onDragLeave={handleDragLeave}
            onDrop={handleDrop}
          >
            {draggingUpload ? (
              <div className="pointer-events-none absolute inset-2 z-40 grid place-items-center rounded-md border border-dashed border-[#7dc8cf] bg-[#101315]/88 text-center shadow-[0_0_0_999px_rgba(17,19,21,0.42)] backdrop-blur-sm">
                <div>
                  <FileUp className="mx-auto mb-2 size-7 text-[#7dc8cf]" />
                  <div className="text-sm font-semibold text-[#f2f4f4]">拖放文件上传</div>
                  <div className="mt-1 text-[12px] text-[#a9b0b6]">支持图片、STEP/STP 和其他参考文件</div>
                </div>
              </div>
            ) : null}
            <div className="flex h-10 shrink-0 items-center justify-between border-b border-[#2d3339] px-3">
              <div className="flex items-center gap-2 text-xs font-semibold">
                <MessageSquare className="size-4 text-[#7dc8cf]" aria-hidden="true" />
                Agent
              </div>
              <input
                ref={fileInputRef}
                type="file"
                multiple
                className="hidden"
                onChange={(event) => void handleUpload(event.target.files)}
              />
            </div>

            <ThreadPrimitive.Root className="flex min-h-0 flex-1 flex-col">
              <ThreadPrimitive.Viewport className="scrollbar-thin min-h-0 flex-1 overflow-y-auto px-3 py-3">
                {error ? (
                  <div className="mb-3 flex gap-2 rounded border border-[#7a3438] bg-[#2a171a] p-2 text-[12px] text-[#f2b0b2]">
                    <CircleAlert className="mt-0.5 size-4 shrink-0" />
                    <span>{error}</span>
                  </div>
                ) : null}

                {attachments.length ? (
                  <div className="mb-3 border-b border-[#252b30] pb-2">
                    {attachments.slice(0, 5).map((attachment) => (
                      <UploadedAttachmentCard key={attachment.id} attachment={attachment} />
                    ))}
                  </div>
                ) : null}

                <div className="divide-y divide-[#252b30]">
                  <ThreadPrimitive.Messages components={{ Message: AnyMessage, UserMessage, AssistantMessage }} />
                </div>
              </ThreadPrimitive.Viewport>

              <div className="shrink-0 border-t border-[#2d3339] p-3">
                <ComposerPrimitive.Root className="grid gap-2">
                  <ComposerPrimitive.Input
                    placeholder="描述要生成或修改的 CAD 模型..."
                    submitMode="enter"
                    className="max-h-44 min-h-24 w-full resize-none rounded border border-[#2d3339] bg-[#111315] p-2 text-[12px] leading-5 text-[#e8e8e3] outline-none focus:border-[#7dc8cf]"
                  />
                  <div className="flex items-center justify-between gap-2">
                    <div className="flex min-w-0 items-center gap-1.5 text-[11px] text-[#969b9f]">
                      <Check className="size-3.5" />
                      <span className="truncate">Enter 发送，Shift + Enter 换行</span>
                    </div>
                    <div className="flex items-center gap-2">
                      <button
                        type="button"
                        className="inline-flex h-8 items-center gap-1 rounded border border-[#2d3339] px-2 text-[11px] text-[#c9cccf] hover:bg-[#22272c] disabled:opacity-60"
                        onClick={() => fileInputRef.current?.click()}
                        disabled={uploading}
                        title="上传文件"
                        aria-label="上传文件"
                      >
                        {uploading ? <Loader2 className="size-3.5 animate-spin" /> : <Paperclip className="size-3.5" />}
                        上传
                      </button>
                      <ComposerPrimitive.Send className="inline-flex h-8 items-center gap-1.5 rounded bg-[#7dc8cf] px-3 text-[12px] font-semibold text-[#101315] hover:bg-[#91d4da] disabled:cursor-not-allowed disabled:opacity-50">
                        <Send className="size-3.5" />
                        发送
                      </ComposerPrimitive.Send>
                    </div>
                  </div>
                </ComposerPrimitive.Root>
              </div>
            </ThreadPrimitive.Root>
          </section>

          <section className="relative min-h-[42vh] bg-[#0d0f11] md:min-h-0">
            {viewerAssetUrl ? (
              <CadViewerPreview
                taskId={currentTaskId}
                artifactPath={artifactPath}
                viewerAssetUrl={viewerAssetUrl}
                activeEditToolId={activeEditToolId}
                aiSelectionMode={aiSelectionMode}
                onSelectTool={setActiveEditToolId}
                onSelectionModeChange={setAiSelectionMode}
                onEditIntent={handleEditIntent}
                onGeometrySelection={handleGeometrySelection}
              />
            ) : (
              <div className="flex h-full min-h-[42vh] items-center justify-center p-6">
                <div className="max-w-md rounded border border-[#2d3339] bg-[#15181b] p-5 text-center">
                  <Square className="mx-auto mb-3 size-7 text-[#7dc8cf]" />
                  <div className="text-sm font-semibold">3D 预览等待模型</div>
                  <div className="mt-2 text-[12px] leading-5 text-[#969b9f]">
                    上传文件会进入 Agent 上下文。生成、重建或修改完成后会切换到真实输出结果。
                  </div>
                </div>
              </div>
            )}
            {artifactUrl ? (
              <a
                className="absolute right-24 top-4 z-30 inline-flex h-7 items-center gap-1 rounded-md border border-[#2d3339] bg-[#15181b]/95 px-2 text-[11px] text-[#c9cccf] hover:bg-[#20252a]"
                href={artifactUrl}
                download
              >
                <FileUp className="size-3.5" />
                STEP
              </a>
            ) : null}
          </section>
        </div>
      </main>
    </AssistantRuntimeProvider>
  );
}