package fileserver

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"image"
	"image/color"
	"image/png"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/labstack/echo/v5"
	"github.com/stretchr/testify/require"
	"google.golang.org/protobuf/types/known/timestamppb"

	"github.com/usememos/memos/internal/markdown"
	"github.com/usememos/memos/internal/profile"
	"github.com/usememos/memos/internal/testutil"
	"github.com/usememos/memos/internal/util"
	apiv1 "github.com/usememos/memos/proto/gen/api/v1"
	storepb "github.com/usememos/memos/proto/gen/store"
	"github.com/usememos/memos/server/auth"
	apiv1service "github.com/usememos/memos/server/router/api/v1"
	"github.com/usememos/memos/store"
	teststore "github.com/usememos/memos/store/test"
)

func TestServeAttachmentFile_ShareTokenAllowsDirectMemoAttachment(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "share-parent-owner",
		Role:     store.RoleUser,
		Email:    "share-parent-owner@example.com",
	})
	require.NoError(t, err)

	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "memo.txt",
			Type:     "text/plain",
			Content:  []byte("memo attachment"),
		},
	})
	require.NoError(t, err)

	parentMemo, err := svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "shared parent",
			Visibility: apiv1.Visibility_PROTECTED,
			Attachments: []*apiv1.Attachment{
				{Name: attachment.Name},
			},
		},
	})
	require.NoError(t, err)

	share, err := svc.CreateMemoShare(creatorCtx, &apiv1.CreateMemoShareRequest{
		Parent:    parentMemo.Name,
		MemoShare: &apiv1.MemoShare{},
	})
	require.NoError(t, err)
	shareToken := share.Name[strings.LastIndex(share.Name, "/")+1:]

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s?share_token=%s", attachment.Name, attachment.Filename, shareToken), nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "memo attachment", rec.Body.String())
}

func TestServeAttachmentFile_LocalStaticFileSupportsRangeRequests(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "range-owner",
		Role:     store.RoleUser,
		Email:    "range-owner@example.com",
	})
	require.NoError(t, err)
	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "range.txt",
			Type:     "text/plain",
			Content:  []byte("0123456789"),
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "range memo",
			Visibility: apiv1.Visibility_PUBLIC,
			Attachments: []*apiv1.Attachment{
				{Name: attachment.Name},
			},
		},
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s", attachment.Name, attachment.Filename), nil)
	req.Header.Set("Range", "bytes=2-5")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusPartialContent, rec.Code)
	require.Equal(t, "2345", rec.Body.String())
	require.Equal(t, "bytes 2-5/10", rec.Header().Get("Content-Range"))
}

func TestServeAttachmentFile_ShareTokenRejectsCommentAttachment(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "private-parent-owner",
		Role:     store.RoleUser,
		Email:    "private-parent-owner@example.com",
	})
	require.NoError(t, err)

	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)
	commenter, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "share-commenter",
		Role:     store.RoleUser,
		Email:    "share-commenter@example.com",
	})
	require.NoError(t, err)
	commenterCtx := context.WithValue(ctx, auth.UserIDContextKey, commenter.ID)

	parentMemo, err := svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "shared parent",
			Visibility: apiv1.Visibility_PROTECTED,
		},
	})
	require.NoError(t, err)

	commentAttachment, err := svc.CreateAttachment(commenterCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "comment.txt",
			Type:     "text/plain",
			Content:  []byte("comment attachment"),
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemoComment(commenterCtx, &apiv1.CreateMemoCommentRequest{
		Name: parentMemo.Name,
		Comment: &apiv1.Memo{
			Content:    "comment with attachment",
			Visibility: apiv1.Visibility_PROTECTED,
			Attachments: []*apiv1.Attachment{
				{Name: commentAttachment.Name},
			},
		},
	})
	require.NoError(t, err)

	share, err := svc.CreateMemoShare(creatorCtx, &apiv1.CreateMemoShareRequest{
		Parent:    parentMemo.Name,
		MemoShare: &apiv1.MemoShare{},
	})
	require.NoError(t, err)
	shareToken := share.Name[strings.LastIndex(share.Name, "/")+1:]

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s?share_token=%s", commentAttachment.Name, commentAttachment.Filename, shareToken), nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusUnauthorized, rec.Code)
}

