package main

import (
	"archive/zip"
	"bufio"
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	_ "embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"hash"
	"html"
	"html/template"
	"io"
	"io/fs"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

const ver = "26.4.26"

// rowTemplate 定义文件列表中的每一行数据结构
type rowTemplate struct {
	Name string       // 文件/文件夹名称
	Href template.URL // 链接地址
	Size string       // 项目大小
	Ext  string       // 项目扩展名 (文件夹是folder)
	Desc string       // 项目描述
	Mime string       // 项目mime
}

// pageTemplate 定义整个页面渲染所需的数据结构
type pageTemplate struct {
	Title       string        // 页面标题
	Path        template.HTML // 页面路径
	ExtraPath   template.HTML // URL前缀路径
	Ro          bool          // 是否只读模式（不允许上传、重命名、移动等操作）
	RowsFiles   []rowTemplate // 文件列表
	RowsFolders []rowTemplate // 文件夹列表
	TimeStamp   string        //时间戳
}

// 命令行参数定义
var host = flag.String("h", "127.0.0.1", "监听的主机地址")                                               // 监听的主机地址
var port = flag.String("p", "8001", "监听的端口")                                                      // 监听的端口
var extraPath = flag.String("prefix", "/", "afi 的 URL 访问路径前缀, 例如 /afi/ (斜杠很重要)")                  // URL前缀路径
var symlinks = flag.Bool("symlinks", false, "跟随符号链接 警告: 启用后符号链接可能会跳出已定义的\"根目录\"") // 是否跟随符号链接
var verb = flag.Bool("verb", false, "详细输出")                                                       // 是否输出详细日志
var skipHidden = flag.Bool("k", false, "跳过隐藏文件")                                                 // 是否跳过隐藏文件（以.开头）
// var ro = flag.Bool("ro", false, "只读模式(无法修改文件系统)")                                                         // 是否只读模式
var title = flag.String("title", "%PATH%", "页面标题, 用 %PATH% 指代完整路径, %ITEM% 指代末端文件/目录名, 不会泄露根目录目录名")
var authstr = flag.String("auth", `{"admin": "password"}`, "可写用户的认证数据 (也可以用环境变量 AFI_AUTH 设置, AFI_AUTH 是最优先且最安全的)") // 这个flag只能用来调试 不然服务器一个ps就泄露了
var dav = flag.Bool("dav", false, "启用 WebDAV 接口")

// rpcCall 定义RPC调用的JSON数据结构
// 用于处理前端发送的远程调用请求（创建目录、移动、删除、计算校验和）
type rpcCall struct {
	Call string   `json:"call"` // 调用的方法名
	Args []string `json:"args"` // 方法参数
}

var rootPath = ""        // 共享目录的根路径(文件系统中的绝对路径)
var authTable = make(map[string]string)

// gzipWriterPool 复用 gzip.Writer，减少每次请求的分配开销
var gzipWriterPool = sync.Pool{
	New: func() any {
		w, err := gzip.NewWriterLevel(nil, gzip.BestSpeed)
		if err != nil {
			panic(err)
		}
		return w
	},
}

// dirCacheEntry 缓存目录列表数据，避免重复 ReadDir + Stat + 排序
type dirCacheEntry struct {
	etag    string        // ETag 值
	folders []rowTemplate // 文件夹列表（不含 "../"）
	files   []rowTemplate // 文件列表
	mtime   time.Time     // 目录修改时间（用于验证）
}

var (
	dirCacheMu sync.RWMutex
	dirCache   = make(map[string]*dirCacheEntry, 128)
)

// check 检查错误, 如果错误不为nil则panic
// 用于简化错误处理
func check(e error) {
	if e != nil {
		panic(e)
	}
}

// exitPath 延迟执行的错误恢复和日志记录函数
// 用于捕获panic并返回500错误, 同时记录日志
// w: HTTP响应写入器
// s: 日志信息参数
func exitPath(w http.ResponseWriter, s ...any) {
	if r := recover(); r != nil {
		// 发生panic时记录错误并返回500状态码
		log.Println("error", s, r)
		w.WriteHeader(500)
	} else if *verb {
		// 详细模式下记录正常日志
		log.Println(s...)
	}
}

// humanize 将字节数转换为人类可读的格式
// 例如: 1024 -> "1.0k", 1048576 -> "1.0M"
// bytes: 文件大小（字节）
// 返回: 格式化后的字符串（如"1.5M"）
var humanUnits = []string{"B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"}

func humanize(bytes int64) string {
	if bytes < 1024 {
		return strconv.FormatInt(bytes, 10) + "B"
	}
	f := float64(bytes) / 1024
	u := 1
	for f >= 1024 && u < len(humanUnits)-1 {
		f /= 1024
		u++
	}
	return strconv.FormatFloat(f, 'f', 1, 64) + humanUnits[u]
}

// parseDescriptIon 只打开一次 DESCRIPT.ION 并返回 map，大幅提升目录列表性能
func parseDescriptIon(dirPath string) map[string]string {
	descMap := make(map[string]string)
	descFilePath := filepath.Join(dirPath, "DESCRIPT.ION")

	// 默认描述, 可被覆写
	descMap["DESCRIPT.ION"] = "备注元数据"
	descMap["README.md"] = "自述文件"
	descMap["LICENSE"] = "许可证"
	descMap[".gitignore"] = "Git/文件忽略模式"
	descMap[".git"] = "Git/裸仓库目录"

	if file, err := os.Open(descFilePath); err == nil {
		defer file.Close()

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" {
				continue
			}

			var targetName, description string

			// 处理带双引号的文件名格式，例如："My Photo.jpg" This is a description
			if strings.HasPrefix(line, `"`) {
				end := strings.Index(line[1:], `"`)
				if end != -1 {
					targetName = line[1 : end+1]
					description = strings.TrimSpace(line[end+2:])
				}
			} else {
				// 普通格式：filename description
				parts := strings.SplitN(line, " ", 2)
				targetName = parts[0]
				if len(parts) > 1 {
					description = strings.TrimSpace(parts[1])
				}
			}

			if targetName != "" {
				// 清理可能的控制字符
				description = strings.Split(description, "\x04")[0]
				// 严格大小写：去掉可能的引号后，直接使用原始名称作为 key
				cleanName := strings.Trim(targetName, `"`)
				descMap[cleanName] = strings.TrimSpace(description)
			}
		}
	}

	return descMap
}

