Pre_Seg_Server/src/components/VideoWorkspace.tsx

import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { useStore } from '../store/useStore';
import {
  annotationToMask,
  buildAnnotationPayload,
  deleteAnnotation,
  exportCoco,
  getProjectAnnotations,
  getProjectFrames,
  getTask,
  getTemplates,
  parseMedia,
  saveAnnotation,
  updateAnnotation,
} from '../lib/api';
import { CanvasArea } from './CanvasArea';
import { ToolsPalette } from './ToolsPalette';
import { OntologyInspector } from './OntologyInspector';
import { FrameTimeline } from './FrameTimeline';
import { ModelStatusBadge } from './ModelStatusBadge';
import type { Frame } from '../store/useStore';

function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export function VideoWorkspace({ onNavigateToAI }: { onNavigateToAI?: () => void }) {
  const activeTool = useStore((state) => state.activeTool);
  const setActiveTool = useStore((state) => state.setActiveTool);
  const currentProject = useStore((state) => state.currentProject);
  const frames = useStore((state) => state.frames);
  const currentFrameIndex = useStore((state) => state.currentFrameIndex);
  const masks = useStore((state) => state.masks);
  const activeTemplateId = useStore((state) => state.activeTemplateId);
  const setFrames = useStore((state) => state.setFrames);
  const setCurrentFrame = useStore((state) => state.setCurrentFrame);
  const setMasks = useStore((state) => state.setMasks);
  const [isSaving, setIsSaving] = useState(false);
  const [isExporting, setIsExporting] = useState(false);
  const [statusMessage, setStatusMessage] = useState('');

  const hydrateSavedAnnotations = useCallback(async (projectId: string, projectFrames: Frame[]) => {
    const frameById = new Map(projectFrames.map((frame) => [frame.id, frame]));
    const annotations = await getProjectAnnotations(projectId);
    const savedMasks = annotations
      .map((annotation) => {
        const frame = annotation.frame_id ? frameById.get(String(annotation.frame_id)) : null;
        return frame ? annotationToMask(annotation, frame) : null;
      })
      .filter((mask): mask is NonNullable<typeof mask> => Boolean(mask));
    setMasks(savedMasks);
  }, [setMasks]);

  useEffect(() => {
    if (!currentProject?.id) return;
    let cancelled = false;

    const loadFrames = async () => {
      try {
        const data = await getProjectFrames(String(currentProject.id));
        if (cancelled) return;

        if (data.length === 0 && currentProject.video_path) {
          // No frames yet but video exists -> queue parsing and poll the task.
          try {
            const task = await parseMedia(String(currentProject.id));
            if (cancelled) return;
            setStatusMessage(`解析任务已入队 #${task.id}`);
            let completed = false;
            for (let attempt = 0; attempt < 60; attempt += 1) {
              const freshTask = await getTask(task.id);
              if (cancelled) return;
              setStatusMessage(freshTask.message || `解析进度 ${freshTask.progress}%`);
              if (freshTask.status === 'success') {
                completed = true;
                break;
              }
              if (freshTask.status === 'failed') {
                setStatusMessage(freshTask.error || '解析任务失败');
                return;
              }
              await sleep(2000);
            }
            if (!completed) {
              setStatusMessage('解析仍在后台运行，可稍后刷新工作区');
              return;
            }
            const fresh = await getProjectFrames(String(currentProject.id));
            if (cancelled) return;
            const mappedFrames = fresh.map((f) => ({
              id: String(f.id),
              projectId: String(f.project_id),
              index: f.frame_index,
              url: f.image_url,
              width: f.width ?? 0,
              height: f.height ?? 0,
            }));
            setFrames(mappedFrames);
            setCurrentFrame(0);
            await hydrateSavedAnnotations(String(currentProject.id), mappedFrames);
          } catch (err) {
            console.error('Parse failed:', err);
          }
        } else {
          const mappedFrames = data.map((f) => ({
            id: String(f.id),
            projectId: String(f.project_id),
            index: f.frame_index,
            url: f.image_url,
            width: f.width ?? 0,
            height: f.height ?? 0,
          }));
          setFrames(mappedFrames);
          setCurrentFrame(0);
          await hydrateSavedAnnotations(String(currentProject.id), mappedFrames);
        }
      } catch (err) {
        console.error('Failed to load frames:', err);
      }
    };

    loadFrames();
    return () => { cancelled = true; };
  }, [currentProject?.id, currentProject?.video_path, hydrateSavedAnnotations, setFrames, setCurrentFrame]);

  const templates = useStore((state) => state.templates);
  const setTemplates = useStore((state) => state.setTemplates);

  useEffect(() => {
    if (templates.length === 0) {
      getTemplates().then((data) => setTemplates(data)).catch(console.error);
    }
  }, [templates.length, setTemplates]);

  const currentFrame = frames[currentFrameIndex] || null;
  const frameById = useMemo(() => new Map(frames.map((frame) => [frame.id, frame])), [frames]);
  const projectFrameIds = useMemo(() => new Set(frames.map((frame) => frame.id)), [frames]);

  const savePendingAnnotations = useCallback(async ({ silent = false } = {}) => {
    if (!currentProject?.id) return 0;
    const projectMasks = masks.filter((mask) => projectFrameIds.has(mask.frameId));
    const pendingMasks = projectMasks.filter((mask) => !mask.annotationId);
    const dirtyMasks = projectMasks.filter((mask) => mask.annotationId && mask.saveStatus === 'dirty');
    if (pendingMasks.length === 0 && dirtyMasks.length === 0) {
      if (!silent) setStatusMessage('没有待保存标注');
      return 0;
    }

    setIsSaving(true);
    setStatusMessage('正在保存标注...');
    try {
      const createPayloads = pendingMasks
        .map((mask) => {
          const frame = frameById.get(mask.frameId);
          return frame ? buildAnnotationPayload(currentProject.id, mask, frame, activeTemplateId) : null;
        })
        .filter((payload): payload is NonNullable<typeof payload> => Boolean(payload));

      const updatePayloads = dirtyMasks
        .map((mask) => {
          const frame = frameById.get(mask.frameId);
          const payload = frame ? buildAnnotationPayload(currentProject.id, mask, frame, activeTemplateId) : null;
          if (!payload || !mask.annotationId) return null;
          const updatePayload = {
            template_id: payload.template_id,
            mask_data: payload.mask_data,
            points: payload.points,
            bbox: payload.bbox,
          };
          return { annotationId: mask.annotationId, payload: updatePayload };
        })
        .filter((item): item is NonNullable<typeof item> => Boolean(item));

      if (createPayloads.length === 0 && updatePayloads.length === 0) {
        setStatusMessage('没有可保存的标注数据');
        return 0;
      }

      await Promise.all([
        ...createPayloads.map((payload) => saveAnnotation(payload)),
        ...updatePayloads.map(({ annotationId, payload }) => updateAnnotation(annotationId, payload)),
      ]);
      await hydrateSavedAnnotations(currentProject.id, frames);
      const savedCount = createPayloads.length + updatePayloads.length;
      setStatusMessage(`已保存 ${savedCount} 个标注`);
      return savedCount;
    } catch (err) {
      console.error('Save annotations failed:', err);
      setStatusMessage('保存失败，请检查后端服务');
      throw err;
    } finally {
      setIsSaving(false);
    }
  }, [activeTemplateId, currentProject?.id, frameById, frames, hydrateSavedAnnotations, masks, projectFrameIds]);

  const handleClearCurrentFrameMasks = useCallback(async () => {
    if (!currentFrame) return;
    const frameMasks = masks.filter((mask) => mask.frameId === currentFrame.id);
    const annotationIds = frameMasks
      .map((mask) => mask.annotationId)
      .filter((annotationId): annotationId is string => Boolean(annotationId));

    setIsSaving(true);
    setStatusMessage(annotationIds.length > 0 ? '正在删除已保存标注...' : '正在清空本帧遮罩...');
    try {
      await Promise.all(annotationIds.map((annotationId) => deleteAnnotation(annotationId)));
      setMasks(masks.filter((mask) => mask.frameId !== currentFrame.id));
      setStatusMessage(annotationIds.length > 0
        ? `已删除 ${annotationIds.length} 个后端标注`
        : '已清空本帧未保存遮罩');
    } catch (err) {
      console.error('Delete annotations failed:', err);
      setStatusMessage('删除失败，请检查后端服务');
    } finally {
      setIsSaving(false);
    }
  }, [currentFrame, masks, setMasks]);

  const handleSave = async () => {
    try {
      await savePendingAnnotations();
    } catch {
      // status message is set in savePendingAnnotations
    }
  };

  const handleExport = async () => {
    if (!currentProject?.id) return;
    setIsExporting(true);
    setStatusMessage('正在准备导出...');
    try {
      await savePendingAnnotations({ silent: true });
      const blob = await exportCoco(currentProject.id);
      const url = URL.createObjectURL(blob);
      const link = document.createElement('a');
      link.href = url;
      link.download = `project_${currentProject.id}_coco.json`;
      document.body.appendChild(link);
      link.click();
      link.remove();
      URL.revokeObjectURL(url);
      setStatusMessage('COCO JSON 已导出');
    } catch (err) {
      console.error('Export failed:', err);
      setStatusMessage('导出失败，请检查后端服务');
    } finally {
      setIsExporting(false);
    }
  };

  return (
    <div className="w-full h-full flex flex-col bg-[#0a0a0a]">
      {/* Top Header / Status bar */}
      <div className="h-14 border-b border-white/5 bg-[#111] flex items-center justify-between px-6 shrink-0">
        <div className="flex items-center gap-4">
          <h2 className="text-xs font-semibold uppercase tracking-widest text-gray-400">核心分割工作区</h2>
          <div className="h-4 w-px bg-white/10"></div>
          <span className="text-sm text-white font-mono">{currentProject?.name || '未选择项目'}</span>
        </div>
        <div className="flex items-center gap-3">
          {statusMessage && (
            <span className="text-[10px] text-gray-500 font-mono max-w-48 truncate" title={statusMessage}>
              {statusMessage}
            </span>
          )}
          <ModelStatusBadge />
          <button
            onClick={handleExport}
            disabled={!currentProject?.id || isExporting || isSaving}
            className="px-4 py-1.5 bg-white/5 hover:bg-white/10 border border-white/10 rounded-md text-xs transition-colors text-white disabled:opacity-40 disabled:cursor-not-allowed"
          >
            {isExporting ? '导出中...' : '导出 JSON 标注集'}
          </button>
          <button
            onClick={handleSave}
            disabled={!currentProject?.id || isSaving || isExporting}
            className="px-4 py-1.5 bg-cyan-600 hover:bg-cyan-500 text-white text-xs font-medium rounded-md transition-shadow shadow-lg shadow-cyan-900/20 disabled:opacity-40 disabled:cursor-not-allowed"
          >
            {isSaving ? '保存中...' : '结构化归档保存'}
          </button>
        </div>
      </div>

      {/* Main Workspace Area */}
      <div className="flex-1 flex overflow-hidden">
        <ToolsPalette activeTool={activeTool} setActiveTool={setActiveTool} onTriggerAI={onNavigateToAI} />
        
        <div className="flex-1 relative flex items-center justify-center p-8 bg-[#151515] overflow-hidden">
          <div className="relative w-full h-full bg-[#1e1e1e] border border-white/5 shadow-2xl rounded-sm">
            <CanvasArea activeTool={activeTool} frame={currentFrame} onClearMasks={handleClearCurrentFrameMasks} />
          </div>
        </div>
        
        <OntologyInspector />
      </div>

      {/* Bottom Timeline */}
      <FrameTimeline />
    </div>
  );
}