package utils

import (
	"archive/zip"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
)

type ExtractedFile struct {
	RelPath string
	AbsPath string
}

type ZipLimits struct {
	MaxFiles      int
	MaxTotalBytes int64
	MaxFileBytes  int64
}

func ExtractZip(zipPath string, destDir string, allowedExt map[string]struct{}, limits ZipLimits) ([]ExtractedFile, error) {
	/*
		ZIP 解压安全策略（默认拒绝风险输入）：
		- 路径穿越防护（ZipSlip）：拒绝包含 “..” 的条目，并校验最终落盘路径必须在 destDir 内
		- 总展开大小限制：防止解压炸弹
		- 单文件展开大小限制：防止单文件巨量占用磁盘
		- 文件数量限制：防止过多条目导致资源耗尽
		- 加密压缩包：不支持，发现加密条目直接失败（避免交互式密码输入/不确定性）
	*/
	r, err := zip.OpenReader(zipPath)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return nil, err
	}

	var (
		out          []ExtractedFile
		totalWritten int64
	)

	for _, f := range r.File {
		if f.FileInfo().IsDir() {
			continue
		}

		if limits.MaxFiles > 0 && len(out) >= limits.MaxFiles {
			return nil, fmt.Errorf("文件数量超出限制")
		}

		rel, err := cleanZipPath(f.Name)
		if err != nil {
			return nil, err
		}

		ext := strings.ToLower(filepath.Ext(rel))
		if _, ok := allowedExt[ext]; !ok {
			continue
		}

		dst := filepath.Join(destDir, filepath.FromSlash(rel))
		if !isWithinDir(destDir, dst) {
			return nil, fmt.Errorf("压缩包存在不安全路径")
		}

		if limits.MaxFileBytes > 0 && int64(f.UncompressedSize64) > limits.MaxFileBytes {
			return nil, fmt.Errorf("压缩包内文件过大")
		}

		if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
			return nil, err
		}

		rc, err := f.Open()
		if err != nil {
			if strings.Contains(strings.ToLower(err.Error()), "password") || strings.Contains(strings.ToLower(err.Error()), "encrypted") {
				return nil, fmt.Errorf("不支持加密压缩包")
			}
			return nil, err
		}

		written, err := writeFileLimited(dst, rc, limits.MaxFileBytes)
		_ = rc.Close()
		if err != nil {
			if errors.Is(err, errFileTooLarge) {
				return nil, fmt.Errorf("压缩包内文件过大")
			}
			return nil, err
		}

		totalWritten += written
		if limits.MaxTotalBytes > 0 && totalWritten > limits.MaxTotalBytes {
			return nil, fmt.Errorf("解压总大小超出限制")
		}

		out = append(out, ExtractedFile{RelPath: rel, AbsPath: dst})
	}

	if len(out) == 0 {
		return nil, fmt.Errorf("压缩包中未发现可识别的图片文件")
	}
	return out, nil
}

func cleanZipPath(name string) (string, error) {
	// ZIP 条目路径清洗：把 Windows 分隔符统一成 “/”，并拒绝任何包含 “..” 的段。
	s := strings.ReplaceAll(name, "\\", "/")
	for _, part := range strings.Split(s, "/") {
		if part == ".." {
			return "", fmt.Errorf("压缩包存在不安全路径")
		}
	}
	clean := path.Clean("/" + s)
	clean = strings.TrimPrefix(clean, "/")
	if clean == "" || clean == "." {
		return "", fmt.Errorf("压缩包路径非法")
	}
	if strings.Contains(clean, ":") {
		return "", fmt.Errorf("压缩包存在不安全路径")
	}
	return clean, nil
}

func isWithinDir(root, target string) bool {
	// 再次兜底校验：target 必须位于 root 下（防止 clean 规则被绕过）。
	rootClean := filepath.Clean(root)
	targetClean := filepath.Clean(target)
	rel, err := filepath.Rel(rootClean, targetClean)
	if err != nil {
		return false
	}
	return rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
}

var errFileTooLarge = errors.New("file too large")

func writeFileLimited(dst string, r io.Reader, maxBytes int64) (int64, error) {
	f, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	if maxBytes > 0 {
		n, err := io.CopyN(f, r, maxBytes+1)
		if err == nil {
			return n, errFileTooLarge
		}
		if errors.Is(err, io.EOF) {
			return n, nil
		}
		return n, err
	}

	return io.Copy(f, r)
}