package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// 说明文件候选名（小写）。需求中文件名为 reademe.md，同时兼容常见的 readme.md。
var readmeCandidates = []string{"reademe.md", "readme.md"}

// 允许预览的文本/代码类扩展名：以纯文本返回，避免 HTML 等同源 XSS。
var previewTextExts = map[string]bool{
	".txt": true, ".md": true, ".markdown": true, ".log": true, ".csv": true,
	".json": true, ".xml": true, ".yml": true, ".yaml": true, ".toml": true,
	".ini": true, ".conf": true, ".cfg": true, ".sql": true, ".env": true,
	".html": true, ".htm": true, ".css": true, ".js": true, ".ts": true,
	".jsx": true, ".tsx": true, ".vue": true, ".go": true, ".py": true,
	".java": true, ".c": true, ".h": true, ".cpp": true, ".hpp": true,
	".cc": true, ".cs": true, ".rb": true, ".rs": true, ".php": true,
	".sh": true, ".bash": true, ".bat": true, ".ps1": true, ".lua": true,
	".kt": true, ".swift": true, ".scala": true, ".r": true, ".dart": true,
	".gradle": true, ".gitignore": true,
}

// 浏览器原生支持内联展示的类型。
var previewInlineExts = map[string]string{
	".png":  "image/png",
	".jpg":  "image/jpeg",
	".jpeg": "image/jpeg",
	".gif":  "image/gif",
	".webp": "image/webp",
	".bmp":  "image/bmp",
	".ico":  "image/x-icon",
	".svg":  "image/svg+xml",
	".pdf":  "application/pdf",
}

// 预览文本文件的大小上限（10 MiB），防止把超大文件整体读入/传输。
const maxPreviewTextSize = 10 * 1024 * 1024

type server struct {
	root  string      // share 目录的绝对路径
	codec *tokenCodec // 路径 <-> 分享 key 的加解密器
}

type entry struct {
	Name    string    `json:"name"`
	Key     string    `json:"key"`  // 该条目的分享 key，传给任意接口即可定位
	Type    string    `json:"type"` // file | dir
	Size    int64     `json:"size"`
	ModTime time.Time `json:"mod_time"`
}

type dirDetail struct {
	Key     string  `json:"key"` // 当前目录的 key；根目录为空串（不传 key 即根目录）
	Name    string  `json:"name"`
	Readme  *string `json:"readme"` // 目录说明，无说明文件时为 null
	Entries []entry `json:"entries"`
}

func main() {
	addr := flag.String("addr", ":8080", "监听地址")
	root := flag.String("root", "./share", "共享目录（仅认领该目录及其子目录）")
	secretPath := flag.String("secret", "./.secret.key", "分享 key 的加密密钥文件（不存在则自动生成）")
	flag.Parse()

	absRoot, err := filepath.Abs(*root)
	if err != nil {
		log.Fatalf("解析共享目录失败: %v", err)
	}
	if err := os.MkdirAll(absRoot, 0o755); err != nil {
		log.Fatalf("创建共享目录失败: %v", err)
	}

	secret, err := loadOrCreateSecret(filepath.Clean(*secretPath))
	if err != nil {
		log.Fatalf("加载加密密钥失败: %v", err)
	}
	codec, err := newTokenCodec(secret)
	if err != nil {
		log.Fatalf("初始化加解密器失败: %v", err)
	}

	s := &server{root: absRoot, codec: codec}

	mux := http.NewServeMux()
	mux.HandleFunc("/api/get_dict_detail", s.handleDirDetail)
	mux.HandleFunc("/api/download", s.handleDownload)
	mux.HandleFunc("/api/preview", s.handlePreview)

	log.Printf("文件传输服务已启动: http://localhost%s", strings.TrimPrefix(*addr, ":"))
	log.Printf("共享目录: %s", absRoot)
	log.Printf("接口: /api/get_dict_detail | /api/download | /api/preview")
	if err := http.ListenAndServe(*addr, logRequests(mux)); err != nil {
		log.Fatal(err)
	}
}

func logRequests(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		next.ServeHTTP(w, r)
		log.Printf("%s %s", r.Method, r.URL.RequestURI())
	})
}

// resolve 把接口传入的相对路径限定在 share 根目录内，防止路径穿越。
func (s *server) resolve(rel string) (string, error) {
	rel = filepath.FromSlash(strings.TrimSpace(rel))
	if filepath.IsAbs(rel) || rel == "." {
		if rel == "." {
			rel = ""
		} else {
			return "", fmt.Errorf("非法路径")
		}
	}

	cleaned := filepath.Clean(rel)
	if cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("非法路径")
	}

	target := filepath.Join(s.root, cleaned)

	// 解析符号链接后再校验一次，确保仍在根目录内。
	realRoot, err := filepath.EvalSymlinks(s.root)
	if err != nil {
		return "", err
	}
	realTarget, err := filepath.EvalSymlinks(target)
	if err != nil {
		if os.IsNotExist(err) {
			return target, nil // 不存在的路径交由调用方通过 os.Stat 返回 404
		}
		return "", err
	}
	relToRoot, err := filepath.Rel(realRoot, realTarget)
	if err != nil || relToRoot == ".." || strings.HasPrefix(relToRoot, ".."+string(filepath.Separator)) {
		return "", fmt.Errorf("非法路径")
	}
	return target, nil
}