// replyList 生成目录列表页面并返回
// 读取指定目录下的文件和文件夹, 构建页面模板数据并渲染
// w: HTTP响应写入器
// r: HTTP请求
// fullPath: 目录的绝对路径
// path: 请求的URL路径
// dirFi: 目录的 FileInfo（由调用者提供，避免重复 stat）
func replyList(w http.ResponseWriter, r *http.Request, fullPath string, path string, dirFi fs.FileInfo) {
	dirMtime := dirFi.ModTime()
	// ETag 基于目录 mtime，目录内容变化时自动失效
	etag := fmt.Sprintf(`W/"%x"`, dirMtime.UnixNano())

	// 缓存查找
	dirCacheMu.RLock()
	entry, ok := dirCache[fullPath]
	dirCacheMu.RUnlock()

	var folders, files []rowTemplate

	if ok && entry.mtime.Equal(dirMtime) {
		// 缓存命中 → ETag 检查
		if r.Header.Get("If-None-Match") == entry.etag {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		w.Header().Set("ETag", entry.etag)
		folders = entry.folders
		files = entry.files
	} else {
		// 缓存未命中 → 读取目录并构建
		entries, err := os.ReadDir(fullPath)
		check(err)

		descMap := parseDescriptIon(fullPath)

		sort.Slice(entries, func(i, j int) bool {
			nameI := entries[i].Name()
			nameJ := entries[j].Name()

			if strings.EqualFold(nameI, "DESCRIPT.ION") && !strings.EqualFold(nameJ, "DESCRIPT.ION") {
				return false
			}
			if strings.EqualFold(nameJ, "DESCRIPT.ION") && !strings.EqualFold(nameI, "DESCRIPT.ION") {
				return true
			}

			return nameI < nameJ
		})

		mimeCache := make(map[string]string, 16)

		folders = make([]rowTemplate, 0, len(entries))
		files = make([]rowTemplate, 0, len(entries))

		for _, entry := range entries {
			name := entry.Name()

			if *skipHidden && strings.HasPrefix(name, ".") {
				continue
			}

			if !*symlinks && entry.Type()&os.ModeSymlink != 0 {
				continue
			}

			var info fs.FileInfo
			if !entry.IsDir() {
				var statErr error
				info, statErr = entry.Info()
				if statErr != nil {
					log.Println("error - cant stat a file", statErr)
					continue
				}
			}

			href := url.PathEscape(name)
			if entry.IsDir() && strings.HasPrefix(href, "/") {
				href = strings.TrimPrefix(href, "/")
			}

			desc := descMap[name]

			if entry.IsDir() {
				folders = append(folders, rowTemplate{
					Name: html.EscapeString(name + "/"),
					Href: template.URL(href),
					Size: "4.0kB",
					Ext:  "folder",
					Desc: html.EscapeString(desc),
					Mime: "inode/directory",
				})
			} else {
				ext := strings.TrimPrefix(filepath.Ext(name), ".")
				filemime, ok := mimeCache[ext]
				if !ok {
					filemime = mime.TypeByExtension("." + ext)
					if filemime != "" {
						if idx := strings.IndexByte(filemime, ';'); idx >= 0 {
							filemime = strings.TrimSpace(filemime[:idx])
						}
					} else {
						if name == "DESCRIPT.ION" {
							filemime = "text/plain"
						} else {
							filemime = "unknown/unknown"
						}
					}
					mimeCache[ext] = filemime
				}

				files = append(files, rowTemplate{
					Name: html.EscapeString(name),
					Href: template.URL(href),
					Size: humanize(info.Size()),
					Ext:  html.EscapeString(ext),
					Desc: html.EscapeString(desc),
					Mime: filemime,
				})
			}
		}

		// 入缓存
		dirCacheMu.Lock()
		dirCache[fullPath] = &dirCacheEntry{etag: etag, folders: folders, files: files, mtime: dirMtime}
		dirCacheMu.Unlock()
		w.Header().Set("ETag", etag)
	}

	// --- 构建页面（请求相关部分）---
	if !strings.HasSuffix(path, "/") {
		path += "/"
	}

	pathtext := "/" + strings.TrimPrefix(path, *extraPath)

	p := pageTemplate{}
	now := time.Now().UTC()
	p.TimeStamp = now.Format(time.RFC3339)

	rowCap := len(folders) + len(files) + 1
	p.RowsFolders = make([]rowTemplate, 0, rowCap)
	p.RowsFiles = make([]rowTemplate, 0, rowCap)

	if path != *extraPath {
		p.RowsFolders = append(p.RowsFolders, rowTemplate{"../", "..", "4.0kB", "folder", "返回上级目录", "inode/directory"})
	}

	p.RowsFolders = append(p.RowsFolders, folders...)
	p.RowsFiles = append(p.RowsFiles, files...)

	p.ExtraPath = template.HTML(html.EscapeString(*extraPath))
	p.Path = template.HTML(html.EscapeString(pathtext))
	if pathtext == "/" {
		p.Title = strings.ReplaceAll(strings.ReplaceAll(*title, "%PATH%", pathtext), "%ITEM%", "/")
	} else {
		p.Title = strings.ReplaceAll(strings.ReplaceAll(*title, "%PATH%", pathtext), "%ITEM%", filepath.Base(fullPath))
	}

	// 渲染到 buffer，按大小决定是否 gzip（<1KB 直接发，压缩了更慢）
	var buf bytes.Buffer
	tmpl.Execute(&buf, p)
	body := buf.Bytes()

	w.Header().Set("Content-Type", "text/html")

	if len(body) > 1024 && strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Add("Content-Encoding", "gzip")
		gz := gzipWriterPool.Get().(*gzip.Writer)
		gz.Reset(w)
		gz.Write(body)
		gz.Close()
		gzipWriterPool.Put(gz)
	} else {
		w.Write(body)
	}
}

