package main

import (
	"bytes"
	"embed"
	"fmt"
	"io"
	"log"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path"
	"slices"
	"strconv"
	"strings"
	"text/template"
	"time"
)

var (
	//go:embed private/*
	privateDirectory embed.FS
	serveHash        = fmt.Sprintf("%x", time.Now().Unix())
	serveDirectory   = os.Getenv("HTTP_DIRECTORY")
	serveAddress     = os.Getenv("HTTP_ADDRESS")
	tmpl             = template.Must(
		template.
			New("directory.html").
			Funcs(template.FuncMap{
				"EncodeURI":         url.PathEscape,
				"GetEntryIcon":      getEntryIcon,
				"GetEntrySizeHuman": getEntrySizeHuman,
			}).
			ParseFS(privateDirectory, "private/directory.html"),
	)
)

type Entry struct {
	Name    string
	IsDir   bool
	Size    int64
	ModTime time.Time
}

func serveFileError(w http.ResponseWriter, err error) {
	if os.IsNotExist(err) {
		w.WriteHeader(http.StatusNotFound)
		return
	}
	if os.IsPermission(err) {
		w.WriteHeader(http.StatusForbidden)
		return
	}
	log.Println("Filesystem Error:", err)
	http.Error(w, "Server Error", http.StatusInternalServerError)
}

func parentPath(urlPath string) string {
	trimmed := strings.TrimSuffix(urlPath, "/")
	if trimmed == "" {
		return ""
	}
	idx := strings.LastIndex(trimmed, "/")
	if idx <= 0 {
		return "/"
	}
	return trimmed[:idx+1]
}

func getEntrySizeHuman(b int64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "KMGTPE"[exp])
}

func getEntryIcon(name string, isDir bool) string {
	if isDir {
		return "icon_folder.png"
	}
	switch strings.ToLower(path.Ext(name)) {
	// Custom Filetypes
	case ".redir":
		return "icon_redirect.png"

	// Default Filetypes
	case ".mp3", ".flac", ".ogg", ".wav", ".aac", ".opus", ".m4a":
		return "icon_audio.png"
	case ".mp4", ".mkv", ".webm", ".avi", ".mov", ".m4v":
		return "icon_video.png"
	case ".jpg", ".jpeg", ".png", ".gif", ".webp", ".avif", ".bmp", ".tiff":
		return "icon_photo.png"
	case ".zip", ".tar", ".gz", ".7z", ".rar", ".zst", ".xz", ".bz2":
		return "icon_archive.png"
	case ".exe", ".elf", ".bin", ".out", ".appimage", ".deb", ".rpm":
		return "icon_program.png"
	case ".pdf", ".txt", ".md", ".srt", ".vtt", ".ass", ".log":
		return "icon_document.png"
	case ".html", ".ts", ".js", ".go", ".rs", ".pug":
		return "icon_script.png"
	case ".json", ".xml", ".yaml", ".toml", ".ini", ".cfg":
		return "icon_config.png"
	default:
		return "icon_generic.png"
	}
}

func init() {
	if serveAddress == "" {
		serveAddress = "127.0.0.1:8080"
	}
	if serveDirectory == "" {
		serveDirectory = "./public"
	}
}

