main.go

package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/shirou/gopsutil/v3/disk"
)

//go:embed index.html style.css script.js
var content embed.FS

type DiskInfo struct {
	Path         string  `json:"path"`
	TotalSpace   uint64  `json:"totalSpace"`
	UsedSpace    uint64  `json:"usedSpace"`
	FreeSpace    uint64  `json:"freeSpace"`
	UsagePercent float64 `json:"usagePercent"`
}

type FileTypeStats struct {
	Extension  string  `json:"extension"`
	Count      int     `json:"count"`
	TotalSize  int64   `json:"totalSize"`
	Percentage float64 `json:"percentage"`
}

type FileInfo struct {
	Path      string    `json:"path"`
	Size      int64     `json:"size"`
	ModTime   time.Time `json:"modTime"`
	Extension string    `json:"extension"`
}

type DirectoryStats struct {
	Path      string `json:"path"`
	FileCount int    `json:"fileCount"`
	DirCount  int    `json:"dirCount"`
	TotalSize int64  `json:"totalSize"`
	Depth     int    `json:"depth"`
}

type ScanResult struct {
	Disks       []DiskInfo       `json:"disks"`
	FileTypes   []FileTypeStats  `json:"fileTypes"`
	TotalFiles  int              `json:"totalFiles"`
	TotalSize   int64            `json:"totalSize"`
	TopFiles    []FileInfo       `json:"topFiles"`
	RecentFiles []FileInfo       `json:"recentFiles"`
	TopDirs     []DirectoryStats `json:"topDirs"`
	MaxDepth    int              `json:"maxDepth"`
}

type scanState struct {
	mutex        sync.RWMutex
	fileTypes    sync.Map
	topFiles     []FileInfo
	recentFiles  []FileInfo
	dirStats     sync.Map
	totalFiles   int32
	totalSize    int64
	maxDepth     int
	fileInfoPool sync.Pool
}

type ScanProgress struct {
	CurrentPath  string `json:"currentPath"`
	ScannedFiles int32  `json:"scannedFiles"`
	CurrentDisk  string `json:"currentDisk"`
}

var (
	diskInfoMutex     sync.RWMutex
	progressChan      = make(chan ScanProgress, 100)
	diskInfoCache     []DiskInfo
	diskInfoCacheTime time.Time
)

const (
	maxScanDepth     = 100
	maxFileSize      = 100 << 30
	scanTimeout      = 15 * time.Minute
	maxWorkers       = 8
	bufferSize       = 10000
	progressInterval = time.Second
)

type fileTask struct {
	path string
	info os.FileInfo
}

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
}

func newScanState() *scanState {
	s := &scanState{
		topFiles:    make([]FileInfo, 0, 100),
		recentFiles: make([]FileInfo, 0, 100),
	}

	s.fileInfoPool.New = func() interface{} {
		return &FileInfo{}
	}

	return s
}

func calculateDirDepth(path string) int {
	return len(strings.Split(path, string(os.PathSeparator))) - 1
}

func updateTopFiles(files []FileInfo, newFile FileInfo, limit int) []FileInfo {
	files = append(files, newFile)
	sort.Slice(files, func(i, j int) bool {
		return files[i].Size > files[j].Size
	})
	if len(files) > limit {
		files = files[:limit]
	}
	return files
}

func updateRecentFiles(files []FileInfo, newFile FileInfo, limit int) []FileInfo {
	files = append(files, newFile)
	sort.Slice(files, func(i, j int) bool {
		return files[i].ModTime.After(files[j].ModTime)
	})
	if len(files) > limit {
		files = files[:limit]
	}
	return files
}

func getDiskInfo() []DiskInfo {
	diskInfoMutex.RLock()
	cacheDuration := 5 * time.Minute
	if diskInfoCache != nil && time.Since(diskInfoCacheTime) < cacheDuration {
		defer diskInfoMutex.RUnlock()
		return diskInfoCache
	}
	diskInfoMutex.RUnlock()

	diskInfoMutex.Lock()
	defer diskInfoMutex.Unlock()

	if diskInfoCache != nil && time.Since(diskInfoCacheTime) < cacheDuration {
		return diskInfoCache
	}

	var disks []DiskInfo

	if runtime.GOOS == "windows" {
		for _, drive := range "ABCDEFGHIJKLMNOPQRSTUVWXYZ" {
			path := string(drive) + ":\\"
			_, err := os.Stat(path)
			if err == nil {
				usage, err := disk.Usage(path)
				if err != nil {
					continue
				}
				total := usage.Total
				free := usage.Free
				used := total - free
				usagePercent := float64(used) / float64(total) * 100

				disks = append(disks, DiskInfo{
					Path:         path,
					TotalSpace:   total,
					UsedSpace:    used,
					FreeSpace:    free,
					UsagePercent: usagePercent,
				})
			}
		}
	} else {
		path := "/"
		usage, err := disk.Usage(path)
		if err == nil {
			total := usage.Total
			free := usage.Free
			used := total - free
			usagePercent := float64(used) / float64(total) * 100

			disks = append(disks, DiskInfo{
				Path:         path,
				TotalSpace:   total,
				UsedSpace:    used,
				FreeSpace:    free,
				UsagePercent: usagePercent,
			})
		}
	}

	diskInfoCache = disks
	diskInfoCacheTime = time.Now()
	return disks
}