// doContent 处理所有内容请求的核心函数
// 根据请求路径判断是目录还是文件, 目录则列出文件列表, 文件则提供下载
// w: HTTP响应写入器
// r: HTTP请求
func doContent(w http.ResponseWriter, r *http.Request) {
	// 如果URL不以配置的前缀开头, 重定向到正确的前缀
	if !strings.HasPrefix(r.URL.Path, *extraPath) {
		http.Redirect(w, r, *extraPath, http.StatusFound)
		return
	}

	// 解码URL中的HTML实体
	path := html.UnescapeString(r.URL.Path)
	defer exitPath(w, "get content", path)
	// 获取安全的绝对路径
	fullPath := enforcePath(path)

	// 只 Open + Stat 一次，避免 FileServer 内部二次 stat
	f, err := os.Open(fullPath)
	check(err)
	fi, err := f.Stat()
	check(err)

	if fi.IsDir() {
		f.Close()
		replyList(w, r, fullPath, path, fi)
	} else {
		w.Header().Set("Content-Type", mime.TypeByExtension(filepath.Ext(fi.Name())))
		http.ServeContent(w, r, fi.Name(), fi.ModTime(), f)
		f.Close()
	}
}

func checkAuth(username string, password string) bool {
	return authTable[username] == password
	// WARN: 这么搞不能防御 timing attack
	// 切勿用于生产环境
}