func TestServeAttachmentFile_MotionClip(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "motion-owner",
		Role:     store.RoleUser,
		Email:    "motion-owner@example.com",
	})
	require.NoError(t, err)
	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "motion.jpg",
			Type:     "image/jpeg",
			Content:  testutil.BuildMotionPhotoJPEG(),
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "motion memo",
			Visibility: apiv1.Visibility_PUBLIC,
			Attachments: []*apiv1.Attachment{
				{Name: attachment.Name},
			},
		},
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s?motion=true", attachment.Name, attachment.Filename), nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "video/mp4", rec.Header().Get("Content-Type"))
	require.Contains(t, rec.Body.String(), "ftyp")
}

func TestServeAttachmentFile_SVGThumbnailServedAsImageWithSecurityHeaders(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "svg-owner",
		Role:     store.RoleUser,
		Email:    "svg-owner@example.com",
	})
	require.NoError(t, err)
	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	svgContent := []byte(`<svg xmlns="http://www.w3.org/2000/svg" width="120" height="40"><text x="0" y="20">memos</text></svg>`)
	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "preview.svg",
			Type:     "image/svg+xml",
			Content:  svgContent,
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "svg memo",
			Visibility: apiv1.Visibility_PUBLIC,
			Attachments: []*apiv1.Attachment{
				{Name: attachment.Name},
			},
		},
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s?thumbnail=true", attachment.Name, attachment.Filename), nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "image/svg+xml", rec.Header().Get("Content-Type"))
	require.Empty(t, rec.Header().Get("Content-Disposition"))
	require.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options"))
	require.Equal(t, "default-src 'none'; style-src 'unsafe-inline';", rec.Header().Get("Content-Security-Policy"))
	require.Equal(t, svgContent, rec.Body.Bytes())
}

func TestServeAttachmentFile_ThumbnailWithSensitiveMetadataServesOriginal(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "hdr-owner",
		Role:     store.RoleUser,
		Email:    "hdr-owner@example.com",
	})
	require.NoError(t, err)
	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	imageContent := testPNGWithChunk(t, "cICP", []byte{9, 16, 9, 1})
	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "hdr.png",
			Type:     "image/png",
			Content:  imageContent,
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:    "hdr memo",
			Visibility: apiv1.Visibility_PUBLIC,
			Attachments: []*apiv1.Attachment{
				{Name: attachment.Name},
			},
		},
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)

	req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/file/%s/%s?thumbnail=true", attachment.Name, attachment.Filename), nil)
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "image/png", rec.Header().Get("Content-Type"))
	require.Equal(t, imageContent, rec.Body.Bytes())
}

func TestHasThumbnailSensitiveMetadata(t *testing.T) {
	tests := []struct {
		name     string
		mimeType string
		data     []byte
		want     bool
	}{
		{
			name:     "jpeg hdr gain map",
			mimeType: "image/jpeg",
			data:     []byte("xmp hdrgm:Version=\"1.0\""),
			want:     true,
		},
		{
			name:     "jpeg icc profile",
			mimeType: "image/jpeg",
			data:     []byte("ICC_PROFILE"),
			want:     true,
		},
		{
			name:     "png cicp chunk",
			mimeType: "image/png",
			data:     []byte("cICP"),
			want:     true,
		},
		{
			name:     "heic",
			mimeType: "image/heic",
			data:     nil,
			want:     true,
		},
		{
			name:     "plain jpeg",
			mimeType: "image/jpeg",
			data:     []byte("plain image data"),
			want:     false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, hasThumbnailSensitiveMetadata(tt.mimeType, tt.data))
		})
	}
}

