package main

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

const (
	chunkDir       = "./data/chunks"
	chunkThreshold = 5 << 30 // >5GB 才走分片
	chunkSize      = 1 << 30 // 每片 1GB
)

var uploadIDRe = regexp.MustCompile(`^[a-zA-Z0-9_-]{8,64}$`)

type chunkMeta struct {
	UploadID    string       `json:"upload_id"`
	Filename    string       `json:"filename"`
	TotalSize   int64        `json:"total_size"`
	TotalChunks int          `json:"total_chunks"`
	Received    map[int]bool `json:"received"`
	CreatedAt   time.Time    `json:"created_at"`
	Reserved    int64        `json:"reserved"`
}

// 全局 map 存 in-flight 上传; 持锁只做元数据操作, 1GB 分片 io.Copy 在锁外进行.
// ponytail: 全局锁, 若并发大文件上传成为瓶颈再换 per-upload 锁.
var (
	chunkMu      sync.RWMutex
	chunkMetaMap = map[string]*chunkMeta{}
)

func chunkDirPath(id string) string { return filepath.Join(chunkDir, id) }
func chunkFilePath(id string, idx int) string {
	return filepath.Join(chunkDir, id, strconv.Itoa(idx))
}

func metaPath(id string) string { return filepath.Join(chunkDir, id, "meta.json") }

// saveMeta 原子写 meta.json (.tmp + rename).
func saveMeta(m *chunkMeta) error {
	b, err := json.Marshal(m)
	if err != nil {
		return err
	}
	tmp := metaPath(m.UploadID) + ".tmp"
	if err := os.WriteFile(tmp, b, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, metaPath(m.UploadID))
}

func loadMeta(id string) *chunkMeta {
	b, err := os.ReadFile(metaPath(id))
	if err != nil {
		return nil
	}
	var m chunkMeta
	if err := json.Unmarshal(b, &m); err != nil {
		return nil
	}
	return &m
}

// expectedChunks totalSize 按 chunkSize 分片应有的片数
func expectedChunks(totalSize int64) int {
	return int((totalSize + chunkSize - 1) / chunkSize)
}

// expectedChunkSize 第 i 片应有的字节数 (最后一片是余数)
func expectedChunkSize(totalSize int64, i int) int64 {
	if n := totalSize - int64(i)*chunkSize; n < chunkSize {
		return n
	}
	return chunkSize
}

func uploadChunkHandler(c *gin.Context) {
	uploadID := c.GetHeader("X-Upload-Id")
	if !uploadIDRe.MatchString(uploadID) {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid upload id"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_upload_id"})
		return
	}
	idx, err := strconv.Atoi(c.GetHeader("X-Chunk-Index"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid chunk index"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_chunk_index"})
		return
	}
	totalChunks, err := strconv.Atoi(c.GetHeader("X-Total-Chunks"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid total chunks"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_total_chunks"})
		return
	}
	filename := c.GetHeader("X-Filename")
	if !checkFilenameSafe(filename) {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid filename"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_filename"})
		return
	}
	totalSize, err := strconv.ParseInt(c.GetHeader("X-Total-Size"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid total size"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_total_size"})
		return
	}
	if totalSize <= chunkThreshold {
		c.JSON(http.StatusBadRequest, gin.H{"error": "chunked upload is for files > 5GB"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "below_threshold"})
		return
	}
	if totalChunks < 1 || idx < 0 || idx >= totalChunks {
		c.JSON(http.StatusBadRequest, gin.H{"error": "chunk index out of range"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "total": totalChunks, "result": "index_out_of_range"})
		return
	}
	if e := expectedChunks(totalSize); totalChunks < e-1 || totalChunks > e+1 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "total chunks mismatch"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "total": totalChunks, "expected": e, "result": "chunk_count_mismatch"})
		return
	}
	if totalSize > quotaRemaining() {
		c.JSON(http.StatusInsufficientStorage, gin.H{"error": "quota exceeded"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "size": totalSize, "result": "quota_exceeded"})
		return
	}

	chunk, err := c.FormFile("chunk")
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "missing form field 'chunk'"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "missing_form_field"})
		return
	}

	dir := chunkDirPath(uploadID)
	chunkMu.RLock()
	meta, inMap := chunkMetaMap[uploadID]
	chunkMu.RUnlock()
	if !inMap {
		meta = loadMeta(uploadID) // 进程重启后从磁盘恢复
		if meta != nil {
			chunkMu.Lock()
			chunkMetaMap[uploadID] = meta
			chunkMu.Unlock()
		}
	}
	if meta == nil {
		if idx != 0 {
			c.JSON(http.StatusNotFound, gin.H{"error": "upload not found"})
			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "upload_not_found"})
			return
		}
		if err := os.MkdirAll(dir, 0o755); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "mkdir_failed"})
			return
		}
		if !quotaReserve(totalSize) {
			_ = os.RemoveAll(dir)
			c.JSON(http.StatusInsufficientStorage, gin.H{"error": "quota exceeded"})
			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "size": totalSize, "result": "quota_exceeded"})
			return
		}
		meta = &chunkMeta{
			UploadID: uploadID, Filename: filename, TotalSize: totalSize,
			TotalChunks: totalChunks, Received: map[int]bool{},
			CreatedAt: time.Now(), Reserved: totalSize,
		}
		chunkMu.Lock()
		chunkMetaMap[uploadID] = meta
		err = saveMeta(meta)
		chunkMu.Unlock()
		if err != nil {
			abortChunkUpload(uploadID, meta)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "meta_save_failed"})
			return
		}
	} else if idx > 0 {
		if _, err := os.Stat(dir); err != nil {
			c.JSON(http.StatusNotFound, gin.H{"error": "upload not found"})
			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "upload_not_found"})
			return
		}
	}

	// 写分片: multipart → 流式 io.Copy 落盘, 不进内存
	dst := chunkFilePath(uploadID, idx)
	in, err := chunk.Open()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "open_failed"})
		return
	}
	out, err := os.Create(dst)
	if err != nil {
		in.Close()
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "create_failed"})
		return
	}
	_, cpErr := io.Copy(out, in)
	in.Close()
	if cerr := out.Close(); cerr != nil && cpErr == nil {
		cpErr = cerr
	}
	if cpErr != nil {
		_ = os.Remove(dst)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "write_failed"})
		return
	}

	// 更新 Received + 持久化 (持写锁, 防并发分片时 meta.json 丢进度)
	chunkMu.Lock()
	meta.Received[idx] = true
	err = saveMeta(meta)
	complete := len(meta.Received) == meta.TotalChunks
	chunkMu.Unlock()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "meta_save_failed"})
		return
	}

	if !complete {
		c.JSON(http.StatusAccepted, gin.H{"received": len(meta.Received), "total": meta.TotalChunks})
		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "received": len(meta.Received), "total": meta.TotalChunks, "result": "ok"})
		return
	}
	assembleChunks(c, uploadID, meta)
}