func checkAuthRequest(w http.ResponseWriter, r *http.Request) bool {
	username, password, ok := r.BasicAuth()
	if !ok {
		return false
	}
	if checkAuth(username, password) {
		return true
	} else {
		return false
	}
}

// 拒绝现代JWT, 古法auth, 我们有先进的https
// 这样不利好防御 XSS 虽说已经把模板搞安全了
// 但是至于XSS, 用户都在浏览器随便运行tampermonkey了我还能说什么呢, 说不定是用户想写脚本调试呢, 我把xss预防了让他们费劲去找token反而不好
// 这才叫利好curl, 而且还防 CSRF
func authLogin(w http.ResponseWriter, r *http.Request) {
	username, password, ok := r.BasicAuth()
	if !ok {
		w.Header().Set("WWW-Authenticate", `Basic realm="AFI File Manager"`)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	if checkAuth(username, password) {
		http.Redirect(w, r, *extraPath, http.StatusFound)
	} else {
		w.Header().Set("WWW-Authenticate", `Basic realm="AFI File Manager"`)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
	}
}
func authState(w http.ResponseWriter, r *http.Request) {
	username, password, ok := r.BasicAuth()
	if !ok {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	if checkAuth(username, password) {
		w.WriteHeader(http.StatusOK)
	} else {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
	}
}

// upload 处理文件上传请求
// 从multipart请求中读取文件内容并保存到指定路径
// w: HTTP响应写入器
// r: HTTP请求
func upload(w http.ResponseWriter, r *http.Request) {
	if !checkAuthRequest(w, r) {
		return
	}
	// 从header中获取目标路径
	path := r.Header.Get("afi-path")
	defer exitPath(w, "upload", path)

	path, err := url.PathUnescape(path)
	check(err)
	// 创建multipart reader解析上传的文件
	reader, err := r.MultipartReader()
	check(err)
	part, err := reader.NextPart()
	if err != nil && err != io.EOF {
		check(err)
	}
	// 创建目标文件
	dst, err := os.Create(enforcePath(path))
	check(err)
	// 将上传内容复制到目标文件
	io.Copy(dst, part)
	w.Write([]byte("ok"))
}

// zipRPC 处理ZIP打包下载请求
// 将指定目录或文件打包成ZIP并下载
// w: HTTP响应写入器
// r: HTTP请求
func zipRPC(w http.ResponseWriter, r *http.Request) {
	if !checkAuthRequest(w, r) {
		return
	}
	// 获取要打包的路径和ZIP文件名
	zipPath := r.URL.Query().Get("zipPath")
	zipName := r.URL.Query().Get("zipName")
	defer exitPath(w, "zip", zipPath)
	zipFullPath := enforcePath(zipPath)
	_, err := os.Lstat(zipFullPath)
	check(err)

	// 设置下载响应头
	w.Header().Add("Content-Disposition", "attachment; filename=\""+zipName+".zip\"")
	zipWriter := zip.NewWriter(w)
	defer zipWriter.Close()

	// 遍历目录树, 将文件添加到ZIP中
	err = filepath.Walk(zipFullPath, func(path string, f fs.FileInfo, err error) error {
		check(err)
		// 跳过目录（目录本身不加入ZIP）
		if f.IsDir() {
			return nil
		}

		// 计算相对路径作为ZIP内的路径
		rel, err := filepath.Rel(zipFullPath, path)
		check(err)
		// 跳过隐藏文件
		if *skipHidden && (strings.HasPrefix(rel, ".") || strings.HasPrefix(f.Name(), ".")) {
			return nil
		}
		// 不允许打包符号链接
		if f.Mode()&os.ModeSymlink != 0 {
			panic(errors.New("symlink not allowed in zip downloads"))
		}

		// 创建ZIP文件头
		header, err := zip.FileInfoHeader(f)
		check(err)
		header.Name = filepath.ToSlash(rel) // 统一使用斜杠分隔
		header.Method = zip.Store           // 使用存储模式 (不压缩)
		headerWriter, err := zipWriter.CreateHeader(header)
		check(err)
		// 打开文件并复制到ZIP中
		file, err := os.Open(path)
		check(err)
		_, err = io.Copy(headerWriter, file)
		file.Close()
		check(err)
		return nil
	})

	check(err)
}

// rpc 处理RPC请求
// 支持创建目录、移动/重命名、删除文件、计算校验和等操作
// w: HTTP响应写入器
// r: HTTP请求
func rpc(w http.ResponseWriter, r *http.Request) {
	if !checkAuthRequest(w, r) {
		return
	}
	var err error
	var rpc rpcCall
	defer exitPath(w, "rpc", &rpc)

	// 读取并解析JSON请求体
	bodyBytes, err := io.ReadAll(r.Body)
	check(err)
	json.Unmarshal(bodyBytes, &rpc)
	ret := []byte("ok")

	// 根据调用的方法执行相应操作
	switch rpc.Call {
	case "touch":
		// 创建目录（包括所有父目录）
		_, err = os.OpenFile(enforcePath(rpc.Args[0]), os.O_RDONLY|os.O_CREATE, 0777)
	case "mkdirp":
		// 创建目录（包括所有父目录）
		err = os.MkdirAll(enforcePath(rpc.Args[0]), os.ModePerm)
	case "mv":
		// 移动或重命名文件/目录
		err = os.Rename(enforcePath(rpc.Args[0]), enforcePath(rpc.Args[1]))
	case "rm":
		// 删除文件或目录（递归删除）
		err = os.RemoveAll(enforcePath(rpc.Args[0]))
	case "sum":
		// 计算文件的校验和（md5, sha1, sha256, sha512）
		file, err := os.Open(enforcePath(rpc.Args[0]))
		check(err)
		var hash hash.Hash
		// 根据参数选择哈希算法
		switch rpc.Args[1] {
		case "md5":
			hash = md5.New()
		case "sha1":
			hash = sha1.New()
		case "sha256":
			hash = sha256.New()
		case "sha512":
			hash = sha512.New()
		}
		_, err = io.Copy(hash, file)
		check(err)
		checksum := hash.Sum(nil)
		// 将哈希值转换为十六进制字符串
		ret = make([]byte, hex.EncodedLen(len(checksum)))
		hex.Encode(ret, checksum)
	}

	check(err)
	w.Write(ret)
}

// enforcePath 确保请求的路径是安全的
// 防止路径遍历攻击, 确保访问路径在共享目录内
// p: 请求的路径
// 返回: 安全的绝对路径
func enforcePath(p string) string {
	joined := filepath.Join(rootPath, strings.TrimPrefix(p, *extraPath))
	fp := filepath.Clean(joined)

	if !strings.HasPrefix(fp, rootPath) || (*skipHidden && strings.Contains(p, "/.")) {
		panic(errors.New("invalid path"))
	}

	if *symlinks {
		if sl, err := filepath.EvalSymlinks(fp); err == nil && !strings.HasPrefix(sl, rootPath) {
			panic(errors.New("invalid path"))
		}
	}

	return fp
}

func main() {
	// 解析命令行参数
	flag.Parse()
	if envAuth := os.Getenv("AFI_AUTH"); envAuth != "" {
		*authstr = envAuth
	}
	if len(flag.Args()) == 1 {
		rootPath = flag.Args()[0]
	} else {
		fmt.Printf("Usage: ./afi [OPTION]... [ROOTDIR]\n\n")
		flag.PrintDefaults()
		os.Exit(1)
	}
	err := json.Unmarshal([]byte(*authstr), &authTable)
	check(err)

	// 获取共享目录的绝对路径
	rootPath, err = filepath.Abs(rootPath)
	check(err)

	// 创建HTTP服务器
	server := &http.Server{
		Addr:              *host + ":" + *port,
		Handler:           nil,
		ReadHeaderTimeout: 10 * time.Second,
		ReadTimeout:       60 * time.Second,
		WriteTimeout:      60 * time.Second,
		IdleTimeout:       120 * time.Second,
	}

	// 注册路由处理器
	http.HandleFunc(*extraPath+"rpc", rpc)         // RPC接口 (增删移与校验和)
	http.HandleFunc(*extraPath+"post", upload)     // 文件上传接口
	if *dav {
		http.HandleFunc(*extraPath+"dav/", webdavHandler) // WebDAV 接口
	}
	http.HandleFunc(*extraPath+"auth", authState)  // 仅检查登录接口(不会弹窗)
	http.HandleFunc(*extraPath+"login", authLogin) // 登录接口
	http.HandleFunc(*extraPath+"zip", zipRPC)      // ZIP打包下载接口
	http.HandleFunc("/", doContent)                // 主内容处理接口

	// 输出启动信息
	fmt.Printf("Agile File Interface 版本 %s\n", ver)
	fmt.Printf("AFI 是 Gossa 的增强分支, 但是并不完全向下兼容原版 Gossa 的参数\n")
	fmt.Printf("AFI 已启动, 根目录为 %s\n", rootPath)
	fmt.Printf("详细输出: %t, 符号链接跟随: %t, 跳过隐藏文件: %t, WebDAV: %t\n", *verb, *symlinks, *skipHidden, *dav)
	fmt.Printf("身份验证信息: %s\n", *authstr)
	fmt.Printf("正在监听: http://%s:%s%s\n", *host, *port, *extraPath)

	// 启动HTTP服务
	if err = server.ListenAndServe(); err != http.ErrServerClosed {
		check(err)
	}
}