func testPNGWithChunk(t *testing.T, chunkType string, chunkData []byte) []byte {
	t.Helper()

	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
	img.Set(0, 0, color.RGBA{R: 255, A: 255})

	var encoded bytes.Buffer
	require.NoError(t, png.Encode(&encoded, img))

	pngData := encoded.Bytes()
	iendIndex := bytes.LastIndex(pngData, []byte("IEND"))
	require.GreaterOrEqual(t, iendIndex, 4)

	chunkStart := iendIndex - 4
	var chunk bytes.Buffer
	require.NoError(t, binary.Write(&chunk, binary.BigEndian, uint32(len(chunkData))))
	chunk.WriteString(chunkType)
	chunk.Write(chunkData)
	checksum := crc32.ChecksumIEEE(append([]byte(chunkType), chunkData...))
	require.NoError(t, binary.Write(&chunk, binary.BigEndian, checksum))

	result := make([]byte, 0, len(pngData)+chunk.Len())
	result = append(result, pngData[:chunkStart]...)
	result = append(result, chunk.Bytes()...)
	result = append(result, pngData[chunkStart:]...)
	return result
}

func newShareAttachmentTestServices(ctx context.Context, t *testing.T) (*apiv1service.APIV1Service, *FileServerService, *store.Store, func()) {
	t.Helper()

	testStore := teststore.NewTestingStore(ctx, t)
	testProfile := &profile.Profile{
		Demo:        true,
		Version:     "test-1.0.0",
		InstanceURL: "http://localhost:8080",
		Driver:      "sqlite",
		DSN:         ":memory:",
		Data:        t.TempDir(),
	}
	secret := "test-secret"
	markdownService := markdown.NewService(markdown.WithTagExtension())
	apiService := &apiv1service.APIV1Service{
		Secret:          secret,
		Profile:         testProfile,
		Store:           testStore,
		MarkdownService: markdownService,
		SSEHub:          apiv1service.NewSSEHub(),
	}
	fileService := NewFileServerService(testProfile, testStore, secret)

	return apiService, fileService, testStore, func() {
		testStore.Close()
	}
}

// makePNGDataURI returns a minimal valid PNG encoded as a data URI, suitable for a
// user avatar.
func makePNGDataURI(t *testing.T) string {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
	img.Set(0, 0, color.RGBA{R: 1, G: 2, B: 3, A: 255})
	var buf bytes.Buffer
	require.NoError(t, png.Encode(&buf, img))
	return "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())
}

// TestServeAttachmentFile_PrivateInstanceDeniesAnonymous verifies that a public
// memo's attachment is served to anonymous visitors on an open instance but denied
// on a private instance (no InstanceURL configured).
func TestServeAttachmentFile_PrivateInstanceDeniesAnonymous(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	creator, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "private-attachment-owner",
		Role:     store.RoleUser,
		Email:    "private-attachment-owner@example.com",
	})
	require.NoError(t, err)
	creatorCtx := context.WithValue(ctx, auth.UserIDContextKey, creator.ID)

	attachment, err := svc.CreateAttachment(creatorCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "public.txt",
			Type:     "text/plain",
			Content:  []byte("public content"),
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(creatorCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:     "public memo",
			Visibility:  apiv1.Visibility_PUBLIC,
			Attachments: []*apiv1.Attachment{{Name: attachment.Name}},
		},
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)
	url := fmt.Sprintf("/file/%s/%s", attachment.Name, attachment.Filename)

	anonymousGet := func() int {
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, url, nil))
		return rec.Code
	}

	// Open instance: anonymous access to a public memo's attachment is allowed.
	fs.Profile.InstanceURL = "http://localhost:8080"
	require.Equal(t, http.StatusOK, anonymousGet())

	// Private instance: the same anonymous request is denied.
	fs.Profile.InstanceURL = ""
	require.Equal(t, http.StatusUnauthorized, anonymousGet())
}