func scanDirectory(path string) (*scanState, error) {
	state := newScanState()

	ctx, cancel := context.WithTimeout(context.Background(), scanTimeout)
	defer cancel()

	tasks := make(chan fileTask, bufferSize)
	results := make(chan error, maxWorkers)
	var wg sync.WaitGroup

	for i := 0; i < maxWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for task := range tasks {
				if err := processFile(ctx, state, task.path, task.info); err != nil {
					select {
					case results <- err:
					default:
					}
				}
			}
		}()
	}
	errChan := make(chan error, 1)
	go func() {
		for err := range results {
			if err != nil {
				select {
				case errChan <- err:
				default:
				}
			}
		}
	}()

	progressTicker := time.NewTicker(progressInterval)
	defer progressTicker.Stop()

	go func() {
		for {
			select {
			case <-progressTicker.C:
				select {
				case progressChan <- ScanProgress{
					CurrentPath:  path,
					ScannedFiles: atomic.LoadInt32(&state.totalFiles),
					CurrentDisk:  filepath.VolumeName(path),
				}:
				default:
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		if err != nil {
			return nil
		}

		if len(path) > 255 {
			log.Printf("跳过路径过长的文件: %s", path)
			return nil
		}

		if info.IsDir() {
			depth := calculateDirDepth(path)
			if depth > maxScanDepth {
				return filepath.SkipDir
			}

			if depth > state.maxDepth {
				state.maxDepth = depth
			}

			state.dirStats.Store(path, &DirectoryStats{
				Path:      path,
				Depth:     depth,
				FileCount: 0,
				DirCount:  0,
				TotalSize: 0,
			})
			return nil
		}

		select {
		case tasks <- fileTask{path: path, info: info}:
		case <-ctx.Done():
			return ctx.Err()
		}

		return nil
	})

	close(tasks)
	wg.Wait()
	close(results)

	select {
	case err := <-errChan:
		return nil, err
	default:
	}

	if err != nil {
		return nil, err
	}

	return state, nil
}

func processFile(ctx context.Context, state *scanState, path string, info os.FileInfo) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	if info.Size() > maxFileSize {
		log.Printf("跳过大文件: %s (大小: %d)", path, info.Size())
		return nil
	}

	ext := filepath.Ext(path)
	if ext == "" {
		ext = "no extension"
	}

	if value, exists := state.fileTypes.Load(ext); exists {
		stats := value.(*FileTypeStats)
		stats.Count++
		stats.TotalSize += info.Size()
	} else {
		state.fileTypes.Store(ext, &FileTypeStats{
			Extension: ext,
			Count:     1,
			TotalSize: info.Size(),
		})
	}

	atomic.AddInt64(&state.totalSize, info.Size())
	atomic.AddInt32(&state.totalFiles, 1)

	fileInfo := FileInfo{
		Path:      path,
		Size:      info.Size(),
		ModTime:   info.ModTime(),
		Extension: ext,
	}

	state.mutex.Lock()
	state.topFiles = updateTopFiles(state.topFiles, fileInfo, 30)
	state.recentFiles = updateRecentFiles(state.recentFiles, fileInfo, 10)
	state.mutex.Unlock()

	dirPath := filepath.Dir(path)
	currentPath := dirPath
	for {
		if dirStat, exists := state.dirStats.Load(currentPath); exists {
			state.mutex.Lock()
			dirStat := dirStat.(*DirectoryStats)
			dirStat.FileCount++
			dirStat.TotalSize += info.Size()
			state.mutex.Unlock()
		}
		parent := filepath.Dir(currentPath)
		if parent == currentPath {
			break
		}
		currentPath = parent
	}
	if atomic.LoadInt32(&state.totalFiles)%1000 == 0 {
		select {
		case progressChan <- ScanProgress{
			CurrentPath:  path,
			ScannedFiles: state.totalFiles,
			CurrentDisk:  filepath.VolumeName(path),
		}:
		default:
		}
	}
	return nil
}

func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}
		next(w, r)
	}
}

