package static

import (
	"bytes"
	"errors"
	"fmt"
	"io/fs"
	"net/url"
	"os"
	pathpkg "path"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"

	"github.com/gofiber/utils/v2"
	"github.com/valyala/fasthttp"

	"github.com/gofiber/fiber/v3"
)

var ErrInvalidPath = errors.New("invalid path")

const invalidPathSentinel = "/__fiber_invalid__"

func bytesToPathString(p []byte) string {
	if bytes.IndexByte(p, '\\') >= 0 {
		b := make([]byte, len(p))
		copy(b, p)
		for i := range b {
			if b[i] == '\\' {
				b[i] = '/'
			}
		}
		return utils.UnsafeString(b)
	}

	return utils.UnsafeString(p)
}

func unescapePathString(s string) (string, error) {
	for strings.IndexByte(s, '%') >= 0 {
		us, err := url.PathUnescape(s)
		if err != nil {
			return "", ErrInvalidPath
		}
		if us == s {
			break
		}
		s = us
	}

	return s, nil
}

func hasParentDirSegment(p []byte) (bool, error) {
	s, err := unescapePathString(bytesToPathString(p))
	if err != nil {
		return false, err
	}

	trimmed := utils.TrimLeft(filepath.ToSlash(s), '/')
	for trimmed != "" {
		segment, rest, found := strings.Cut(trimmed, "/")
		if segment == ".." {
			return true, nil
		}
		if !found {
			return false, nil
		}
		trimmed = rest
	}

	return false, nil
}

// sanitizePath validates and cleans the requested path.
// It returns an error if the path attempts to traverse directories.
func sanitizePath(p []byte, filesystem fs.FS) ([]byte, error) {
	hasTrailingSlash := len(p) > 0 && p[len(p)-1] == '/'

	s, err := unescapePathString(bytesToPathString(p))
	if err != nil {
		return nil, err
	}

	if strings.IndexByte(s, '\\') >= 0 {
		return nil, ErrInvalidPath
	}

	// reject any null bytes
	if strings.IndexByte(s, '\x00') >= 0 {
		return nil, ErrInvalidPath
	}

	normalized := filepath.ToSlash(s)
	if filesystem == nil && strings.HasPrefix(normalized, "//") {
		return nil, ErrInvalidPath
	}

	s = pathpkg.Clean("/" + normalized)

	trimmed := utils.TrimLeft(s, '/')
	if trimmed != "" {
		if slices.Contains(strings.Split(trimmed, "/"), "..") {
			return nil, ErrInvalidPath
		}
	}

	if filesystem == nil {
		normalizedClean := filepath.ToSlash(trimmed)
		if strings.HasPrefix(normalizedClean, "//") {
			return nil, ErrInvalidPath
		}
		if volume := filepath.VolumeName(normalizedClean); volume != "" {
			return nil, ErrInvalidPath
		}
		if len(normalizedClean) >= 2 && normalizedClean[1] == ':' {
			drive := normalizedClean[0]
			if (drive >= 'a' && drive <= 'z') || (drive >= 'A' && drive <= 'Z') {
				return nil, ErrInvalidPath
			}
		}
		if strings.HasPrefix(filepath.ToSlash(s), "//") {
			return nil, ErrInvalidPath
		}
	}

	if filesystem != nil {
		s = trimmed
		if s == "" {
			return []byte("/"), nil
		}
		if !fs.ValidPath(s) {
			return nil, ErrInvalidPath
		}
		s = "/" + s
	}

	if hasTrailingSlash && len(s) > 1 && s[len(s)-1] != '/' {
		s += "/"
	}

	return utils.UnsafeBytes(s), nil
}

