import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { useStore } from '../store/useStore';
import {
  annotationToMask,
  buildAnnotationPayload,
  deleteAnnotation,
  exportCoco,
  exportMasks,
  getProjectAnnotations,
  getProjectFrames,
  getTemplates,
  importGtMask,
  propagateMasks,
  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, Mask } from '../store/useStore';

type PropagationDirection = 'forward' | 'backward';

export function VideoWorkspace({ onNavigateToAI }: { onNavigateToAI?: () => void }) {
  const gtMaskInputRef = React.useRef<HTMLInputElement>(null);
  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 maskHistory = useStore((state) => state.maskHistory);
  const maskFuture = useStore((state) => state.maskFuture);
  const activeTemplateId = useStore((state) => state.activeTemplateId);
  const aiModel = useStore((state) => state.aiModel);
  const selectedMaskIds = useStore((state) => state.selectedMaskIds);
  const latestSelectedMaskIdsRef = React.useRef<string[]>(selectedMaskIds);
  if (selectedMaskIds.length > 0) {
    latestSelectedMaskIdsRef.current = selectedMaskIds;
  }
  const setFrames = useStore((state) => state.setFrames);
  const setCurrentFrame = useStore((state) => state.setCurrentFrame);
  const setMasks = useStore((state) => state.setMasks);
  const setSelectedMaskIds = useStore((state) => state.setSelectedMaskIds);
  const undoMasks = useStore((state) => state.undoMasks);
  const redoMasks = useStore((state) => state.redoMasks);
  const [isSaving, setIsSaving] = useState(false);
  const [isExporting, setIsExporting] = useState(false);
  const [isImportingGt, setIsImportingGt] = useState(false);
  const [isPropagating, setIsPropagating] = useState(false);
  const [statusMessage, setStatusMessage] = useState('');
  const [propagationStartFrame, setPropagationStartFrame] = useState(1);
  const [propagationEndFrame, setPropagationEndFrame] = useState(1);

  const hydrateSavedAnnotations = useCallback(async (
    projectId: string,
    projectFrames: Frame[],
    preserveSelectedIds: string[] = [],
    excludeUnsavedMaskIds: string[] = [],
  ) => {
    const frameById = new Map(projectFrames.map((frame) => [frame.id, frame]));
    const projectFrameIds = new Set(projectFrames.map((frame) => frame.id));
    const excludedDraftIds = new Set(excludeUnsavedMaskIds);
    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));
    const currentMasks = useStore.getState().masks;
    const unsavedMasks = currentMasks.filter((mask) => (
      !projectFrameIds.has(mask.frameId) || (!mask.annotationId && !excludedDraftIds.has(mask.id))
    ));
    const mergedMasks = [...unsavedMasks, ...savedMasks];
    setMasks(mergedMasks);
    if (preserveSelectedIds.length > 0) {
      const mergedMaskIds = new Set(mergedMasks.map((mask) => mask.id));
      const nextSelectedIds = preserveSelectedIds.filter((id) => mergedMaskIds.has(id));
      if (nextSelectedIds.length > 0) {
        setSelectedMaskIds(nextSelectedIds);
      }
    }
  }, [setMasks, setSelectedMaskIds]);

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

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

        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,
          timestampMs: f.timestamp_ms ?? undefined,
          sourceFrameNumber: f.source_frame_number ?? undefined,
        }));
        setFrames(mappedFrames);
        setCurrentFrame(0);
        if (mappedFrames.length === 0) {
          setMasks([]);
          if (currentProject.status === 'parsing') {
            setStatusMessage('生成帧任务正在后台运行，可在 Dashboard 查看进度');
          } else if (currentProject.video_path) {
            setStatusMessage('该项目已导入视频但尚未生成帧，请在项目库点击“生成帧”');
          } else {
            setStatusMessage('当前项目没有可显示帧');
          }
          return;
        }
        setStatusMessage('');
        await hydrateSavedAnnotations(String(currentProject.id), mappedFrames, selectedIdsBeforeLoad);
      } 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 totalFrames = frames.length;
  const frameById = useMemo(() => new Map(frames.map((frame) => [frame.id, frame])), [frames]);
  const projectFrameIds = useMemo(() => new Set(frames.map((frame) => frame.id)), [frames]);
  const currentFrameNumber = currentFrameIndex + 1;

  useEffect(() => {
    if (totalFrames === 0) {
      setPropagationStartFrame(1);
      setPropagationEndFrame(1);
      return;
    }
    setPropagationStartFrame(currentFrameNumber);
    setPropagationEndFrame(Math.min(totalFrames, currentFrameNumber + 29));
  }, [currentFrameNumber, totalFrames]);

  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 createItems = pendingMasks
        .map((mask) => {
          const frame = frameById.get(mask.frameId);
          const payload = frame ? buildAnnotationPayload(currentProject.id, mask, frame, activeTemplateId) : null;
          return payload ? { maskId: mask.id, payload } : null;
        })
        .filter((item): item is NonNullable<typeof item> => Boolean(item));

      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 (createItems.length === 0 && updatePayloads.length === 0) {
        setStatusMessage('没有可保存的标注数据');
        return 0;
      }

      await Promise.all([
        ...createItems.map(({ payload }) => saveAnnotation(payload)),
        ...updatePayloads.map(({ annotationId, payload }) => updateAnnotation(annotationId, payload)),
      ]);
      await hydrateSavedAnnotations(
        currentProject.id,
        frames,
        useStore.getState().selectedMaskIds,
        createItems.map(({ maskId }) => maskId),
      );
      const savedCount = createItems.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 handleDeleteMaskAnnotations = useCallback(async (annotationIds: string[]) => {
    if (annotationIds.length === 0) return;
    try {
      await Promise.all(annotationIds.map((annotationId) => deleteAnnotation(annotationId)));
      setStatusMessage(`已删除 ${annotationIds.length} 个被合并标注`);
    } catch (err) {
      console.error('Delete merged annotations failed:', err);
      setStatusMessage('合并后删除原标注失败，请检查后端服务');
      throw err;
    }
  }, []);

  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);
    }
  };

  const downloadBlob = (blob: Blob, filename: string) => {
    const url = URL.createObjectURL(blob);
    const link = document.createElement('a');
    link.href = url;
    link.download = filename;
    document.body.appendChild(link);
    link.click();
    link.remove();
    URL.revokeObjectURL(url);
  };

  const handleExportMasks = async () => {
    if (!currentProject?.id) return;
    setIsExporting(true);
    setStatusMessage('正在准备导出语义 Mask ZIP...');
    try {
      await savePendingAnnotations({ silent: true });
      const blob = await exportMasks(currentProject.id);
      downloadBlob(blob, `project_${currentProject.id}_masks.zip`);
      setStatusMessage('PNG Mask ZIP 已导出');
    } catch (err) {
      console.error('Mask export failed:', err);
      setStatusMessage('Mask 导出失败，请检查后端服务');
    } finally {
      setIsExporting(false);
    }
  };

  const handleImportGtMask = async (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0];
    if (!file || !currentProject?.id || !currentFrame?.id) return;
    setIsImportingGt(true);
    setStatusMessage('正在导入 GT Mask...');
    try {
      const imported = await importGtMask(file, currentProject.id, currentFrame.id);
      await hydrateSavedAnnotations(currentProject.id, frames);
      setStatusMessage(`已导入 ${imported.length} 个 GT 区域`);
    } catch (err) {
      console.error('GT mask import failed:', err);
      setStatusMessage('GT Mask 导入失败，请检查文件或后端服务');
    } finally {
      setIsImportingGt(false);
      event.target.value = '';
    }
  };

  const clampFrameNumber = useCallback((value: number) => {
    if (totalFrames <= 0) return 1;
    return Math.min(Math.max(value, 1), totalFrames);
  }, [totalFrames]);

  const buildSeedPayload = useCallback((seedMask: Mask) => {
    if (!currentProject?.id || !currentFrame) return null;
    const seedPayload = buildAnnotationPayload(currentProject.id, seedMask, currentFrame, activeTemplateId);
    if (!seedPayload?.mask_data?.polygons?.length && !seedPayload?.bbox) {
      return null;
    }
    return {
      polygons: seedPayload.mask_data?.polygons,
      bbox: seedPayload.bbox,
      points: seedPayload.points,
      label: seedPayload.mask_data?.label,
      color: seedPayload.mask_data?.color,
      class_metadata: seedPayload.mask_data?.class,
      template_id: seedPayload.template_id,
    };
  }, [activeTemplateId, currentFrame, currentProject?.id]);

  const handleAutoPropagate = async () => {
    if (!currentProject?.id || !currentFrame?.id) return;
    const seedMasks = masks.filter((mask) => mask.frameId === currentFrame.id);
    if (seedMasks.length === 0) {
      setStatusMessage('请先在当前参考帧创建或保存至少一个 mask');
      return;
    }

    const startFrameNumber = clampFrameNumber(propagationStartFrame);
    const endFrameNumber = clampFrameNumber(propagationEndFrame);
    const rangeStartIndex = Math.min(startFrameNumber, endFrameNumber) - 1;
    const rangeEndIndex = Math.max(startFrameNumber, endFrameNumber) - 1;
    const propagationDirections: Array<{ direction: PropagationDirection; maxFrames: number }> = [];
    if (rangeStartIndex < currentFrameIndex) {
      propagationDirections.push({
        direction: 'backward',
        maxFrames: currentFrameIndex - rangeStartIndex + 1,
      });
    }
    if (rangeEndIndex > currentFrameIndex) {
      propagationDirections.push({
        direction: 'forward',
        maxFrames: rangeEndIndex - currentFrameIndex + 1,
      });
    }
    if (propagationDirections.length === 0) {
      setStatusMessage('传播范围只包含当前帧，请选择前后至少一帧');
      return;
    }

    const seeds = seedMasks
      .map((mask) => ({ mask, seed: buildSeedPayload(mask) }))
      .filter((item): item is { mask: Mask; seed: NonNullable<ReturnType<typeof buildSeedPayload>> } => Boolean(item.seed));
    if (seeds.length === 0) {
      setStatusMessage('所选区域缺少可传播的 polygon 或 bbox');
      return;
    }

    setIsPropagating(true);
    setStatusMessage(`${aiModel.toUpperCase()} 正在以第 ${currentFrameNumber} 帧为参考，自动传播 ${seeds.length} 个 mask 到第 ${rangeStartIndex + 1}-${rangeEndIndex + 1} 帧...`);
    try {
      let createdCount = 0;
      let processedCount = 0;
      for (const { seed } of seeds) {
        for (const { direction, maxFrames } of propagationDirections) {
          const result = await propagateMasks({
            project_id: Number(currentProject.id),
            frame_id: Number(currentFrame.id),
            model: aiModel,
            direction,
            max_frames: maxFrames,
            include_source: false,
            save_annotations: true,
            seed,
          });
          createdCount += result.created_annotation_count;
          processedCount += result.processed_frame_count;
        }
      }
      await hydrateSavedAnnotations(currentProject.id, frames);
      setStatusMessage(`已自动传播 ${seeds.length} 个参考 mask，处理 ${processedCount} 帧次，保存 ${createdCount} 个区域`);
    } catch (err) {
      console.error('Propagation failed:', err);
      setStatusMessage('传播失败，请检查模型状态或后端日志');
    } finally {
      setIsPropagating(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 />
          <input
            ref={gtMaskInputRef}
            type="file"
            accept="image/png,image/jpeg,image/bmp,image/tiff"
            className="hidden"
            onChange={handleImportGtMask}
          />
          <button
            onClick={() => gtMaskInputRef.current?.click()}
            disabled={!currentProject?.id || !currentFrame?.id || isImportingGt || isSaving || isExporting || isPropagating}
            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"
          >
            {isImportingGt ? '导入中...' : '导入 GT Mask'}
          </button>
          <div className="flex items-center gap-1 rounded-md border border-white/10 bg-white/[0.03] px-2 py-1">
            <span className="text-[10px] text-gray-500 whitespace-nowrap">参考帧 {currentFrameNumber || 0}</span>
            <span className="text-[10px] text-gray-600">帧</span>
            <input
              aria-label="传播起始帧"
              type="number"
              min={1}
              max={Math.max(totalFrames, 1)}
              value={propagationStartFrame}
              onChange={(event) => setPropagationStartFrame(clampFrameNumber(Number(event.target.value) || 1))}
              disabled={isPropagating || isSaving || isExporting || isImportingGt || totalFrames === 0}
              className="h-6 w-14 rounded bg-black/20 border border-white/10 px-1 text-[10px] text-gray-300 outline-none focus:border-cyan-500/50 disabled:opacity-40"
            />
            <span className="text-[10px] text-gray-600">-</span>
            <input
              aria-label="传播结束帧"
              type="number"
              min={1}
              max={Math.max(totalFrames, 1)}
              value={propagationEndFrame}
              onChange={(event) => setPropagationEndFrame(clampFrameNumber(Number(event.target.value) || 1))}
              disabled={isPropagating || isSaving || isExporting || isImportingGt || totalFrames === 0}
              className="h-6 w-14 rounded bg-black/20 border border-white/10 px-1 text-[10px] text-gray-300 outline-none focus:border-cyan-500/50 disabled:opacity-40"
            />
          </div>
          <button
            onClick={handleAutoPropagate}
            disabled={!currentProject?.id || !currentFrame?.id || isSaving || isExporting || isImportingGt || isPropagating}
            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"
          >
            {isPropagating ? '传播中...' : '自动传播'}
          </button>
          <button
            onClick={handleExportMasks}
            disabled={!currentProject?.id || isExporting || isSaving || isPropagating}
            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 ? '导出中...' : '导出 PNG Mask ZIP'}
          </button>
          <button
            onClick={handleExport}
            disabled={!currentProject?.id || isExporting || isSaving || isPropagating}
            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 || isPropagating}
            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}
          onUndo={undoMasks}
          onRedo={redoMasks}
          canUndo={maskHistory.length > 0}
          canRedo={maskFuture.length > 0}
        />
        
        <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}
              onDeleteMaskAnnotations={handleDeleteMaskAnnotations}
            />
          </div>
        </div>
        
        <OntologyInspector />
      </div>

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