func handleScan(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	var cancel context.CancelFunc
	ctx, cancel = context.WithTimeout(ctx, scanTimeout)
	defer cancel()

	if r.Method != http.MethodGet {
		http.Error(w, "方法不允许", http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	log.Println("开始扫描磁盘...")
	disks := getDiskInfo()
	if len(disks) == 0 {
		log.Println("错误：未找到可用磁盘")
		http.Error(w, "未找到可用磁盘", http.StatusInternalServerError)
		return
	}

	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	if m.Alloc > uint64(runtime.NumCPU()*1024*1024*1024) {
		http.Error(w, "系统资源不足", http.StatusServiceUnavailable)
		return
	}

	var combinedState scanState

	for _, disk := range disks {
		select {
		case <-ctx.Done():
			http.Error(w, "扫描超时", http.StatusGatewayTimeout)
			return
		default:
		}

		log.Printf("正在扫描磁盘: %s", disk.Path)
		state, err := scanDirectory(disk.Path)
		if err != nil {
			log.Printf("扫描磁盘 %s 时出错: %v", disk.Path, err)
			continue
		}
		combinedState.totalFiles += state.totalFiles
		combinedState.totalSize += state.totalSize
		if state.maxDepth > combinedState.maxDepth {
			combinedState.maxDepth = state.maxDepth
		}
		for _, file := range state.topFiles {
			combinedState.topFiles = updateTopFiles(combinedState.topFiles, file, 30)
		}
		for _, file := range state.recentFiles {
			combinedState.recentFiles = updateRecentFiles(combinedState.recentFiles, file, 10)
		}
		state.fileTypes.Range(func(key, value interface{}) bool {
			ext := key.(string)
			stats := value.(*FileTypeStats)
			if existing, ok := combinedState.fileTypes.Load(ext); ok {
				existingStats := existing.(*FileTypeStats)
				existingStats.Count += stats.Count
				existingStats.TotalSize += stats.TotalSize
			} else {
				combinedState.fileTypes.Store(ext, stats)
			}
			return true
		})
		state.dirStats.Range(func(key, value interface{}) bool {
			path := key.(string)
			stats := value.(*DirectoryStats)
			if existing, ok := combinedState.dirStats.Load(path); ok {
				existingStats := existing.(*DirectoryStats)
				existingStats.FileCount += stats.FileCount
				existingStats.TotalSize += stats.TotalSize
				if stats.Depth > existingStats.Depth {
					existingStats.Depth = stats.Depth
				}
			} else {
				combinedState.dirStats.Store(path, stats)
			}
			return true
		})
	}

	var fileTypes []FileTypeStats
	combinedState.fileTypes.Range(func(key, value interface{}) bool {
		stats := value.(*FileTypeStats)
		stats.Percentage = float64(stats.TotalSize) / float64(combinedState.totalSize) * 100
		fileTypes = append(fileTypes, *stats)
		return true
	})

	sort.Slice(fileTypes, func(i, j int) bool {
		return fileTypes[i].TotalSize > fileTypes[j].TotalSize
	})

	var topDirs []DirectoryStats
	combinedState.dirStats.Range(func(key, value interface{}) bool {
		dir := value.(*DirectoryStats)
		topDirs = append(topDirs, *dir)
		return true
	})

	sort.Slice(topDirs, func(i, j int) bool {
		return topDirs[i].TotalSize > topDirs[j].TotalSize
	})

	if len(topDirs) > 30 {
		topDirs = topDirs[:30]
	}

	if len(combinedState.topFiles) > 1000 {
		combinedState.topFiles = combinedState.topFiles[:1000]
	}

	result := ScanResult{
		Disks:       disks,
		FileTypes:   fileTypes,
		TotalFiles:  int(combinedState.totalFiles),
		TotalSize:   combinedState.totalSize,
		TopFiles:    combinedState.topFiles,
		RecentFiles: combinedState.recentFiles,
		MaxDepth:    combinedState.maxDepth,
		TopDirs:     topDirs,
	}

	if err := json.NewEncoder(w).Encode(result); err != nil {
		log.Printf("JSON编码错误: %v", err)
		http.Error(w, "内部服务器错误", http.StatusInternalServerError)
		return
	}
	log.Printf("扫描完成: 总文件数 %d, 总大小 %d bytes", combinedState.totalFiles, combinedState.totalSize)
}

func handleProgress(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	_, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "SSE not supported", http.StatusInternalServerError)
		return
	}

	done := make(chan struct{})
	defer close(done)

	go func() {
		<-ctx.Done()
		select {
		case <-done:
		default:
			close(done)
		}
	}()

	for progress := range progressChan {
		select {
		case <-done:
			return
		default:
			data, err := json.Marshal(progress)
			if err != nil {
				log.Printf("进度序列化错误: %v", err)
				continue
			}
			fmt.Fprintf(w, "data: %s\n\n", data)
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}
}

func main() {
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
	fsys := http.FileServer(http.FS(content))
	http.HandleFunc("/scan", corsMiddleware(handleScan))
	http.HandleFunc("/progress", corsMiddleware(handleProgress))
	http.Handle("/", fsys)

	log.Println("Server starting on http://localhost:8080")
	server := &http.Server{
		Addr:              ":8080",
		ReadHeaderTimeout: 3 * time.Second,
		ReadTimeout:       20 * time.Minute,
		WriteTimeout:      20 * time.Minute,
	}

	go func() {
		<-stop
		log.Println("正在关闭服务器...")
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil {
			log.Printf("服务器关闭出错: %v", err)
		}
	}()

	if err := server.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatal(err)
	}
	log.Println("服务器已关闭")
}