// toSlashRel 返回相对 share 根的 slash 风格路径。
func (s *server) toSlashRel(abs string) string {
	rel, err := filepath.Rel(s.root, abs)
	if err != nil {
		return ""
	}
	return filepath.ToSlash(rel)
}

// resolveKey 解密分享 key 并把解密出的路径限定在 share 根目录内。
// allowEmpty 为 true 时（目录浏览），空 key 表示根目录。
func (s *server) resolveKey(key string, allowEmpty bool) (string, error) {
	if key == "" {
		if allowEmpty {
			return s.resolve("")
		}
		return "", fmt.Errorf("缺少 key")
	}
	rel, err := s.codec.Open(strings.TrimSpace(key))
	if err != nil {
		// 不向前端泄露具体的解密/篡改细节
		return "", fmt.Errorf("无效的分享链接")
	}
	return s.resolve(rel)
}

func (s *server) handleDirDetail(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "仅支持 GET")
		return
	}

	dir, err := s.resolveKey(r.URL.Query().Get("key"), true)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}

	info, err := os.Stat(dir)
	if err != nil {
		writeError(w, http.StatusNotFound, "目录不存在")
		return
	}
	if !info.IsDir() {
		writeError(w, http.StatusBadRequest, "目标不是目录")
		return
	}

	items, err := os.ReadDir(dir)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "读取目录失败: "+err.Error())
		return
	}

	detail := dirDetail{
		Name:    info.Name(),
		Entries: make([]entry, 0, len(items)),
	}

	// 当前目录的相对路径；根目录的 key 留空（不传 key 即可访问根）。
	curRel := s.toSlashRel(dir)
	if curRel == "." || curRel == "" {
		curRel = ""
		detail.Name = "share"
	} else {
		detail.Key = s.codec.Seal(curRel)
	}

	for _, item := range items {
		full := filepath.Join(dir, item.Name())
		fi, err := item.Info()
		if err != nil {
			continue
		}
		t := "file"
		if item.IsDir() {
			t = "dir"
		}
		detail.Entries = append(detail.Entries, entry{
			Name:    item.Name(),
			Key:     s.codec.Seal(s.toSlashRel(full)),
			Type:    t,
			Size:    fi.Size(),
			ModTime: fi.ModTime(),
		})
	}

	// 目录在前、文件在后，各自按名称排序。
	sort.Slice(detail.Entries, func(i, j int) bool {
		a, b := detail.Entries[i], detail.Entries[j]
		if (a.Type == "dir") != (b.Type == "dir") {
			return a.Type == "dir"
		}
		return strings.ToLower(a.Name) < strings.ToLower(b.Name)
	})

	if readmePath := findReadme(dir, items); readmePath != "" {
		content, err := os.ReadFile(readmePath)
		if err == nil {
			text := string(content)
			detail.Readme = &text
		}
	}

	writeJSON(w, http.StatusOK, detail)
}

// findReadme 在目录条目里查找说明文件，大小写不敏感，优先 reademe.md。
func findReadme(dir string, items []os.DirEntry) string {
	lowerNames := make(map[string]string)
	for _, item := range items {
		if !item.IsDir() {
			lowerNames[strings.ToLower(item.Name())] = item.Name()
		}
	}
	for _, candidate := range readmeCandidates {
		if name, ok := lowerNames[candidate]; ok {
			return filepath.Join(dir, name)
		}
	}
	return ""
}

func (s *server) handleDownload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "仅支持 GET")
		return
	}

	file, err := s.resolveKey(r.URL.Query().Get("key"), false)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}

	info, err := os.Stat(file)
	if err != nil {
		writeError(w, http.StatusNotFound, "文件不存在")
		return
	}
	if info.IsDir() {
		writeError(w, http.StatusBadRequest, "不支持下载目录，仅可下载单个文件")
		return
	}

	f, err := os.Open(file)
	cd := mime.FormatMediaType("attachment", map[string]string{"filename": info.Name()})
	w.Header().Set("Content-Disposition", cd)
	http.ServeContent(w, r, info.Name(), info.ModTime(), f)
}

func (s *server) handlePreview(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "仅支持 GET")
		return
	}

	file, err := s.resolveKey(r.URL.Query().Get("key"), false)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}

	info, err := os.Stat(file)
	if err != nil {
		writeError(w, http.StatusNotFound, "文件不存在")
		return
	}
	if info.IsDir() {
		writeError(w, http.StatusBadRequest, "目录不支持预览")
		return
	}

	ext := strings.ToLower(filepath.Ext(file))

	f, err := os.Open(file)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "打开文件失败")
		return
	}
	defer f.Close()

	if contentType, ok := previewInlineExts[ext]; ok {
		w.Header().Set("Content-Type", contentType)
		w.Header().Set("Content-Disposition", "inline")
		http.ServeContent(w, r, info.Name(), info.ModTime(), f)
		return
	}

	if previewTextExts[ext] {
		if info.Size() > maxPreviewTextSize {
			writeError(w, http.StatusRequestEntityTooLarge, "文件过大，不支持预览，请直接下载")
			return
		}
		// 统一按纯文本返回，避免 html/svg 等在服务同源下执行造成 XSS。
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.Header().Set("Content-Disposition", "inline")
		http.ServeContent(w, r, info.Name(), info.ModTime(), f)
		return
	}

	writeError(w, http.StatusUnsupportedMediaType, "该文件类型不支持预览，请下载后查看")
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, msg string) {
	writeJSON(w, status, map[string]any{"code": status, "message": msg})
}