package storage

import (
	"context"
	"io"
	"mime"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
	"time"
)

func TestReadOptions(t *testing.T) {
	for _, name := range []string{"", "report.pdf", "报告 2026.pdf", `a";b.pdf`, "100%.pdf"} {
		h, err := (ReadOptions{Disposition: DispositionInline, FileName: name, ContentType: "Application/PDF"}).ResponseHeaders()
		if err != nil {
			t.Fatal(err)
		}
		disposition, params, err := mime.ParseMediaType(h.Get("Content-Disposition"))
		if err != nil || disposition != "inline" || params["filename"] != name || h.Get("Content-Type") != "application/pdf" || h.Get("Cache-Control") != "private, no-store" {
			t.Fatal("invalid read headers", h, err)
		}
	}
	h, err := (ReadOptions{}).ResponseHeaders()
	if err != nil || h.Get("Content-Disposition") != "attachment" || h.Get("Content-Type") != "" {
		t.Fatal(h, err)
	}
	for _, options := range []ReadOptions{
		{Disposition: "form-data"}, {Disposition: "Inline"}, {Disposition: "inline\r\nX: y"},
		{FileName: "x\r\nX: y"}, {FileName: "a\x00b"}, {FileName: "a\tb"}, {FileName: "bad\xff"}, {FileName: "a/b"}, {FileName: `a\b`}, {FileName: ".."}, {FileName: strings.Repeat("a", 1025)},
		{ContentType: "bad"}, {ContentType: "image/*"}, {ContentType: "application/pdf\r\nX: y"}, {ContentType: "application/pdf\x00"}, {ContentType: strings.Repeat("x", 257)},
	} {
		if _, err := options.ResponseHeaders(); err == nil {
			t.Fatal("invalid options accepted", options)
		}
	}
}

type readTestStore struct {
	calls, bytes int
	streaming    bool
}

func (*readTestStore) Put(context.Context, string, io.Reader, int64, string) (Info, error) {
	panic("unused")
}
func (*readTestStore) Delete(context.Context, string) error { panic("unused") }
func (s *readTestStore) Get(context.Context, string) (Object, error) {
	s.calls++
	body := &readTestBody{Reader: strings.NewReader("0123456789"), store: s}
	var stream io.ReadCloser = body
	if s.streaming {
		stream = struct {
			io.Reader
			io.Closer
		}{body, body}
	}
	return Object{Info: Info{Size: 10, ContentType: "application/pdf", ModifiedAt: time.Unix(1000, 0)}, Body: stream}, nil
}

type readTestBody struct {
	*strings.Reader
	store *readTestStore
}

func (b *readTestBody) Read(p []byte) (int, error) {
	n, err := b.Reader.Read(p)
	b.store.bytes += n
	return n, err
}
func (*readTestBody) Close() error { return nil }

