import React, { useState, useCallback } from 'react';
import { Target, PlusCircle, MinusCircle, SquareDashed, Sparkles, SendToBack, Image as ImageIcon, Undo, Redo, Loader2 } from 'lucide-react';
import { cn } from '../lib/utils';
import { Stage, Layer, Image as KonvaImage, Circle, Path, Group } from 'react-konva';
import useImage from 'use-image';
import { OntologyInspector } from './OntologyInspector';
import { useStore } from '../store/useStore';
import { predictMask } from '../lib/api';

interface AISegmentationProps {
    onSendToWorkspace: () => void;
}

export function AISegmentation({ onSendToWorkspace }: AISegmentationProps) {
    const storeActiveTool = useStore((state) => state.activeTool);
    const setActiveTool = useStore((state) => state.setActiveTool);
    const masks = useStore((state) => state.masks);
    const addMask = useStore((state) => state.addMask);
    const clearMasks = useStore((state) => state.clearMasks);

    const [modelSize, setModelSize] = useState('vit_l');
    const [semanticText, setSemanticText] = useState('');
    const [autoDeleteBg, setAutoDeleteBg] = useState(true);
    const [cropMode, setCropMode] = useState(false);
    const [isInferencing, setIsInferencing] = useState(false);
    
    // Canvas state
    const [scale, setScale] = useState(1);
    const [position, setPosition] = useState({ x: 0, y: 0 });
    const [points, setPoints] = useState<{ x: number, y: number, type: 'pos'|'neg' }[]>([]);
    const [cursorPos, setCursorPos] = useState({ x: 0, y: 0 });
    const [image] = useImage('https://images.unsplash.com/photo-1549317661-bd32c8ce0be2?q=80&w=2070&auto=format&fit=crop');

    const effectiveTool = storeActiveTool;

    const handleWheel = (e: any) => {
      e.evt.preventDefault();
      const scaleBy = 1.1;
      const stage = e.target.getStage();
      const oldScale = stage.scaleX();
      const mousePointTo = {
        x: stage.getPointerPosition().x / oldScale - stage.x() / oldScale,
        y: stage.getPointerPosition().y / oldScale - stage.y() / oldScale,
      };
      const newScale = e.evt.deltaY < 0 ? oldScale * scaleBy : oldScale / scaleBy;
      setScale(newScale);
      setPosition({
        x: -(mousePointTo.x - stage.getPointerPosition().x / newScale) * newScale,
        y: -(mousePointTo.y - stage.getPointerPosition().y / newScale) * newScale,
      });
    };

    const handleMouseMove = (e: any) => {
      const stage = e.target.getStage();
      if (!stage) return;
      const pos = stage.getPointerPosition();
      if (pos) {
         const imageX = (pos.x - position.x) / scale;
         const imageY = (pos.y - position.y) / scale;
         setCursorPos({ x: imageX, y: imageY });
      }
    };

    const runInference = useCallback(async () => {
      if (points.length === 0 && !semanticText.trim()) return;
      setIsInferencing(true);
      try {
        const result = await predictMask({
          imageUrl: 'https://images.unsplash.com/photo-1549317661-bd32c8ce0be2?q=80&w=2070&auto=format&fit=crop',
          points: points.map((p) => ({ x: p.x, y: p.y, type: p.type })),
          text: semanticText.trim() || undefined,
          modelSize,
        });

        result.masks.forEach((m) => {
          addMask({
            id: m.id,
            frameId: 'frame-ai-1',
            pathData: m.pathData,
            label: m.label,
            color: m.color,
            segmentation: m.segmentation,
            bbox: m.bbox,
            area: m.area,
          });
        });
      } catch (err) {
        console.error('AI inference failed:', err);
      } finally {
        setIsInferencing(false);
      }
    }, [points, semanticText, modelSize, addMask]);

    const handleStageClick = (e: any) => {
      if (effectiveTool === 'move') return;
      if (effectiveTool === 'point_pos' || effectiveTool === 'point_neg') {
        const stage = e.target.getStage();
        const pos = stage.getRelativePointerPosition();
        if (pos) {
          setPoints([...points, { x: pos.x, y: pos.y, type: effectiveTool === 'point_pos' ? 'pos' : 'neg' }]);
        }
      }
    };

    return (
        <div className="w-full h-full flex bg-[#0a0a0a]">
            {/* Left AI Controller Panel */}
            <aside className="w-80 bg-[#0d0d0d] flex flex-col border-r border-white/5 shrink-0 z-10 overflow-hidden">
                <div className="h-16 border-b border-white/5 flex items-center px-6 shrink-0 justify-between">
                    <div className="flex items-center font-medium text-[11px] uppercase tracking-widest text-cyan-400">
                        <Sparkles size={16} className="mr-3 text-cyan-400" />
                        AI智能分割引擎
                    </div>
                </div>

                <div className="flex-1 overflow-y-auto p-6 flex flex-col gap-8">
                    {/* Model Select */}
                    <div>
                        <h3 className="text-[10px] font-bold text-gray-500 uppercase tracking-widest mb-3">视觉基础模型选型</h3>
                        <div className="bg-[#111] border border-white/5 flex p-1 rounded-lg">
                        {['vit_b', 'vit_l', 'vit_h'].map(m => (
                            <button 
                                key={m}
                                className={cn("flex-1 text-xs py-2 rounded-md transition-colors text-center uppercase tracking-wider font-mono", modelSize === m ? "bg-white/10 text-white font-medium shadow-sm" : "text-gray-500 hover:text-gray-300 hover:bg-white/5")}
                                onClick={() => setModelSize(m)}
                            >
                                {m.split('_')[1]}
                            </button>
                        ))}
                        </div>
                    </div>

                    {/* Prompt Tools */}
                    <div>
                        <h3 className="text-[10px] font-bold text-gray-500 uppercase tracking-widest mb-3">交互式提示工具</h3>
                        <div className="grid grid-cols-2 gap-3">
                        <button 
                            onClick={() => setActiveTool('point_pos')}
                            className={cn("flex flex-col items-center justify-center p-4 rounded-lg border transition-all", effectiveTool === 'point_pos' ? "bg-green-500/10 border-green-500/30 text-green-400 shadow-[0_0_15px_rgba(34,197,94,0.1)]" : "bg-[#111] border-white/5 text-gray-400 hover:bg-white/5 hover:text-gray-200")}
                        >
                            <PlusCircle size={22} className="mb-3" />
                            <span className="text-[10px] uppercase tracking-wider font-semibold">正向选点</span>
                        </button>
                        
                        <button 
                            onClick={() => setActiveTool('point_neg')}
                            className={cn("flex flex-col items-center justify-center p-4 rounded-lg border transition-all", effectiveTool === 'point_neg' ? "bg-red-500/10 border-red-500/30 text-red-500 shadow-[0_0_15px_rgba(239,68,68,0.1)]" : "bg-[#111] border-white/5 text-gray-400 hover:bg-white/5 hover:text-gray-200")}
                        >
                            <MinusCircle size={22} className="mb-3" />
                            <span className="text-[10px] uppercase tracking-wider font-semibold">反向选点</span>
                        </button>
                        
                        <button 
                            onClick={() => setActiveTool('box_select')}
                            className={cn("flex flex-col items-center justify-center p-4 rounded-lg border transition-all", effectiveTool === 'box_select' ? "bg-blue-500/10 border-blue-500/30 text-blue-400 shadow-[0_0_15px_rgba(59,130,246,0.1)]" : "bg-[#111] border-white/5 text-gray-400 hover:bg-white/5 hover:text-gray-200")}
                        >
                            <SquareDashed size={22} className="mb-3" />
                            <span className="text-[10px] uppercase tracking-wider font-semibold">边界框选</span>
                        </button>

                        <button 
                            onClick={() => setActiveTool('move')}
                            className={cn("flex flex-col items-center justify-center p-4 rounded-lg border transition-all", effectiveTool === 'move' ? "bg-white/10 border-white/20 text-white shadow-[0_0_15px_rgba(255,255,255,0.05)]" : "bg-[#111] border-white/5 text-gray-400 hover:bg-white/5 hover:text-gray-200")}
                        >
                            <Target size={22} className="mb-3" />
                            <span className="text-[10px] uppercase tracking-wider font-semibold">视口控制</span>
                        </button>
                        </div>
                    </div>

                    {/* Semantic Description */}
                    <div>
                        <div className="flex justify-between items-center mb-3">
                        <h3 className="text-[10px] font-bold text-gray-500 uppercase tracking-widest">语义引导描述</h3>
                        <span className="text-[9px] bg-cyan-500/10 text-cyan-400 px-1.5 py-0.5 rounded border border-cyan-500/20 font-mono">零样本推理</span>
                        </div>
                        <textarea 
                        value={semanticText}
                        onChange={e => setSemanticText(e.target.value)}
                        placeholder="例如：'分割出左侧车道上行驶的所有红色汽车'..."
                        className="w-full bg-[#111] border border-white/5 rounded-lg p-3 text-sm text-white placeholder-gray-600 focus:outline-none focus:border-cyan-500/50 focus:ring-1 focus:ring-cyan-500/50 transition-all font-sans min-h-[100px] resize-none hover:border-white/10"
                        />
                    </div>

                    {/* Parameters */}
                    <div>
                        <h3 className="text-[10px] font-bold text-gray-500 uppercase tracking-widest mb-3 flex items-center gap-2">核心参数设定</h3>
                        <div className="space-y-4 bg-[#111] rounded-lg p-5 border border-white/5">
                        <div className="flex items-center justify-between cursor-pointer group" onClick={() => setCropMode(!cropMode)}>
                            <span className="text-[11px] text-gray-400 uppercase tracking-wider font-medium group-hover:text-gray-200 transition-colors">自动裁剪无锚区域</span>
                            <button className={cn("w-8 h-4 rounded-full transition-colors relative", cropMode ? "bg-cyan-500" : "bg-white/20")}>
                                <div className={cn("absolute top-0.5 left-0.5 w-3 h-3 bg-white rounded-full transition-transform shadow-sm", cropMode ? "translate-x-4" : "")} />
                            </button>
                        </div>
                        
                        <div className="flex items-center justify-between cursor-pointer group" onClick={() => setAutoDeleteBg(!autoDeleteBg)}>
                            <span className="text-[11px] text-gray-400 uppercase tracking-wider font-medium group-hover:text-gray-200 transition-colors">自动清理干涉点</span>
                            <button className={cn("w-8 h-4 rounded-full transition-colors relative", autoDeleteBg ? "bg-cyan-500" : "bg-white/20")}>
                                <div className={cn("absolute top-0.5 left-0.5 w-3 h-3 bg-white rounded-full transition-transform shadow-sm", autoDeleteBg ? "translate-x-4" : "")} />
                            </button>
                        </div>
                        </div>
                    </div>
                </div>

                <div className="p-6 bg-[#0a0a0a] border-t border-white/5 shrink-0 flex flex-col gap-3">
                    <button 
                        onClick={runInference}
                        disabled={isInferencing}
                        className={cn(
                            "w-full py-3.5 rounded-lg flex items-center justify-center gap-2 transition-all shadow-lg font-medium tracking-wide text-xs uppercase",
                            isInferencing 
                                ? "bg-cyan-500/50 text-black/70 cursor-not-allowed"
                                : "bg-cyan-500 hover:bg-cyan-400 text-black shadow-cyan-500/20 hover:shadow-cyan-500/40"
                        )}
                    >
                        {isInferencing ? <Loader2 size={16} className="animate-spin" /> : <Sparkles size={16} />}
                        {isInferencing ? '推理中...' : '执行高精度语义分割'}
                    </button>
                    <button 
                        onClick={onSendToWorkspace}
                        className="w-full py-3.5 rounded-lg flex items-center justify-center gap-2 transition-all font-medium tracking-wide text-xs uppercase bg-white/5 hover:bg-white/10 text-gray-300 border border-white/5 hover:border-white/10"
                    >
                        <SendToBack size={16} /> 退档推送至工作区重组
                    </button>
                </div>
            </aside>

            {/* Right Canvas Area */}
            <main className="flex-1 bg-[#151515] relative overflow-hidden flex flex-col">
                <header className="h-16 border-b border-white/5 bg-[#111] flex items-center justify-between px-6 shrink-0">
                    <div className="flex flex-col">
                       <h2 className="text-sm font-semibold tracking-wide text-white">模型端推理侧可视化 (Visualizer)</h2>
                       <span className="text-[10px] text-gray-500 uppercase tracking-widest font-mono">SAM 3 内核级动态即时渲染</span>
                    </div>
                    <div className="flex items-center gap-4">
                       <button className="w-8 h-8 rounded text-gray-400 hover:bg-white/5 hover:text-white flex items-center justify-center transition-colors" title="撤销操作 (Ctrl+Z)">
                          <Undo size={14} />
                       </button>
                       <button className="w-8 h-8 rounded text-gray-400 hover:bg-white/5 hover:text-white flex items-center justify-center transition-colors" title="重做操作 (Ctrl+Shift+Z)">
                          <Redo size={14} />
                       </button>
                       <div className="w-px h-4 bg-white/10 mx-1"></div>
                       <button className="flex items-center gap-2 text-xs text-gray-400 hover:text-white transition-colors bg-white/5 hover:bg-white/10 px-3 py-1.5 rounded-md border border-white/5">
                          <ImageIcon size={14} /> 上传替换底图
                       </button>
                       <button className="text-xs text-gray-400 hover:text-white transition-colors px-3 py-1.5" onClick={() => { setPoints([]); clearMasks(); }}>
                          清空全体锚点
                       </button>
                    </div>
                </header>
                
                <div className="flex-1 relative p-8">
                    <div className="w-full h-full relative border border-white/5 rounded shadow-2xl bg-[#1e1e1e] overflow-hidden cursor-crosshair">
                        <Stage
                            width={window.innerWidth - 320 - 64}
                            height={window.innerHeight - 64 - 64}
                            onWheel={handleWheel}
                            onMouseMove={handleMouseMove}
                            onClick={handleStageClick}
                            scaleX={scale}
                            scaleY={scale}
                            x={position.x}
                            y={position.y}
                            draggable={effectiveTool === 'move'}
                        >
                            <Layer>
                                {/* Background Image */}
                                {image && (
                                    <KonvaImage 
                                        image={image} 
                                        x={0} 
                                        y={0}
                                        opacity={0.8}
                                    />
                                )}
                                
                                {/* AI Returned Masks */}
                                {masks.map((mask) => (
                                    <Group key={mask.id} opacity={0.45}>
                                        <Path
                                            data={mask.pathData}
                                            fill={mask.color}
                                            stroke={mask.color}
                                            strokeWidth={1 / scale}
                                        />
                                    </Group>
                                ))}

                                {/* Points */}
                                {points.map((p, i) => (
                                <Group key={i} x={p.x} y={p.y}>
                                    <Circle 
                                        radius={6 / scale} 
                                        fill={p.type === 'pos' ? '#22c55e' : '#ef4444'}
                                        stroke="#ffffff"
                                        strokeWidth={2 / scale}
                                        shadowColor="black"
                                        shadowBlur={4}
                                    />
                                    <Circle 
                                        radius={1.5 / scale} 
                                        fill="#ffffff"
                                    />
                                </Group>
                                ))}
                            </Layer>
                        </Stage>
                        <div className="absolute bottom-4 left-4 flex gap-4 text-[10px] font-mono text-gray-500 pointer-events-none">
                            <span>光标坐标: {cursorPos.x.toFixed(2)}, {cursorPos.y.toFixed(2)}</span>
                            <span>缩放比率: {(scale * 100).toFixed(0)}%</span>
                            <span>遮罩数: {masks.length}</span>
                        </div>
                    </div>
                </div>
            </main>

            {/* Right Ontology / Label Assignment Panel */}
            <OntologyInspector />
        </div>
    );
}
