next1/components/data-collection-step.tsx

"use client"

import type React from "react"

import { useState, useRef, useEffect, useCallback } from "react"
import { Button } from "@/components/ui/button"
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"
import { Badge } from "@/components/ui/badge"
import { Progress } from "@/components/ui/progress"
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"
import { Upload, ImageIcon, Trash2, Eye, Info, CheckCircle } from "lucide-react"

interface ImageData {
  id: string
  name: string
  url: string
  size: number
  type: string
}

interface DataCollectionStepProps {
  onNext?: () => void
  selectedImages: string[]
  onImagesChange: (images: string[]) => void
}

const sampleImageCategories = {
  landscape: {
    name: "风景画",
    images: [
      { id: "landscape-1", name: "山水湖泊.jpg", url: "/mountain-lake-painting.png", size: 245760, type: "image/jpeg" },
      { id: "landscape-2", name: "日出山峰.jpg", url: "/sunrise-mountain-peaks.png", size: 234560, type: "image/jpeg" },
      { id: "landscape-3", name: "森林小径.jpg", url: "/forest-path-painting.png", size: 267890, type: "image/jpeg" },
      { id: "landscape-4", name: "海边日落.jpg", url: "/sunset-beach-painting.png", size: 298765, type: "image/jpeg" },
      {
        id: "landscape-5",
        name: "雪山风光.jpg",
        url: "/snowy-mountain-landscape.png",
        size: 312450,
        type: "image/jpeg",
      },
      { id: "landscape-6", name: "田园风光.jpg", url: "/pastoral-landscape.png", size: 287650, type: "image/jpeg" },
      { id: "landscape-7", name: "瀑布奇观.jpg", url: "/waterfall-landscape.png", size: 345670, type: "image/jpeg" },
      { id: "landscape-8", name: "沙漠绿洲.jpg", url: "/desert-oasis-painting.png", size: 276540, type: "image/jpeg" },
      {
        id: "landscape-9",
        name: "湖光山色.jpg",
        url: "/lake-mountain-reflection-painting.png",
        size: 298760,
        type: "image/jpeg",
      },
      {
        id: "landscape-10",
        name: "秋叶满山.jpg",
        url: "/autumn-mountain-foliage-painting.png",
        size: 321450,
        type: "image/jpeg",
      },
    ],
  },
  portrait: {
    name: "人物肖像",
    images: [
      { id: "portrait-1", name: "经典肖像.jpg", url: "/artistic-portrait.png", size: 312320, type: "image/jpeg" },
      { id: "portrait-2", name: "女性肖像.jpg", url: "/elegant-woman-portrait.png", size: 298450, type: "image/jpeg" },
      { id: "portrait-3", name: "老人肖像.jpg", url: "/elderly-man-portrait.png", size: 287650, type: "image/jpeg" },
      { id: "portrait-4", name: "儿童肖像.jpg", url: "/child-portrait-painting.png", size: 276540, type: "image/jpeg" },
      {
        id: "portrait-5",
        name: "艺术家自画像.jpg",
        url: "/artist-self-portrait.png",
        size: 345670,
        type: "image/jpeg",
      },
      {
        id: "portrait-6",
        name: "古典美人.jpg",
        url: "/classical-beauty-portrait.png",
        size: 321450,
        type: "image/jpeg",
      },
      {
        id: "portrait-7",
        name: "现代肖像.jpg",
        url: "/modern-portrait-painting.png",
        size: 298760,
        type: "image/jpeg",
      },
      {
        id: "portrait-8",
        name: "侧面肖像.jpg",
        url: "/profile-portrait-painting.png",
        size: 287650,
        type: "image/jpeg",
      },
      {
        id: "portrait-9",
        name: "双人肖像.jpg",
        url: "/couple-portrait-painting.png",
        size: 356780,
        type: "image/jpeg",
      },
      {
        id: "portrait-10",
        name: "表情肖像.jpg",
        url: "/expressive-portrait-painting.png",
        size: 312450,
        type: "image/jpeg",
      },
    ],
  },
  abstract: {
    name: "抽象艺术",
    images: [
      {
        id: "abstract-1",
        name: "几何抽象.jpg",
        url: "/colorful-geometric-abstract.png",
        size: 189440,
        type: "image/jpeg",
      },
      {
        id: "abstract-2",
        name: "色彩流动.jpg",
        url: "/flowing-colors-abstract.png",
        size: 234560,
        type: "image/jpeg",
      },
      {
        id: "abstract-3",
        name: "线条构成.jpg",
        url: "/linear-composition-abstract.png",
        size: 267890,
        type: "image/jpeg",
      },
      {
        id: "abstract-4",
        name: "色块拼接.jpg",
        url: "/color-blocks-abstract.png",
        size: 298765,
        type: "image/jpeg",
      },
      {
        id: "abstract-5",
        name: "动感抽象.jpg",
        url: "/dynamic-abstract-painting.png",
        size: 312450,
        type: "image/jpeg",
      },
      {
        id: "abstract-6",
        name: "纹理抽象.jpg",
        url: "/textured-abstract-art.png",
        size: 287650,
        type: "image/jpeg",
      },
      {
        id: "abstract-7",
        name: "光影抽象.jpg",
        url: "/light-shadow-abstract.png",
        size: 345670,
        type: "image/jpeg",
      },
      {
        id: "abstract-8",
        name: "螺旋构图.jpg",
        url: "/spiral-composition-abstract.png",
        size: 276540,
        type: "image/jpeg",
      },
      {
        id: "abstract-9",
        name: "对比抽象.jpg",
        url: "/contrast-abstract-painting.png",
        size: 298760,
        type: "image/jpeg",
      },
      {
        id: "abstract-10",
        name: "渐变抽象.jpg",
        url: "/gradient-abstract-art.png",
        size: 321450,
        type: "image/jpeg",
      },
    ],
  },
  stilllife: {
    name: "静物画",
    images: [
      { id: "stilllife-1", name: "花卉静物.jpg", url: "/floral-still-life.png", size: 278528, type: "image/jpeg" },
      {
        id: "stilllife-2",
        name: "水果静物.jpg",
        url: "/fruit-still-life.png",
        size: 234560,
        type: "image/jpeg",
      },
      {
        id: "stilllife-3",
        name: "花瓶静物.jpg",
        url: "/vase-still-life.png",
        size: 267890,
        type: "image/jpeg",
      },
      {
        id: "stilllife-4",
        name: "书籍静物.jpg",
        url: "/books-still-life.png",
        size: 298765,
        type: "image/jpeg",
      },
      {
        id: "stilllife-5",
        name: "茶具静物.jpg",
        url: "/tea-set-still-life.png",
        size: 312450,
        type: "image/jpeg",
      },
      {
        id: "stilllife-6",
        name: "蜡烛静物.jpg",
        url: "/candle-still-life.png",
        size: 287650,
        type: "image/jpeg",
      },
      {
        id: "stilllife-7",
        name: "乐器静物.jpg",
        url: "/musical-instrument-still-life.png",
        size: 345670,
        type: "image/jpeg",
      },
      {
        id: "stilllife-8",
        name: "古董静物.jpg",
        url: "/antique-still-life.png",
        size: 276540,
        type: "image/jpeg",
      },
      {
        id: "stilllife-9",
        name: "厨具静物.jpg",
        url: "/kitchen-utensils-still-life.png",
        size: 298760,
        type: "image/jpeg",
      },
      {
        id: "stilllife-10",
        name: "珠宝静物.jpg",
        url: "/jewelry-still-life.png",
        size: 321450,
        type: "image/jpeg",
      },
    ],
  },
}