func TestReadSignatureAndExpiry(t *testing.T) {
	store := &readTestStore{}
	s, _ := NewSigner(store, "https://files.example/read", []byte(strings.Repeat("s", 32)))
	now := time.Unix(1800000000, 0).UTC()
	s.now = func() time.Time { return now }
	link, err := s.PresignGetWithOptions(context.Background(), "folder/file", 30*time.Minute, ReadOptions{Disposition: DispositionInline, FileName: "报告.pdf", ContentType: "application/pdf"})
	if err != nil || !link.ExpiresAt.Equal(now.Add(30*time.Minute)) {
		t.Fatal("incorrect expiry", err)
	}
	serve := func(method, address string, headers http.Header) *httptest.ResponseRecorder {
		r := httptest.NewRequest(method, address, nil)
		if headers != nil {
			r.Header = headers
		}
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		return w
	}
	for _, method := range []string{"GET", "HEAD"} {
		w := serve(method, link.URL, nil)
		if w.Code != 200 || w.Header().Get("Content-Type") != "application/pdf" || w.Header().Get("Cache-Control") != "private, no-store" {
			t.Fatal(w)
		}
		if method == "HEAD" && w.Body.Len() != 0 {
			t.Fatal("HEAD body")
		}
	}
	old, err := s.PresignGet(context.Background(), "folder/file", 30*time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if w := serve("GET", old.String(), nil); w.Code != 200 || w.Header().Get("Content-Disposition") != "attachment" {
		t.Fatal("legacy download changed")
	}
	mutations := []func(url.Values){
		func(q url.Values) { q.Set("key", "other") }, func(q url.Values) { q.Set("expires", "1800001900") }, func(q url.Values) { q.Set("signature", "bad") },
		func(q url.Values) { q.Set("filename", "other.pdf") }, func(q url.Values) { q.Set("disposition", "attachment") }, func(q url.Values) { q.Set("contentType", "text/html") },
		func(q url.Values) { q.Set("v", "3") }, func(q url.Values) { q.Del("v") }, func(q url.Values) {
			for _, k := range []string{"v", "filename", "disposition", "contentType"} {
				q.Del(k)
			}
		},
		func(q url.Values) { q.Set("extra", "x") }, func(q url.Values) { q.Set("filename", "evil\r\nX: yes") },
	}
	for _, key := range []string{"key", "expires", "signature", "v", "disposition", "filename", "contentType"} {
		key := key
		mutations = append(mutations, func(q url.Values) { q.Add(key, q.Get(key)) })
	}
	for _, mutate := range mutations {
		u, _ := url.Parse(link.URL)
		q := u.Query()
		mutate(q)
		u.RawQuery = q.Encode()
		before := store.calls
		if w := serve("GET", u.String(), http.Header{"Range": {"bytes=0-1"}}); w.Code != 403 || store.calls != before {
			t.Fatal("tampered link reached storage")
		}
	}
	u := *old
	q := u.Query()
	q.Set("disposition", "inline")
	u.RawQuery = q.Encode()
	if serve("GET", u.String(), nil).Code != 403 {
		t.Fatal("legacy download upgraded by query")
	}
	if serve("GET", link.URL+"&bad=%ZZ", nil).Code != 403 {
		t.Fatal("malformed query accepted")
	}
	before := store.calls
	now = link.ExpiresAt.Add(-time.Nanosecond)
	if serve("GET", link.URL, nil).Code != 200 {
		t.Fatal("early expiration")
	}
	before = store.calls
	now = link.ExpiresAt
	for _, method := range []string{"GET", "HEAD"} {
		w := serve(method, link.URL, http.Header{"Range": {"bytes=0-1"}, "If-Modified-Since": {"Fri, 01 Jan 2100 00:00:00 GMT"}})
		if w.Code != 403 || store.calls != before || w.Header().Get("Cache-Control") != "private, no-store" {
			t.Fatal("expired link reached storage or cache")
		}
		if method == "HEAD" && w.Body.Len() != 0 {
			t.Fatal("HEAD error has body")
		}
	}
	for _, ttl := range []time.Duration{0, time.Second - time.Nanosecond, time.Hour + time.Nanosecond} {
		if _, err := s.PresignGetWithOptions(context.Background(), "key", ttl, ReadOptions{}); err == nil {
			t.Fatal("invalid TTL")
		}
	}
	for _, ttl := range []time.Duration{time.Second, time.Hour} {
		if _, err := s.PresignGetWithOptions(context.Background(), "key", ttl, ReadOptions{}); err != nil {
			t.Fatal(err)
		}
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := s.PresignGetWithOptions(ctx, "key", time.Minute, ReadOptions{}); err != context.Canceled {
		t.Fatal(err)
	}
}

func TestReadRangeStreamingAndCache(t *testing.T) {
	for _, streaming := range []bool{false, true} {
		store := &readTestStore{streaming: streaming}
		s, _ := NewSigner(store, "https://files.example/read", []byte(strings.Repeat("s", 32)))
		link, _ := s.PresignGetWithOptions(context.Background(), "key", 30*time.Minute, ReadOptions{})
		for _, method := range []string{"GET", "HEAD"} {
			before := store.bytes
			r := httptest.NewRequest(method, link.URL, nil)
			r.Header.Set("Range", "bytes=2-4")
			w := httptest.NewRecorder()
			s.ServeHTTP(w, r)
			if method == "HEAD" {
				if w.Code != 200 || w.Body.Len() != 0 || store.bytes != before {
					t.Fatal("HEAD read body", w)
				}
			} else if streaming {
				if w.Code != 200 || w.Body.String() != "0123456789" || w.Header().Get("Accept-Ranges") != "" {
					t.Fatal("stream fallback", w)
				}
			} else {
				if w.Code != 206 || w.Body.String() != "234" || store.bytes-before != 3 {
					t.Fatal("range buffered full body", w)
				}
			}
			if w.Header().Get("Cache-Control") != "private, no-store" {
				t.Fatal("missing cache policy")
			}
		}
		if !streaming {
			for _, h := range []http.Header{{"Range": {"bytes=99-100"}}, {"If-Modified-Since": {"Fri, 01 Jan 2100 00:00:00 GMT"}}} {
				r := httptest.NewRequest("GET", link.URL, nil)
				r.Header = h
				w := httptest.NewRecorder()
				s.ServeHTTP(w, r)
				if w.Code != 416 && w.Code != 304 {
					t.Fatal(w.Code)
				}
				if w.Header().Get("Cache-Control") != "private, no-store" {
					t.Fatal("cache policy lost on error/conditional")
				}
			}
		}
	}
}