/*
Copyright 2026 linux.do
Modified by Arctel.net, 2026

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package upload

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/Rain-kl/Wavelet/internal/apps/oauth"
	"github.com/Rain-kl/Wavelet/internal/db"
	"github.com/Rain-kl/Wavelet/internal/model"
	"github.com/Rain-kl/Wavelet/internal/storage"
	"github.com/Rain-kl/Wavelet/internal/testhelper"
	"github.com/Rain-kl/Wavelet/internal/util"
	"github.com/gin-gonic/gin"
)

type testResponse struct {
	ErrorMsg string          `json:"error_msg"`
	Data     json.RawMessage `json:"data"`
}

func setupTestRouter(authUser *model.User) *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	uploadGroup := r.Group("/api/v1/upload")

	// Mock authentication middleware
	uploadGroup.Use(func(c *gin.Context) {
		if authUser != nil {
			util.SetToContext(c, oauth.UserObjKey, authUser)
		}
		c.Next()
	})

	uploadGroup.POST("", UploadFile)
	uploadGroup.GET("/download/:id", DownloadFile)
	uploadGroup.POST("/download/batch", BatchDownloadFiles)
	return r
}

func createMultipartRequest(t *testing.T, fieldName, fileName string, fileContent []byte, extraFields map[string]string) (string, *bytes.Buffer) {
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	part, err := writer.CreateFormFile(fieldName, fileName)
	if err != nil {
		t.Fatalf("failed to create form file: %v", err)
	}

	_, err = part.Write(fileContent)
	if err != nil {
		t.Fatalf("failed to write file content: %v", err)
	}

	for k, v := range extraFields {
		err = writer.WriteField(k, v)
		if err != nil {
			t.Fatalf("failed to write form field: %v", err)
		}
	}

	err = writer.Close()
	if err != nil {
		t.Fatalf("failed to close multipart writer: %v", err)
	}

	return writer.FormDataContentType(), body
}

func TestUploadFile(t *testing.T) {
	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
	defer cleanup()
	defer os.RemoveAll("uploads") // Clean up local files created during tests

	authUser := &model.User{ID: 1001, Username: "test_user"}
	router := setupTestRouter(authUser)

	// Mock Storage Client
	mockFiles := make(map[string][]byte)
	var putCount int

	restoreStorage := storage.MockStorage(
		func(ctx context.Context, key string, body io.Reader, size int64, contentType string) error {
			data, err := io.ReadAll(body)
			if err != nil {
				return err
			}
			mockFiles[key] = data
			putCount++
			return nil
		},
		func(ctx context.Context, key string) (*storage.ObjectInfo, error) {
			data, ok := mockFiles[key]
			if !ok {
				return nil, os.ErrNotExist
			}
			return &storage.ObjectInfo{
				Body:          io.NopCloser(bytes.NewReader(data)),
				ContentLength: int64(len(data)),
				ContentType:   "application/octet-stream",
			}, nil
		},
		func(ctx context.Context, key string) error {
			delete(mockFiles, key)
			return nil
		},
	)
	defer restoreStorage()

	// 开启 S3 Storage
	storage.IsEnabledFunc = func() bool { return true }
	defer func() {
		storage.IsEnabledFunc = func() bool { return false }
	}()

	t.Run("upload allowed image file successfully", func(t *testing.T) {
		putCount = 0
		imgContent := []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89") // Valid PNG header
		contentType, body := createMultipartRequest(t, "file", "test.png", imgContent, map[string]string{
			"type":     "avatar",
			"metadata": `{"extra":{"source":"test_runner"}}`,
		})

		req, _ := http.NewRequest("POST", "/api/v1/upload", body)
		req.Header.Set("Content-Type", contentType)

		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String())
		}

		var resp testResponse
		if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
			t.Fatalf("failed to unmarshal response: %v", err)
		}

		if resp.ErrorMsg != "" {
			t.Fatalf("expected success response, got failure: %s", resp.ErrorMsg)
		}

		// Verify database record
		var uploadRecord model.Upload
		if err := json.Unmarshal(resp.Data, &uploadRecord); err != nil {
			t.Fatalf("failed to unmarshal upload record: %v", err)
		}

		var dbRecord model.Upload
		if err := dbConn.First(&dbRecord, uploadRecord.ID).Error; err != nil {
			t.Fatalf("failed to retrieve database record: %v", err)
		}

		if dbRecord.FileName != "test.png" || dbRecord.Extension != "png" {
			t.Errorf("incorrect filename or extension: %s, %s", dbRecord.FileName, dbRecord.Extension)
		}

		if dbRecord.MimeType != "image/png" {
			t.Errorf("incorrect mime type detected: %s", dbRecord.MimeType)
		}

		if dbRecord.StorageDriver != "s3" {
			t.Errorf("expected storage driver s3, got %s", dbRecord.StorageDriver)
		}

		if dbRecord.Metadata.Extra["source"] != "test_runner" {
			t.Errorf("expected extra meta 'source' to be 'test_runner', got %v", dbRecord.Metadata.Extra)
		}

		if putCount != 1 {
			t.Errorf("expected 1 storage Put operation, got %d", putCount)
		}
	})

	t.Run("upload blocked extension file", func(t *testing.T) {
		// System config allowed: jpg,png,webp. Uploading docx should be blocked.
		contentType, body := createMultipartRequest(t, "file", "contract.docx", []byte("fake docx content"), nil)
		req, _ := http.NewRequest("POST", "/api/v1/upload", body)
		req.Header.Set("Content-Type", contentType)

		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String())
		}

		var resp testResponse
		json.Unmarshal(w.Body.Bytes(), &resp)
		if resp.ErrorMsg == "" || !strings.Contains(resp.ErrorMsg, ErrUnsupportedFormat) {
			t.Errorf("expected unsupported format error, got: %v", resp)
		}
	})

	t.Run("instant upload deduplication (秒传)", func(t *testing.T) {
		putCount = 0
		imgContent := []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01")

		// Upload first time
		contentType1, body1 := createMultipartRequest(t, "file", "avatar1.png", imgContent, map[string]string{"type": "avatar"})
		req1, _ := http.NewRequest("POST", "/api/v1/upload", body1)
		req1.Header.Set("Content-Type", contentType1)
		w1 := httptest.NewRecorder()
		router.ServeHTTP(w1, req1)

		if w1.Code != http.StatusOK {
			t.Fatalf("first upload failed: %s", w1.Body.String())
		}
		if putCount != 1 {
			t.Errorf("expected 1 put count on first upload, got %d", putCount)
		}

		// Upload same file second time (different filename, same content)
		contentType2, body2 := createMultipartRequest(t, "file", "avatar2.png", imgContent, map[string]string{"type": "avatar"})
		req2, _ := http.NewRequest("POST", "/api/v1/upload", body2)
		req2.Header.Set("Content-Type", contentType2)
		w2 := httptest.NewRecorder()
		router.ServeHTTP(w2, req2)

		if w2.Code != http.StatusOK {
			t.Fatalf("second upload failed: %s", w2.Body.String())
		}

		var resp2 testResponse
		json.Unmarshal(w2.Body.Bytes(), &resp2)

		if resp2.ErrorMsg != "" {
			t.Fatalf("second upload was unsuccessful: %s", resp2.ErrorMsg)
		}

		var uploadRecord2 model.Upload
		if err := json.Unmarshal(resp2.Data, &uploadRecord2); err != nil {
			t.Fatalf("failed to unmarshal second upload record: %v", err)
		}

		// Check if it triggered another storage put
		if putCount != 1 {
			t.Errorf("PutObject was triggered again! Expected deduplication (putCount=1), got putCount=%d", putCount)
		}

		// Check if database contains both records sharing the same FilePath
		var records []model.Upload
		dbConn.Where("hash = ?", uploadRecord2.Hash).Find(&records)
		if len(records) != 2 {
			t.Errorf("expected 2 database records sharing the same hash, got %d", len(records))
		}
		if records[0].FilePath != records[1].FilePath {
			t.Errorf("file paths are different: %s vs %s", records[0].FilePath, records[1].FilePath)
		}
		if records[0].ID == records[1].ID {
			t.Error("database record IDs should be unique")
		}

		t.Logf("Instant upload success. Record 1: %d, Record 2: %d", records[0].ID, records[1].ID)
	})

	t.Run("upload in local storage fallback mode", func(t *testing.T) {
		// Turn off S3
		storage.IsEnabledFunc = func() bool { return false }

		// Seed allowed extensions configuration to allow txt files
		var sc model.SystemConfig
		dbConn.Where("key = ?", model.ConfigKeyUploadAllowedExtensions).First(&sc)
		sc.Value = "jpg,png,webp,txt"
		dbConn.Save(&sc)
		_ = db.HSetJSON(context.Background(), model.SystemConfigRedisHashKey, sc.Key, &sc)

		contentType, body := createMultipartRequest(t, "file", "doc.txt", []byte("hello world generic document file"), map[string]string{
			"type": "document",
		})
		req, _ := http.NewRequest("POST", "/api/v1/upload", body)
		req.Header.Set("Content-Type", contentType)

		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String())
		}

		var resp testResponse
		json.Unmarshal(w.Body.Bytes(), &resp)

		if resp.ErrorMsg != "" {
			t.Fatalf("local upload failed: %s", resp.ErrorMsg)
		}

		var localRecord model.Upload
		if err := json.Unmarshal(resp.Data, &localRecord); err != nil {
			t.Fatalf("failed to unmarshal local upload record: %v", err)
		}

		if localRecord.StorageDriver != "local" {
			t.Errorf("expected storage driver local, got %s", localRecord.StorageDriver)
		}

		// Confirm file was actually written to local disk
		fileContent, err := os.ReadFile(localRecord.FilePath)
		if err != nil {
			t.Fatalf("failed to read local file: %v", err)
		}

		if string(fileContent) != "hello world generic document file" {
			t.Errorf("unexpected local file contents: %s", string(fileContent))
		}
	})
}

func TestDownloadFile(t *testing.T) {
	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
	defer cleanup()
	defer os.RemoveAll("uploads")

	authUser := &model.User{ID: 1001, Username: "test_user"}
	router := setupTestRouter(authUser)

	// Seed upload records in DB
	localUpload := model.Upload{
		ID:            2001,
		UserID:        1001,
		FileName:      "中文文件名.txt",
		FilePath:      "uploads/test_download.txt",
		FileSize:      12,
		MimeType:      "text/plain",
		Extension:     "txt",
		StorageDriver: "local",
		Status:        model.UploadStatusUsed,
	}

	// Create local file
	err := os.MkdirAll("uploads", 0755)
	if err != nil {
		t.Fatalf("failed to create directory: %v", err)
	}
	err = os.WriteFile(localUpload.FilePath, []byte("hello download"), 0644)
	if err != nil {
		t.Fatalf("failed to write file: %v", err)
	}

	dbConn.Create(&localUpload)

	t.Run("download file successfully", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/api/v1/upload/download/2001", nil)
		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String())
		}

		if w.Body.String() != "hello download" {
			t.Errorf("expected body 'hello download', got '%s'", w.Body.String())
		}

		// Verify Content-Disposition header (supports UTF-8 escaping)
		contentDisp := w.Header().Get("Content-Disposition")
		expectedDisp := "attachment; filename*=UTF-8''%E4%B8%AD%E6%96%87%E6%96%87%E4%BB%B6%E5%90%8D.txt"
		if contentDisp != expectedDisp {
			t.Errorf("expected Content-Disposition header %q, got %q", expectedDisp, contentDisp)
		}

		if w.Header().Get("Content-Type") != "text/plain" {
			t.Errorf("expected Content-Type text/plain, got %s", w.Header().Get("Content-Type"))
		}
	})

	t.Run("download non-existent file", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/api/v1/upload/download/9999", nil)
		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusNotFound {
			t.Errorf("expected status 404, got %d", w.Code)
		}
	})
}

func TestBatchDownloadFiles(t *testing.T) {
	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
	defer cleanup()
	defer os.RemoveAll("uploads")

	authUser := &model.User{ID: 1001, Username: "test_user"}
	router := setupTestRouter(authUser)

	// Create and write files locally
	err := os.MkdirAll("uploads", 0755)
	if err != nil {
		t.Fatalf("failed to create local dir: %v", err)
	}

	_ = os.WriteFile("uploads/f1.txt", []byte("file1 content"), 0644)
	_ = os.WriteFile("uploads/f2.txt", []byte("file2 content"), 0644)
	_ = os.WriteFile("uploads/f3.txt", []byte("duplicate name file content"), 0644)

	// Seed upload records. Note f2 and f3 have the same FileName "file_a.txt" to trigger name collision resolution.
	uploads := []model.Upload{
		{
			ID:            3001,
			UserID:        1001,
			FileName:      "file_a.txt",
			FilePath:      "uploads/f1.txt",
			FileSize:      13,
			MimeType:      "text/plain",
			Extension:     "txt",
			StorageDriver: "local",
			Status:        model.UploadStatusUsed,
		},
		{
			ID:            3002,
			UserID:        1001,
			FileName:      "file_b.txt",
			FilePath:      "uploads/f2.txt",
			FileSize:      13,
			MimeType:      "text/plain",
			Extension:     "txt",
			StorageDriver: "local",
			Status:        model.UploadStatusUsed,
		},
		{
			ID:            3003,
			UserID:        1001,
			FileName:      "file_a.txt", // COLLISION with 3001!
			FilePath:      "uploads/f3.txt",
			FileSize:      28,
			MimeType:      "text/plain",
			Extension:     "txt",
			StorageDriver: "local",
			Status:        model.UploadStatusUsed,
		},
	}

	for _, up := range uploads {
		dbConn.Create(&up)
	}

	t.Run("batch download zip successfully and check duplicate renaming", func(t *testing.T) {
		reqBody, _ := json.Marshal(batchDownloadRequest{
			IDs: []string{"3001", "3002", "3003"},
		})
		req, _ := http.NewRequest("POST", "/api/v1/upload/download/batch", bytes.NewReader(reqBody))
		req.Header.Set("Content-Type", "application/json")

		w := httptest.NewRecorder()
		router.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String())
		}

		if w.Header().Get("Content-Type") != "application/zip" {
			t.Errorf("expected Content-Type application/zip, got %s", w.Header().Get("Content-Type"))
		}

		// Unzip in-memory
		zipReader, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
		if err != nil {
			t.Fatalf("failed to read zip buffer: %v", err)
		}

		if len(zipReader.File) != 3 {
			t.Errorf("expected 3 files inside the ZIP, got %d", len(zipReader.File))
		}

		// Extract files to check their contents and name collision resolutions
		extracted := make(map[string]string)
		for _, f := range zipReader.File {
			rc, err := f.Open()
			if err != nil {
				t.Fatalf("failed to open zip file entry %s: %v", f.Name, err)
			}
			content, _ := io.ReadAll(rc)
			rc.Close()
			extracted[f.Name] = string(content)
		}

		// Checks
		if extracted["file_a.txt"] != "file1 content" {
			t.Errorf("file_a.txt content incorrect: %q", extracted["file_a.txt"])
		}
		if extracted["file_b.txt"] != "file2 content" {
			t.Errorf("file_b.txt content incorrect: %q", extracted["file_b.txt"])
		}
		// The second file_a.txt should be renamed to file_a_1.txt
		if extracted["file_a_1.txt"] != "duplicate name file content" {
			t.Errorf("file_a_1.txt content incorrect: %q. Extracted files: %v", extracted["file_a_1.txt"], extracted)
		}

		t.Logf("Successfully unzipped batch. Extracted files: %+v", extracted)
	})
}