package main

import (
	"crypto/rand"
	"encoding/hex"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"mime"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/msteinert/pam"
)

type Config struct {
	Listen      string
	StorageDir  string
	BaseURL     string
	MaxFileSize int64
	ExpireHours int
  IndexHTML   string
}

var config Config

var indexHTML string

func loadIndexHTML() {
    content, err := os.ReadFile(config.IndexHTML)
    if err != nil {
        panic(err)
    }
    indexHTML = string(content)
}

func init() {
	flag.StringVar(&config.Listen, "listen", ":8080", "Address to listen on")
	flag.StringVar(&config.StorageDir, "storage", "/tmp/share", "Directory to store uploaded files")
	flag.StringVar(&config.BaseURL, "baseurl", "http://localhost:8080", "Base URL for generated file links")
	flag.Int64Var(&config.MaxFileSize, "maxsize", 10*1024*1024, "Maximum allowed file size in bytes")
	flag.IntVar(&config.ExpireHours, "expire", 24, "Number of hours before files are deleted")
  flag.StringVar(&config.IndexHTML, "index", "index.html", "Path to html file to serve as index")
}

func authenticateUser(username, password string) bool {
	t, err := pam.StartFunc("system-auth", username, func(s pam.Style, msg string) (string, error) {
		switch s {
		case pam.PromptEchoOff:
			return password, nil
		}
		return "", fmt.Errorf("unsupported PAM style")
	})
	if err != nil {
		return false
	}
	err = t.Authenticate(0)
	return err == nil
}

func generateRandomFilename() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}

func cleanupOldFiles() {
	ticker := time.NewTicker(1 * time.Hour)
	for range ticker.C {
		now := time.Now()
		err := filepath.Walk(config.StorageDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() && now.Sub(info.ModTime()) > time.Duration(config.ExpireHours)*time.Hour {
				os.Remove(path)
			}
			return nil
		})
		if err != nil {
			log.Printf("Cleanup error: %v", err)
		}
	}
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPut {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	username, password, ok := r.BasicAuth()
	if !ok || !authenticateUser(username, password) {
		w.Header().Set("WWW-Authenticate", `Basic realm="Upload"`)
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	if r.ContentLength > config.MaxFileSize {
		http.Error(w, "File too large", http.StatusRequestEntityTooLarge)
		return
	}

	ext := path.Ext(r.URL.Path)
	if ext == "" {
		contentType := r.Header.Get("Content-Type")
		exts, _ := mime.ExtensionsByType(contentType)
		if len(exts) > 0 {
			ext = exts[0]
		}
	}

	randomName, err := generateRandomFilename()
	if err != nil {
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	filename := randomName + ext
	filepath := path.Join(config.StorageDir, filename)

	f, err := os.Create(filepath)
	if err != nil {
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}
	defer f.Close()

	_, err = io.Copy(f, r.Body)
	if err != nil {
		os.Remove(filepath)
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	fileURL := strings.TrimRight(config.BaseURL, "/") + "/" + filename
	fmt.Fprintf(w, "%s\n", fileURL)
}

func handleIndex(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.ServeFile(w, r, path.Join(config.StorageDir, r.URL.Path))
		return
	}

	tmpl, err := template.New("index").Parse(indexHTML)
	if err != nil {
		http.Error(w, "Internal server error", http.StatusInternalServerError)
		return
	}

	tmpl.Execute(w, config)
}

func main() {
	flag.Parse()
  loadIndexHTML()

	if err := os.MkdirAll(config.StorageDir, 0755); err != nil {
		log.Fatalf("Failed to create storage directory: %v", err)
	}

	go cleanupOldFiles()

	http.HandleFunc("/", handleIndex)
	http.HandleFunc("/upload", handleUpload)

	log.Printf("Starting server on %s", config.Listen)
	log.Fatal(http.ListenAndServe(config.Listen, nil))
}