import React, { useState } from 'react';
import { invoke } from '@tauri-apps/api/core';
import { open } from '@tauri-apps/plugin-dialog';
import { listen } from '@tauri-apps/api/event';
import {
  ImagePlus,
  Upload,
  Download,
  Settings,
  Play,
  CheckCircle,
  XCircle,
  AlertCircle,
  Folder,
  FolderOpen,
  Image as ImageIcon,
  Loader2
} from 'lucide-react';

interface ImageEnhancementProgress {
  current: number;
  total: number;
  percentage: number;
  status: string;
  current_file?: string;
}

interface ImageEnhancementResult {
  success: boolean;
  message: string;
  input_path: string;
  output_path?: string;
  execution_time: number;
  error?: string;
}

interface BatchImageEnhancementResult {
  success: boolean;
  message: string;
  total_processed: number;
  successful_count: number;
  failed_count: number;
  results: ImageEnhancementResult[];
  total_execution_time: number;
}

const AiModelFaceHairFixTool: React.FC = () => {
  // 单张图片处理状态
  const [singleImagePath, setSingleImagePath] = useState<string>('');
  const [singleOutputPath, setSingleOutputPath] = useState<string>('');
  
  // 批量处理状态
  const [batchInputDir, setBatchInputDir] = useState<string>('');
  const [batchOutputDir, setBatchOutputDir] = useState<string>('');
  
  // 通用设置
  const [serverUrl, setServerUrl] = useState<string>('http://192.168.0.193:8188');
  const [facePrompt, setFacePrompt] = useState<string>('beautiful woman, perfect skin, detailed facial features');
  const [faceDenoiseStr, setFaceDenoiseStr] = useState<string>('0.35');
  
  // 处理状态
  const [isProcessing, setIsProcessing] = useState<boolean>(false);
  const [progress, setProgress] = useState<ImageEnhancementProgress | null>(null);
  const [results, setResults] = useState<ImageEnhancementResult[] | null>(null);
  const [batchResults, setBatchResults] = useState<BatchImageEnhancementResult | null>(null);
  
  // 当前模式：single 或 batch
  const [mode, setMode] = useState<'single' | 'batch'>('single');



  // 选择单张图片
  const selectSingleImage = async () => {
    try {
      const selected = await open({
        multiple: false,
        filters: [
          {
            name: 'Image Files',
            extensions: ['png', 'jpg', 'jpeg', 'bmp', 'tiff', 'webp']
          }
        ]
      });
      
      if (selected && typeof selected === 'string') {
        setSingleImagePath(selected);
        // 自动生成输出路径
        const pathParts = selected.split('.');
        const extension = pathParts.pop();
        const basePath = pathParts.join('.');
        setSingleOutputPath(`${basePath}_enhanced.${extension}`);
      }
    } catch (error) {
      console.error('选择图片失败:', error);
    }
  };

  // 选择单张图片输出路径
  const selectSingleOutputPath = async () => {
    try {
      const selected = await open({
        multiple: false,
        filters: [
          {
            name: 'Image Files',
            extensions: ['png', 'jpg', 'jpeg', 'bmp', 'tiff', 'webp']
          }
        ]
      });
      
      if (selected && typeof selected === 'string') {
        setSingleOutputPath(selected);
      }
    } catch (error) {
      console.error('选择输出路径失败:', error);
    }
  };

  // 选择批量输入目录
  const selectBatchInputDir = async () => {
    try {
      const selected = await open({
        directory: true,
        multiple: false
      });
      
      if (selected && typeof selected === 'string') {
        setBatchInputDir(selected);
      }
    } catch (error) {
      console.error('选择输入目录失败:', error);
    }
  };

  // 选择批量输出目录
  const selectBatchOutputDir = async () => {
    try {
      const selected = await open({
        directory: true,
        multiple: false
      });
      
      if (selected && typeof selected === 'string') {
        setBatchOutputDir(selected);
      }
    } catch (error) {
      console.error('选择输出目录失败:', error);
    }
  };

  // 处理单张图片
  const processSingleImage = async () => {
    if (!singleImagePath || !singleOutputPath) {
      alert('请选择输入图片和输出路径');
      return;
    }

    setIsProcessing(true);
    setResults(null);
    setProgress(null);

    try {
      // 监听进度事件
      const unlisten = await listen<ImageEnhancementProgress>('image-enhancement-progress', (event) => {
        setProgress(event.payload);
      });

      const result = await invoke<ImageEnhancementResult>('ai_model_face_hair_fix_single_image', {
        inputImage: singleImagePath,
        serverUrl,
        facePrompt,
        faceDenoise: faceDenoiseStr,
        outputImage: singleOutputPath
      });

      setResults([result]);

      // 清理事件监听器
      unlisten();
    } catch (error) {
      console.error('处理失败:', error);
      alert(`处理失败: ${error}`);
    } finally {
      setIsProcessing(false);
    }
  };

  // 批量处理图片
  const processBatchImages = async () => {
    if (!batchInputDir || !batchOutputDir) {
      alert('请选择输入目录和输出目录');
      return;
    }

    setIsProcessing(true);
    setBatchResults(null);
    setProgress(null);

    try {
      // 监听进度事件
      const unlisten = await listen<ImageEnhancementProgress>('image-enhancement-progress', (event) => {
        setProgress(event.payload);
      });

      const result = await invoke<BatchImageEnhancementResult>('ai_model_face_hair_fix_batch_images', {
        inputDir: batchInputDir,
        serverUrl,
        facePrompt,
        faceDenoise: faceDenoiseStr,
        outputDir: batchOutputDir
      });

      setBatchResults(result);

      // 清理事件监听器
      unlisten();
    } catch (error) {
      console.error('批量处理失败:', error);
      alert(`批量处理失败: ${error}`);
    } finally {
      setIsProcessing(false);
    }
  };

  // 重置状态
  const resetState = () => {
    setProgress(null);
    setResults(null);
    setBatchResults(null);
    setIsProcessing(false);
  };

  return (
    <div className="p-6 max-w-6xl mx-auto">
      <div className="mb-6">
        <h1 className="text-2xl font-bold text-gray-900 mb-2 flex items-center">
          <ImagePlus className="mr-3 text-blue-600" size={28} />
          AI模型面部头发修复工具
        </h1>
        <p className="text-gray-600">
          基于ComfyUI的专业AI模型面部头发修复工具，支持单张图片和批量处理
        </p>
      </div>

      {/* 模式选择 */}
      <div className="mb-6">
        <div className="flex space-x-4">
          <button
            onClick={() => { setMode('single'); resetState(); }}
            className={`px-4 py-2 rounded-lg font-medium transition-colors ${
              mode === 'single'
                ? 'bg-blue-600 text-white'
                : 'bg-gray-200 text-gray-700 hover:bg-gray-300'
            }`}
          >
            <ImageIcon className="inline mr-2" size={16} />
            单张图片处理
          </button>
          <button
            onClick={() => { setMode('batch'); resetState(); }}
            className={`px-4 py-2 rounded-lg font-medium transition-colors ${
              mode === 'batch'
                ? 'bg-blue-600 text-white'
                : 'bg-gray-200 text-gray-700 hover:bg-gray-300'
            }`}
          >
            <Folder className="inline mr-2" size={16} />
            批量处理
          </button>
        </div>
      </div>

      {/* 服务器设置 */}
      <div className="bg-white rounded-lg shadow-sm border p-6 mb-6">
        <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
          <Settings className="mr-2 text-gray-600" size={20} />
          服务器设置
        </h2>
        
        <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
          <div>
            <label className="block text-sm font-medium text-gray-700 mb-2">
              ComfyUI服务器地址
            </label>
            <input
              type="text"
              value={serverUrl}
              onChange={(e) => setServerUrl(e.target.value)}
              className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
              placeholder="http://192.168.0.193:8188"
            />
          </div>
          
          <div>
            <label className="block text-sm font-medium text-gray-700 mb-2">
              去噪强度 (0.0-1.0)
            </label>
            <input
              type="text"
              value={faceDenoiseStr}
              onChange={(e) => setFaceDenoiseStr(e.target.value)}
              className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
              placeholder="0.25"
            />
          </div>
        </div>
        
        <div className="mt-4">
          <label className="block text-sm font-medium text-gray-700 mb-2">
            面部提示词
          </label>
          <textarea
            value={facePrompt}
            onChange={(e) => setFacePrompt(e.target.value)}
            rows={3}
            className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
            placeholder="beautiful woman, perfect skin, detailed facial features"
          />
        </div>
      </div>

      {/* 单张图片处理 */}
      {mode === 'single' && (
        <div className="bg-white rounded-lg shadow-sm border p-6 mb-6">
          <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
            <ImageIcon className="mr-2 text-gray-600" size={20} />
            单张图片处理
          </h2>

          <div className="space-y-4">
            <div>
              <label className="block text-sm font-medium text-gray-700 mb-2">
                输入图片
              </label>
              <div className="flex space-x-2">
                <input
                  type="text"
                  value={singleImagePath}
                  onChange={(e) => setSingleImagePath(e.target.value)}
                  className="flex-1 px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
                  placeholder="选择要处理的图片文件"
                />
                <button
                  onClick={selectSingleImage}
                  className="px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 transition-colors flex items-center"
                >
                  <Upload className="mr-2" size={16} />
                  选择
                </button>
              </div>
            </div>

            <div>
              <label className="block text-sm font-medium text-gray-700 mb-2">
                输出路径
              </label>
              <div className="flex space-x-2">
                <input
                  type="text"
                  value={singleOutputPath}
                  onChange={(e) => setSingleOutputPath(e.target.value)}
                  className="flex-1 px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
                  placeholder="增强后图片的保存路径"
                />
                <button
                  onClick={selectSingleOutputPath}
                  className="px-4 py-2 bg-gray-600 text-white rounded-lg hover:bg-gray-700 transition-colors flex items-center"
                >
                  <Download className="mr-2" size={16} />
                  选择
                </button>
              </div>
            </div>

            <div className="flex justify-end">
              <button
                onClick={processSingleImage}
                disabled={isProcessing || !singleImagePath || !singleOutputPath}
                className="px-6 py-2 bg-green-600 text-white rounded-lg hover:bg-green-700 disabled:bg-gray-400 disabled:cursor-not-allowed transition-colors flex items-center"
              >
                {isProcessing ? (
                  <Loader2 className="mr-2 animate-spin" size={16} />
                ) : (
                  <Play className="mr-2" size={16} />
                )}
                {isProcessing ? '处理中...' : '开始处理'}
              </button>
            </div>
          </div>
        </div>
      )}

      {/* 批量处理 */}
      {mode === 'batch' && (
        <div className="bg-white rounded-lg shadow-sm border p-6 mb-6">
          <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
            <Folder className="mr-2 text-gray-600" size={20} />
            批量图片处理
          </h2>

          <div className="space-y-4">
            <div>
              <label className="block text-sm font-medium text-gray-700 mb-2">
                输入目录
              </label>
              <div className="flex space-x-2">
                <input
                  type="text"
                  value={batchInputDir}
                  onChange={(e) => setBatchInputDir(e.target.value)}
                  className="flex-1 px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
                  placeholder="包含图片文件的目录"
                />
                <button
                  onClick={selectBatchInputDir}
                  className="px-4 py-2 bg-blue-600 text-white rounded-lg hover:bg-blue-700 transition-colors flex items-center"
                >
                  <FolderOpen className="mr-2" size={16} />
                  选择
                </button>
              </div>
            </div>

            <div>
              <label className="block text-sm font-medium text-gray-700 mb-2">
                输出目录
              </label>
              <div className="flex space-x-2">
                <input
                  type="text"
                  value={batchOutputDir}
                  onChange={(e) => setBatchOutputDir(e.target.value)}
                  className="flex-1 px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
                  placeholder="增强后图片的保存目录"
                />
                <button
                  onClick={selectBatchOutputDir}
                  className="px-4 py-2 bg-gray-600 text-white rounded-lg hover:bg-gray-700 transition-colors flex items-center"
                >
                  <FolderOpen className="mr-2" size={16} />
                  选择
                </button>
              </div>
            </div>

            <div className="flex justify-end">
              <button
                onClick={processBatchImages}
                disabled={isProcessing || !batchInputDir || !batchOutputDir}
                className="px-6 py-2 bg-green-600 text-white rounded-lg hover:bg-green-700 disabled:bg-gray-400 disabled:cursor-not-allowed transition-colors flex items-center"
              >
                {isProcessing ? (
                  <Loader2 className="mr-2 animate-spin" size={16} />
                ) : (
                  <Play className="mr-2" size={16} />
                )}
                {isProcessing ? '批量处理中...' : '开始批量处理'}
              </button>
            </div>
          </div>
        </div>
      )}

      {/* 进度显示 */}
      {progress && (
        <div className="bg-white rounded-lg shadow-sm border p-6 mb-6">
          <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
            <Loader2 className="mr-2 text-blue-600 animate-spin" size={20} />
            处理进度
          </h2>

          <div className="space-y-4">
            <div>
              <div className="flex justify-between text-sm text-gray-600 mb-1">
                <span>{progress.status}</span>
                <span>{progress.current}/{progress.total} ({progress.percentage.toFixed(1)}%)</span>
              </div>
              <div className="w-full bg-gray-200 rounded-full h-2">
                <div
                  className="bg-blue-600 h-2 rounded-full transition-all duration-300"
                  style={{ width: `${progress.percentage}%` }}
                />
              </div>
            </div>

            {progress.current_file && (
              <div className="text-sm text-gray-600">
                <span className="font-medium">当前文件:</span> {progress.current_file}
              </div>
            )}
          </div>
        </div>
      )}

      {/* 单张图片结果 */}
      {results && mode === 'single' && (
        <div className="bg-white rounded-lg shadow-sm border p-6 mb-6">
          <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
            <CheckCircle className="mr-2 text-green-600" size={20} />
            处理结果
          </h2>

          {results.map((result, index) => (
            <div key={index} className="border rounded-lg p-4">
              <div className="flex items-center justify-between mb-2">
                <div className="flex items-center">
                  {result.success ? (
                    <CheckCircle className="mr-2 text-green-600" size={16} />
                  ) : (
                    <XCircle className="mr-2 text-red-600" size={16} />
                  )}
                  <span className={`font-medium ${result.success ? 'text-green-700' : 'text-red-700'}`}>
                    {result.success ? '处理成功' : '处理失败'}
                  </span>
                </div>
                <span className="text-sm text-gray-500">
                  耗时: {result.execution_time.toFixed(2)}秒
                </span>
              </div>

              <div className="text-sm text-gray-600 space-y-1">
                <div><span className="font-medium">输入:</span> {result.input_path}</div>
                {result.output_path && (
                  <div><span className="font-medium">输出:</span> {result.output_path}</div>
                )}
                <div><span className="font-medium">消息:</span> {result.message}</div>
                {result.error && (
                  <div className="text-red-600"><span className="font-medium">错误:</span> {result.error}</div>
                )}
              </div>
            </div>
          ))}
        </div>
      )}

      {/* 批量处理结果 */}
      {batchResults && mode === 'batch' && (
        <div className="bg-white rounded-lg shadow-sm border p-6">
          <h2 className="text-lg font-semibold text-gray-900 mb-4 flex items-center">
            {batchResults.success ? (
              <CheckCircle className="mr-2 text-green-600" size={20} />
            ) : (
              <AlertCircle className="mr-2 text-yellow-600" size={20} />
            )}
            批量处理结果
          </h2>

          <div className="mb-4 p-4 bg-gray-50 rounded-lg">
            <div className="grid grid-cols-2 md:grid-cols-4 gap-4 text-sm">
              <div>
                <span className="font-medium text-gray-700">总计:</span>
                <span className="ml-2 text-gray-900">{batchResults.total_processed}</span>
              </div>
              <div>
                <span className="font-medium text-green-700">成功:</span>
                <span className="ml-2 text-green-900">{batchResults.successful_count}</span>
              </div>
              <div>
                <span className="font-medium text-red-700">失败:</span>
                <span className="ml-2 text-red-900">{batchResults.failed_count}</span>
              </div>
              <div>
                <span className="font-medium text-gray-700">总耗时:</span>
                <span className="ml-2 text-gray-900">{batchResults.total_execution_time.toFixed(2)}秒</span>
              </div>
            </div>
            <div className="mt-2 text-sm text-gray-600">
              {batchResults.message}
            </div>
          </div>

          <div className="space-y-2 max-h-96 overflow-y-auto">
            {batchResults.results.map((result, index) => (
              <div key={index} className="border rounded-lg p-3">
                <div className="flex items-center justify-between mb-1">
                  <div className="flex items-center">
                    {result.success ? (
                      <CheckCircle className="mr-2 text-green-600" size={14} />
                    ) : (
                      <XCircle className="mr-2 text-red-600" size={14} />
                    )}
                    <span className={`text-sm font-medium ${result.success ? 'text-green-700' : 'text-red-700'}`}>
                      {result.success ? '成功' : '失败'}
                    </span>
                  </div>
                  <span className="text-xs text-gray-500">
                    {result.execution_time.toFixed(2)}秒
                  </span>
                </div>

                <div className="text-xs text-gray-600 space-y-1">
                  <div className="truncate"><span className="font-medium">输入:</span> {result.input_path}</div>
                  {result.output_path && (
                    <div className="truncate"><span className="font-medium">输出:</span> {result.output_path}</div>
                  )}
                  {result.error && (
                    <div className="text-red-600 truncate"><span className="font-medium">错误:</span> {result.error}</div>
                  )}
                </div>
              </div>
            ))}
          </div>
        </div>
      )}
    </div>
  );
};

export default AiModelFaceHairFixTool;