// New creates a new middleware handler.
// The root argument specifies the root directory from which to serve static assets.
//
// Note: Root has to be string or fs.FS; otherwise, it will panic.
func New(root string, cfg ...Config) fiber.Handler {
	config := configDefault(cfg...)

	var createFS sync.Once
	var fileHandler fasthttp.RequestHandler
	var cacheControlValue string
	var rootCheckErr error
	var rootIsFile bool

	// adjustments for io/fs compatibility
	if config.FS != nil && root == "" {
		root = "."
	}

	return func(c fiber.Ctx) error {
		// Don't execute middleware if Next returns true
		if config.Next != nil && config.Next(c) {
			return c.Next()
		}

		// We only serve static assets on GET or HEAD methods
		method := c.Method()
		if method != fiber.MethodGet && method != fiber.MethodHead {
			return c.Next()
		}

		// Initialize FS
		createFS.Do(func() {
			prefix := c.Route().Path

			rootIsFile, rootCheckErr = isFile(root, config.FS)

			// Is prefix a partial wildcard?
			if before, _, found := strings.Cut(prefix, "*"); found {
				// /john* -> /john
				prefix = before
			}

			prefixLen := len(prefix)
			if prefixLen > 1 && prefix[prefixLen-1:] == "/" {
				// /john/ -> /john
				prefixLen--
			}

			// For io/fs.FS, Root must be empty so fasthttp's pathToFilePath
			// returns clean relative paths without prefixing the root.
			// PathRewrite already handles file-root and subdirectory cases.
			fsRoot := root
			if config.FS != nil {
				fsRoot = ""
			}

			var fsRootPrefix []byte
			if config.FS != nil && root != "" && root != "." && !rootIsFile {
				// fasthttp.FS.Root is forced to "" for io/fs.FS so pathToFilePath
				// returns clean relative paths. PathRewrite prepends the caller's
				// configured subdirectory after sanitizing the request path.
				fsRootPrefix = make([]byte, len(root)+1)
				fsRootPrefix[0] = '/'
				copy(fsRootPrefix[1:], root)
			}

			fileServer := &fasthttp.FS{
				Root:                   fsRoot,
				FS:                     config.FS,
				AllowEmptyRoot:         true,
				GenerateIndexPages:     config.Browse,
				AcceptByteRange:        config.ByteRange,
				Compress:               config.Compress,
				CompressBrotli:         config.Compress, // Brotli compression won't work without this
				CompressZstd:           config.Compress, // Zstd compression won't work without this
				CompressedFileSuffixes: c.App().Config().CompressedFileSuffixes,
				CacheDuration:          config.CacheDuration,
				SkipCache:              config.CacheDuration < 0,
				IndexNames:             config.IndexNames,
				PathNotFound: func(fctx *fasthttp.RequestCtx) {
					fctx.Response.SetStatusCode(fiber.StatusNotFound)
				},
			}

			fileServer.PathRewrite = func(fctx *fasthttp.RequestCtx) []byte {
				path := fctx.Path()

				if len(path) >= prefixLen {
					if rootCheckErr != nil && fileServer.FS != nil {
						return []byte(invalidPathSentinel)
					}

					// If the root is a file, we need to reset the path to "/" always.
					switch {
					case rootIsFile && fileServer.FS == nil:
						path = []byte("/")
					case rootIsFile && fileServer.FS != nil:
						path = utils.UnsafeBytes(root)
					default:
						path = path[prefixLen:]
						if len(fsRootPrefix) > 0 {
							hasTraversal, err := hasParentDirSegment(path)
							if err != nil || hasTraversal {
								return []byte(invalidPathSentinel)
							}
						}
						if len(path) == 0 || path[len(path)-1] != '/' {
							path = append(path, '/')
						}
					}
				}

				if len(path) > 0 && path[0] != '/' {
					path = append([]byte("/"), path...)
				}

				sanitized, err := sanitizePath(path, fileServer.FS)
				if err != nil {
					// return a guaranteed-missing path so fs responds with 404
					return []byte(invalidPathSentinel)
				}
				if len(fsRootPrefix) > 0 {
					rewrittenPath := make([]byte, len(fsRootPrefix)+len(sanitized))
					copy(rewrittenPath, fsRootPrefix)
					copy(rewrittenPath[len(fsRootPrefix):], sanitized)
					return rewrittenPath
				}
				return sanitized
			}

			maxAge := config.MaxAge
			if maxAge > 0 {
				cacheControlValue = "public, max-age=" + strconv.Itoa(maxAge)
			}

			fileHandler = fileServer.NewRequestHandler()
		})

		// Serve file
		fileHandler(c.RequestCtx())

		// Sets the response Content-Disposition header to attachment if the Download option is true
		if config.Download {
			name := filepath.Base(c.Path())
			if rootIsFile {
				name = filepath.Base(root)
			}
			c.Attachment(name)
		}

		// Return request if found and not forbidden
		status := c.RequestCtx().Response.StatusCode()

		if status != fiber.StatusNotFound && status != fiber.StatusForbidden {
			if cacheControlValue != "" {
				c.RequestCtx().Response.Header.Set(fiber.HeaderCacheControl, cacheControlValue)
			}

			if config.ModifyResponse != nil {
				return config.ModifyResponse(c)
			}

			return nil
		}

		// Return custom 404 handler if provided.
		if config.NotFoundHandler != nil {
			return config.NotFoundHandler(c)
		}

		// Reset response to default
		c.RequestCtx().SetContentType("") // Issue #420
		c.RequestCtx().Response.SetStatusCode(fiber.StatusOK)
		c.RequestCtx().Response.SetBodyString("")

		// Next middleware
		return c.Next()
	}
}

// isFile checks if the root is a file.
func isFile(root string, filesystem fs.FS) (bool, error) {
	var file fs.File
	var err error

	if filesystem != nil {
		file, err = filesystem.Open(root)
		if err != nil {
			return false, fmt.Errorf("static: %w", err)
		}
		defer func() {
			_ = file.Close() //nolint:errcheck // not needed
		}()
	} else {
		file, err = os.Open(filepath.Clean(root))
		if err != nil {
			return false, fmt.Errorf("static: %w", err)
		}
		defer func() {
			_ = file.Close() //nolint:errcheck // not needed
		}()
	}

	stat, err := file.Stat()
	if err != nil {
		return false, fmt.Errorf("static: %w", err)
	}

	return stat.Mode().IsRegular(), nil
}