// TestServeUserAvatar_PrivateInstanceRequiresAuth verifies that avatars are exposed
// to anonymous visitors on an open instance but require authentication on a private
// instance.
func TestServeUserAvatar_PrivateInstanceRequiresAuth(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	_, err := svc.Store.CreateUser(ctx, &store.User{
		Username:  "avatar-owner",
		Role:      store.RoleUser,
		Email:     "avatar-owner@example.com",
		AvatarURL: makePNGDataURI(t),
	})
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)

	anonymousGet := func() int {
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/file/users/avatar-owner/avatar", nil))
		return rec.Code
	}

	// Open instance: anonymous avatar access is allowed.
	fs.Profile.InstanceURL = "http://localhost:8080"
	require.Equal(t, http.StatusOK, anonymousGet())

	// Private instance: anonymous avatar access is denied.
	fs.Profile.InstanceURL = ""
	require.Equal(t, http.StatusUnauthorized, anonymousGet())
}

// TestServeAttachmentFile_RefreshCookieAuthenticatesOwner verifies that the file
// server authenticates a request via the refresh-token cookie (the browser <img>
// flow) — the AuthenticateToUser cookie fallback — letting the owner fetch their own
// private memo's attachment without an Authorization header.
func TestServeAttachmentFile_RefreshCookieAuthenticatesOwner(t *testing.T) {
	ctx := context.Background()
	svc, fs, _, cleanup := newShareAttachmentTestServices(ctx, t)
	defer cleanup()

	owner, err := svc.Store.CreateUser(ctx, &store.User{
		Username: "cookie-owner",
		Role:     store.RoleUser,
		Email:    "cookie-owner@example.com",
	})
	require.NoError(t, err)
	ownerCtx := context.WithValue(ctx, auth.UserIDContextKey, owner.ID)

	attachment, err := svc.CreateAttachment(ownerCtx, &apiv1.CreateAttachmentRequest{
		Attachment: &apiv1.Attachment{
			Filename: "secret.txt",
			Type:     "text/plain",
			Content:  []byte("secret content"),
		},
	})
	require.NoError(t, err)

	_, err = svc.CreateMemo(ownerCtx, &apiv1.CreateMemoRequest{
		Memo: &apiv1.Memo{
			Content:     "private memo",
			Visibility:  apiv1.Visibility_PRIVATE,
			Attachments: []*apiv1.Attachment{{Name: attachment.Name}},
		},
	})
	require.NoError(t, err)

	// Mint a valid refresh token for the owner and store its record.
	tokenID := util.GenUUID()
	require.NoError(t, svc.Store.AddUserRefreshToken(ctx, owner.ID, &storepb.RefreshTokensUserSetting_RefreshToken{
		TokenId:   tokenID,
		ExpiresAt: timestamppb.New(time.Now().Add(auth.RefreshTokenDuration)),
		CreatedAt: timestamppb.Now(),
	}))
	refreshToken, _, err := auth.GenerateRefreshToken(owner.ID, tokenID, []byte(svc.Secret))
	require.NoError(t, err)

	e := echo.New()
	fs.RegisterRoutes(e)
	url := fmt.Sprintf("/file/%s/%s", attachment.Name, attachment.Filename)

	// Without credentials, the private attachment is denied.
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, url, nil))
	require.Equal(t, http.StatusUnauthorized, rec.Code)

	// With the refresh-token cookie, the owner is authenticated and served.
	req := httptest.NewRequest(http.MethodGet, url, nil)
	req.AddCookie(&http.Cookie{Name: auth.RefreshTokenCookieName, Value: refreshToken})
	rec = httptest.NewRecorder()
	e.ServeHTTP(rec, req)
	require.Equal(t, http.StatusOK, rec.Code)
	require.Equal(t, "secret content", rec.Body.String())
}
