package main

import (
	"crypto/rand"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

var uploadDir string

func initConfig() {
	if envUploadDir := os.Getenv("UPLOAD_DIR"); envUploadDir != "" {
		uploadDir = envUploadDir
	} else {
		uploadDir = "./uploads/"
	}
}

func main() {
	initConfig()

	if len(os.Args) < 2 {
		fmt.Println("Usage: go run main.go <port>")
		return
	}

	port := os.Args[1]

	uploadEndpoint := "upload"
	if len(os.Args) >= 3 {
		uploadEndpoint = os.Args[2]
	}

	http.HandleFunc(fmt.Sprintf("/%s/", uploadEndpoint), uploadHandler)
	fmt.Printf("Server started at :%s\n", port)
	fmt.Printf("Endpoint: /%s/\n", uploadEndpoint)
	fmt.Printf("Upload Directory: %s\n", uploadDir)
	http.ListenAndServe(":"+port, nil)
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed)
		return
	}

	// Parse the multipart form
	err := r.ParseMultipartForm(10 << 20) // 10 MB
	if err != nil {
		http.Error(w, "Error parsing form data", http.StatusBadRequest)
		return
	}

	file, handler, err := r.FormFile("file")
	if err != nil {
		http.Error(w, "Error retrieving the file", http.StatusBadRequest)
		return
	}
	defer file.Close()

	// Sanitize the filename
	sanitizedFilename := sanitizeFilename(handler.Filename)
	if sanitizedFilename == "" {
		sanitizedFilename = generateRandomFilename()
	}

	filePath := filepath.Join(uploadDir, sanitizedFilename)

	// Ensure unique filename
	for fileExists(filePath) {
		sanitizedFilename = generateRandomFilename()
		filePath = filepath.Join(uploadDir, sanitizedFilename)
	}

	// Create the uploads directory if it doesn't exist
	if _, err := os.Stat(uploadDir); os.IsNotExist(err) {
		os.MkdirAll(uploadDir, os.ModePerm)
	}

	// Save the file
	out, err := os.Create(filePath)
	if err != nil {
		http.Error(w, "Unable to save the file", http.StatusInternalServerError)
		return
	}
	defer out.Close()

	_, err = io.Copy(out, file)
	if err != nil {
		http.Error(w, "Error saving the file", http.StatusInternalServerError)
		return
	}

	w.Write([]byte("File uploaded successfully"))
}

func sanitizeFilename(filename string) string {
	// Allow only alphanumeric, dashes, underscores, and dots
	re := regexp.MustCompile(`[^a-zA-Z0-9._-]`)
	sanitized := re.ReplaceAllString(filename, "")
	sanitized = strings.TrimLeft(sanitized, ".")

	return sanitized
}

func generateRandomFilename() string {
	const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, 16)
	for i := range b {
		b[i] = letters[randInt(len(letters))]
	}
	return string(b) + ".file"
}

func randInt(n int) int {
	b := make([]byte, 1)
	rand.Read(b)
	return int(b[0]) % n
}

func fileExists(filename string) bool {
	info, err := os.Stat(filename)
	if os.IsNotExist(err) {
		return false
	}
	return !info.IsDir()
}