func main() {
	mux := http.NewServeMux()
	prv := http.FileServerFS(privateDirectory)

	mux.HandleFunc("/private/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
		prv.ServeHTTP(w, r)
	})

	mux.HandleFunc("/preferences/sort", func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		http.SetCookie(w, &http.Cookie{
			Name:   "sort",
			Value:  r.FormValue("sort"),
			Path:   "/",
			MaxAge: 31536000,
		})
		http.Redirect(w, r, r.FormValue("redirect"), http.StatusSeeOther)
	})

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {

		filepath := path.Join(serveDirectory, path.Clean(r.URL.Path))
		filestat, err := os.Stat(filepath)
		if err != nil {
			serveFileError(w, err)
			return
		}

		// Use Browser Cache (if possible)
		modtime := filestat.ModTime().UTC()
		modhash := fmt.Sprint(modtime.Unix())
		if r.Header.Get("If-None-Match") == modhash {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		if ts, err := time.Parse(http.TimeFormat, r.Header.Get("If-Modified-Since")); err == nil {
			if !modtime.After(ts) {
				w.WriteHeader(http.StatusNotModified)
				return
			}
		}
		if ts, err := time.Parse(http.TimeFormat, r.Header.Get("If-Unmodified-Since")); err == nil {
			if modtime.After(ts) {
				w.WriteHeader(http.StatusPreconditionFailed)
				return
			}
		}

		if filestat.IsDir() {
			// List Directory Contents

			indexPath := path.Join(filepath, "index.html")
			if _, err := os.Stat(indexPath); err == nil {
				http.ServeFile(w, r, indexPath)
				return
			}

			dirEntries, err := os.ReadDir(filepath)
			if err != nil {
				serveFileError(w, err)
				return
			}

			// Entry Information
			entries := make([]Entry, 0, len(dirEntries))
			for _, entry := range dirEntries {
				info, err := entry.Info()
				if err != nil {
					continue
				}

				// Check Symbolic Link
				isDir := entry.IsDir()
				if entry.Type()&os.ModeSymlink != 0 {
					if stat, err := os.Stat(path.Join(filepath, entry.Name())); err == nil {
						isDir = stat.IsDir()
						info = stat
					}
				}

				entries = append(entries, Entry{
					Name:    entry.Name(),
					IsDir:   isDir,
					Size:    info.Size(),
					ModTime: info.ModTime().UTC(),
				})
			}

			// Entry Sorting
			sortOrder := "alpha"
			if c, err := r.Cookie("sort"); err == nil && c.Value != "" {
				sortOrder = c.Value
			}
			slices.SortFunc(entries, func(a, b Entry) int {
				// Folders are always on top
				if a.IsDir != b.IsDir {
					if a.IsDir {
						return -1
					}
					return 1
				}
				// Use
				if sortOrder == "modified" {
					if a.ModTime.Equal(b.ModTime) {
						return strings.Compare(a.Name, b.Name)
					}
					if a.ModTime.After(b.ModTime) {
						return -1
					}
					return 1
				}
				return strings.Compare(a.Name, b.Name)
			})

			// Directory Listing
			w.Header().Set("Content-Type", "text/html; charset=utf-8")
			if err := tmpl.Execute(w, map[string]any{
				"Path":      r.URL.Path,
				"Parent":    parentPath(r.URL.Path),
				"Item":      entries,
				"ItemCount": len(entries),
				"Version":   serveHash,
				"Sort":      sortOrder,
			}); err != nil {
				serveFileError(w, err)
				return
			}

		} else {
			// Serve File Contents

			f, err := os.Open(filepath)
			if err != nil {
				serveFileError(w, err)
				return
			}
			defer f.Close()

			w.Header().Set("Accept-Ranges", "bytes")
			w.Header().Set("Content-Type", mime.TypeByExtension(path.Ext(filepath)))
			w.Header().Set("Last-Modified", modtime.Format(http.TimeFormat))
			w.Header().Set("ETag", modhash)

			// Special File Handler: .redir
			if strings.EqualFold(path.Ext(filepath), ".redir") {
				raw, _ := io.ReadAll(f)
				uri, err := url.Parse(string(bytes.TrimSpace(raw)))
				if err != nil {
					serveFileError(w, err)
					return
				}
				http.Redirect(w, r, uri.String(), http.StatusTemporaryRedirect)
				return
			}

			// Serve Byte Range (if applicable)
			rangeHeader := r.Header.Get("Range")
			if rangeHeader != "" {
				var start, end int64
				fmt.Sscanf(rangeHeader, "bytes=%d-%d", &start, &end)
				if end == 0 || end >= filestat.Size() {
					end = filestat.Size() - 1
				}
				length := end - start + 1
				w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, filestat.Size()))
				w.Header().Set("Content-Length", strconv.FormatInt(length, 10))
				w.WriteHeader(http.StatusPartialContent)
				f.Seek(start, io.SeekStart)
				io.CopyN(w, f, length)
				return
			}

			// Serve Raw Contents
			if _, err := io.Copy(w, f); err != nil {
				serveFileError(w, err)
				return
			}
		}

	})

	svr := http.Server{
		Handler: mux,
		Addr:    serveAddress,
	}
	log.Printf("Listening @ %s\n", serveAddress)
	if err := svr.ListenAndServe(); err != nil {
		log.Fatalf("Listen Error: %s\n", err)
	}
}