// assembleChunks 所有分片到位后顺序拼装落盘并清理.
// ponytail: 末片重传/并发组装不额外去重, 前端单分片串行重试不会触发.
func assembleChunks(c *gin.Context, uploadID string, meta *chunkMeta) {
	dir := chunkDirPath(uploadID)
	for i := 0; i < meta.TotalChunks; i++ {
		info, err := os.Stat(chunkFilePath(uploadID, i))
		if err != nil || info.Size() != expectedChunkSize(meta.TotalSize, i) {
			abortChunkUpload(uploadID, meta)
			c.JSON(http.StatusBadRequest, gin.H{"error": "chunk size mismatch"})
			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "chunk_size_mismatch"})
			return
		}
	}

	name := uniqueFilename(*uploadDir, meta.Filename)
	dst := filepath.Join(*uploadDir, name)
	out, err := os.Create(dst)
	if err != nil {
		abortChunkUpload(uploadID, meta)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "result": "create_failed"})
		return
	}
	var realSize int64
	for i := 0; i < meta.TotalChunks; i++ {
		in, err := os.Open(chunkFilePath(uploadID, i))
		if err != nil {
			out.Close()
			_ = os.Remove(dst)
			abortChunkUpload(uploadID, meta)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "open_failed"})
			return
		}
		n, cErr := io.Copy(out, in)
		in.Close()
		if cErr != nil {
			out.Close()
			_ = os.Remove(dst)
			abortChunkUpload(uploadID, meta)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "copy_failed"})
			return
		}
		realSize += n
	}
	if err := out.Close(); err != nil {
		_ = os.Remove(dst)
		abortChunkUpload(uploadID, meta)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
		audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "result": "close_failed"})
		return
	}

	// 配额: 释放预留, 按真实大小重新占用
	quotaSub(meta.Reserved)
	quotaReserve(realSize)
	chunkMu.Lock()
	delete(chunkMetaMap, uploadID)
	chunkMu.Unlock()
	_ = os.RemoveAll(dir)

	now := time.Now()
	_ = os.Chtimes(dst, now, now)
	c.JSON(http.StatusOK, fileResponse{
		Filename: name, URL: "/download/" + name, Size: realSize,
		UploadedAt: now, ExpiresAt: now.Add(*fileTTL),
	})
	audit(c, "upload_assembled", gin.H{
		"upload_id": uploadID, "file": name, "size": realSize, "result": "ok",
	})
}

// abortChunkUpload 失败时释放预留配额并清掉分片目录与内存记录
func abortChunkUpload(uploadID string, meta *chunkMeta) {
	quotaSub(meta.Reserved)
	chunkMu.Lock()
	delete(chunkMetaMap, uploadID)
	chunkMu.Unlock()
	_ = os.RemoveAll(chunkDirPath(uploadID))
}

// cleanupChunks 清理超过 maxAge 的孤儿分片目录并释放预留配额
func cleanupChunks(maxAge time.Duration) {
	entries, err := os.ReadDir(chunkDir)
	if err != nil {
		if !os.IsNotExist(err) {
			log.Printf("[cleaner] read chunk dir failed: %v", err)
		}
		return
	}
	now := time.Now()
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		id := e.Name()
		dir := chunkDirPath(id)
		if !uploadIDRe.MatchString(id) {
			_ = os.RemoveAll(dir) // 非法目录名 (非本服务生成), 直接清
			continue
		}
		meta := loadMeta(id)
		if meta == nil {
			// 无 meta 的孤儿目录: 无法确定 reserved, 删目录即可 (quota 由 cleaner 的 rescanQuota 校正)
			_ = os.RemoveAll(dir)
			chunkMu.Lock()
			delete(chunkMetaMap, id)
			chunkMu.Unlock()
			continue
		}
		if now.Sub(meta.CreatedAt) > maxAge {
			quotaSub(meta.Reserved)
			chunkMu.Lock()
			delete(chunkMetaMap, id)
			chunkMu.Unlock()
			_ = os.RemoveAll(dir)
			log.Printf("[cleaner] removed orphan chunks: %s (age=%v)", id, now.Sub(meta.CreatedAt))
		}
	}
}