export default function DataCollectionStep({
  onNext,
  selectedImages: selectedImageUrls,
  onImagesChange,
}: DataCollectionStepProps) {
  const [selectedImages, setSelectedImages] = useState<ImageData[]>([])
  const [uploadedImages, setUploadedImages] = useState<ImageData[]>([])
  const [selectedCategory, setSelectedCategory] = useState<string>("landscape")
  const fileInputRef = useRef<HTMLInputElement>(null)

  useEffect(() => {
    if (selectedImageUrls.length > 0 && selectedImages.length === 0) {
      const matchingImages: ImageData[] = []
      Object.values(sampleImageCategories).forEach((category) => {
        category.images.forEach((image) => {
          if (selectedImageUrls.includes(image.url)) {
            matchingImages.push(image)
          }
        })
      })
      setSelectedImages(matchingImages)
    }
  }, []) // Empty dependency array to run only on mount

  const handleImagesChange = useCallback(() => {
    const allImageUrls = [...selectedImages, ...uploadedImages].map((img) => img.url)
    onImagesChange(allImageUrls)
  }, [selectedImages, uploadedImages, onImagesChange])

  useEffect(() => {
    handleImagesChange()
  }, [handleImagesChange])

  const handleFileUpload = (event: React.ChangeEvent<HTMLInputElement>) => {
    const files = event.target.files
    if (!files) return

    Array.from(files).forEach((file) => {
      if (file.type.startsWith("image/")) {
        const reader = new FileReader()
        reader.onload = (e) => {
          const newImage: ImageData = {
            id: `upload-${Date.now()}-${Math.random()}`,
            name: file.name,
            url: e.target?.result as string,
            size: file.size,
            type: file.type,
          }
          setUploadedImages((prev) => [...prev, newImage])
        }
        reader.readAsDataURL(file)
      }
    })
  }

  const handleSampleSelect = (image: ImageData) => {
    if (!selectedImages.find((img) => img.id === image.id)) {
      setSelectedImages((prev) => [...prev, image])
    }
  }

  const handleSelectCategory = (categoryKey: string) => {
    const category = sampleImageCategories[categoryKey as keyof typeof sampleImageCategories]
    const newImages = category.images.filter((img) => !selectedImages.find((selected) => selected.id === img.id))
    setSelectedImages((prev) => [...prev, ...newImages])
  }

  const handleRemoveImage = (imageId: string) => {
    setSelectedImages((prev) => prev.filter((img) => img.id !== imageId))
    setUploadedImages((prev) => prev.filter((img) => img.id !== imageId))
  }

  const formatFileSize = (bytes: number) => {
    if (bytes === 0) return "0 Bytes"
    const k = 1024
    const sizes = ["Bytes", "KB", "MB", "GB"]
    const i = Math.floor(Math.log(bytes) / Math.log(k))
    return Number.parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + " " + sizes[i]
  }

  const allImages = [...selectedImages, ...uploadedImages]
  const totalSize = allImages.reduce((sum, img) => sum + img.size, 0)

  const handleNext = () => {
    if (onNext && allImages.length > 0) {
      onNext()
    }
  }

  return (
    <div className="space-y-6">
      {/* Introduction */}
      <Card className="bg-primary/5 border-primary/20">
        <CardHeader>
          <div className="flex items-center gap-2">
            <Info className="w-5 h-5 text-primary" />
            <CardTitle className="text-lg font-serif">什么是数据采集？</CardTitle>
          </div>
        </CardHeader>
        <CardContent>
          <p className="text-muted-foreground leading-relaxed">
            数据采集是AI绘画的第一步，也是最关键的步骤。就像人类艺术家需要观察大量的艺术作品来学习绘画技巧一样，
            AI模型也需要"看到"大量的图像数据来学习如何创作。数据的质量和多样性直接影响AI生成图像的效果。
          </p>
        </CardContent>
      </Card>

      {/* Data Collection Interface */}
      <Tabs defaultValue="samples" className="w-full">
        <TabsList className="grid w-full grid-cols-2">
          <TabsTrigger value="samples">示例数据集</TabsTrigger>
          <TabsTrigger value="upload">上传图像</TabsTrigger>
        </TabsList>

        <TabsContent value="samples" className="space-y-4">
          <Card>
            <CardHeader>
              <CardTitle className="font-serif">选择示例图像</CardTitle>
              <CardDescription>
                从我们准备的示例数据集中选择图像作为AI的学习材料，每个类别包含10张精选图像
              </CardDescription>
            </CardHeader>
            <CardContent>
              <Tabs value={selectedCategory} onValueChange={setSelectedCategory} className="w-full">
                <TabsList className="grid w-full grid-cols-4">
                  <TabsTrigger value="landscape">风景画</TabsTrigger>
                  <TabsTrigger value="portrait">人物肖像</TabsTrigger>
                  <TabsTrigger value="abstract">抽象艺术</TabsTrigger>
                  <TabsTrigger value="stilllife">静物画</TabsTrigger>
                </TabsList>

                {Object.entries(sampleImageCategories).map(([categoryKey, category]) => (
                  <TabsContent key={categoryKey} value={categoryKey} className="space-y-4">
                    <div className="flex justify-between items-center">
                      <h3 className="text-lg font-semibold">{category.name} (10张图像)</h3>
                      <Button
                        variant="outline"
                        onClick={() => handleSelectCategory(categoryKey)}
                        disabled={category.images.every((img) =>
                          selectedImages.find((selected) => selected.id === img.id),
                        )}
                      >
                        选择全部
                      </Button>
                    </div>
                    <div className="grid grid-cols-2 md:grid-cols-5 gap-4">
                      {category.images.map((image) => {
                        const isSelected = selectedImages.find((img) => img.id === image.id)
                        return (
                          <div key={image.id} className="relative group">
                            <div
                              className={`border-2 rounded-lg overflow-hidden transition-all ${
                                isSelected ? "border-primary shadow-lg" : "border-border hover:border-primary/50"
                              }`}
                            >
                              <img
                                src={image.url || "/placeholder.svg"}
                                alt={image.name}
                                className="w-full h-24 object-cover"
                              />
                              <div className="p-2">
                                <p className="text-xs font-medium truncate">{image.name}</p>
                                <p className="text-xs text-muted-foreground">{formatFileSize(image.size)}</p>
                              </div>
                            </div>

                            {isSelected ? (
                              <div className="absolute top-1 right-1">
                                <CheckCircle className="w-4 h-4 text-primary bg-white rounded-full" />
                              </div>
                            ) : (
                              <Button
                                size="sm"
                                onClick={() => handleSampleSelect(image)}
                                className="absolute top-1 right-1 opacity-0 group-hover:opacity-100 transition-opacity text-xs px-2 py-1 h-auto"
                              >
                                选择
                              </Button>
                            )}
                          </div>
                        )
                      })}
                    </div>
                  </TabsContent>
                ))}
              </Tabs>
            </CardContent>
          </Card>
        </TabsContent>

        <TabsContent value="upload" className="space-y-4">
          <Card>
            <CardHeader>
              <CardTitle className="font-serif">上传自定义图像</CardTitle>
              <CardDescription>上传你自己的图像来训练AI模型（支持 JPG, PNG, GIF 格式）</CardDescription>
            </CardHeader>
            <CardContent>
              <div
                className="border-2 border-dashed border-border rounded-lg p-8 text-center hover:border-primary/50 transition-colors cursor-pointer"
                onClick={() => fileInputRef.current?.click()}
              >
                <Upload className="w-12 h-12 text-muted-foreground mx-auto mb-4" />
                <h3 className="text-lg font-semibold mb-2">点击上传图像</h3>
                <p className="text-muted-foreground mb-4">或将图像文件拖拽到此区域</p>
                <Button variant="outline">
                  <ImageIcon className="w-4 h-4 mr-2" />
                  选择文件
                </Button>
                <input
                  ref={fileInputRef}
                  type="file"
                  multiple
                  accept="image/*"
                  onChange={handleFileUpload}
                  className="hidden"
                />
              </div>

              {uploadedImages.length > 0 && (
                <div className="mt-6">
                  <h4 className="font-semibold mb-3">已上传的图像</h4>
                  <div className="grid grid-cols-2 md:grid-cols-3 gap-4">
                    {uploadedImages.map((image) => (
                      <div key={image.id} className="relative group border rounded-lg overflow-hidden">
                        <img
                          src={image.url || "/placeholder.svg"}
                          alt={image.name}
                          className="w-full h-32 object-cover"
                        />
                        <div className="p-3">
                          <p className="text-sm font-medium truncate">{image.name}</p>
                          <p className="text-xs text-muted-foreground">{formatFileSize(image.size)}</p>
                        </div>
                        <Button
                          size="sm"
                          variant="destructive"
                          onClick={() => handleRemoveImage(image.id)}
                          className="absolute top-2 right-2 opacity-0 group-hover:opacity-100 transition-opacity"
                        >
                          <Trash2 className="w-4 h-4" />
                        </Button>
                      </div>
                    ))}
                  </div>
                </div>
              )}
            </CardContent>
          </Card>
        </TabsContent>
      </Tabs>

      {/* Dataset Statistics */}
      <Card>
        <CardHeader>
          <CardTitle className="font-serif">数据集统计</CardTitle>
          <CardDescription>当前选择的训练数据概览</CardDescription>
        </CardHeader>
        <CardContent>
          <div className="grid grid-cols-2 md:grid-cols-4 gap-4 mb-6">
            <div className="text-center p-4 bg-muted rounded-lg">
              <div className="text-2xl font-bold text-primary">{allImages.length}</div>
              <div className="text-sm text-muted-foreground">图像总数</div>
            </div>
            <div className="text-center p-4 bg-muted rounded-lg">
              <div className="text-2xl font-bold text-primary">{formatFileSize(totalSize)}</div>
              <div className="text-sm text-muted-foreground">总大小</div>
            </div>
            <div className="text-center p-4 bg-muted rounded-lg">
              <div className="text-2xl font-bold text-primary">{selectedImages.length}</div>
              <div className="text-sm text-muted-foreground">示例图像</div>
            </div>
            <div className="text-center p-4 bg-muted rounded-lg">
              <div className="text-2xl font-bold text-primary">{uploadedImages.length}</div>
              <div className="text-sm text-muted-foreground">上传图像</div>
            </div>
          </div>

          <div className="space-y-3">
            <div className="flex justify-between text-sm">
              <span>数据集完整度</span>
              <span>{Math.min(100, (allImages.length / 10) * 100).toFixed(0)}%</span>
            </div>
            <Progress value={Math.min(100, (allImages.length / 10) * 100)} className="h-2" />
            <p className="text-xs text-muted-foreground">建议至少选择 10 张图像以获得更好的训练效果</p>
          </div>

          {allImages.length > 0 && (
            <div className="mt-6">
              <h4 className="font-semibold mb-3">已选择的图像</h4>
              <div className="flex flex-wrap gap-2">
                {allImages.map((image) => (
                  <Badge key={image.id} variant="secondary" className="flex items-center gap-1">
                    {image.name}
                    <button onClick={() => handleRemoveImage(image.id)} className="ml-1 hover:text-destructive">
                      <Trash2 className="w-3 h-3" />
                    </button>
                  </Badge>
                ))}
              </div>
            </div>
          )}
        </CardContent>
      </Card>

      {/* Next Step Button */}
      <div className="flex justify-end">
        <Button size="lg" disabled={allImages.length === 0} className="font-semibold" onClick={handleNext}>
          继续到数据预处理
          <Eye className="w-4 h-4 ml-2" />
        </Button>
      </div>
    </div>
  )
}