package main

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"embed"
	"encoding/base64"
	"encoding/binary"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/skip2/go-qrcode"
	"roseh.moe/pkg/ccl"
	"roseh.moe/pkg/roseh.moe/internal/pwhash"
	"roseh.moe/pkg/wordlist"
)

var (
	port        = flag.Int("port", 42069, "port to listen on")
	selfURL     = flag.String("self-url", "http://localhost:42069", "base URL of the server")
	secretsFile = flag.String("secrets", "secrets.ccl", "secrets file")
	configFile  = flag.String("config", "config.ccl", "configuration file (ccl format https://pkg.go.dev/roseh.moe/pkg/ccl)")

	serverStartTime = time.Now()
)

var notepadPassword, secretKey []byte

func loadSecrets() error {
	fileBytes, err := os.ReadFile(*secretsFile)
	if err != nil {
		return err
	}
	var secretsConfig struct {
		NotepadPassword []byte
		SecretKey       []byte
	}
	if err := ccl.Unmarshal(fileBytes, &secretsConfig); err != nil {
		return fmt.Errorf("%s: %s", *secretsFile, err)
	}
	if len(secretsConfig.NotepadPassword) == 0 {
		return fmt.Errorf("%s: missing notepad-password", *secretsFile)
	}
	if len(secretsConfig.SecretKey) == 0 {
		return fmt.Errorf("%s: missing secret-key", *secretsFile)
	}
	notepadPassword = secretsConfig.NotepadPassword
	secretKey = secretsConfig.SecretKey
	return nil
}

type serviceConfiguration struct {
	packageRedirects map[string]string
	commandRedirects map[string]string
}

var (
	serviceConfigMu           sync.Mutex
	serviceConfig             *serviceConfiguration
	serviceConfigLastModified time.Time
)

func loadConfig() (*serviceConfiguration, error) {
	stat, err := os.Stat(*configFile)
	if err != nil {
		return nil, err
	}
	serviceConfigMu.Lock()
	oldConfig := serviceConfig
	configLastModified := serviceConfigLastModified
	serviceConfigMu.Unlock()
	newModTime := stat.ModTime()
	if !newModTime.After(configLastModified) {
		return oldConfig, nil
	}
	log.Printf("Reloading config file...")
	fileBytes, err := os.ReadFile(*configFile)
	if err != nil {
		return nil, err
	}
	var fileConfig struct {
		Redirect []struct{ Package, Command, To string }
	}
	if err := ccl.Unmarshal(fileBytes, &fileConfig); err != nil {
		return nil, err
	}
	config := &serviceConfiguration{
		packageRedirects: make(map[string]string),
		commandRedirects: make(map[string]string),
	}
	for _, redirect := range fileConfig.Redirect {
		if redirect.Package != "" {
			config.packageRedirects[redirect.Package] = redirect.To
		} else {
			config.commandRedirects[redirect.Command] = redirect.To
		}
	}
	serviceConfigMu.Lock()
	serviceConfig = config
	serviceConfigLastModified = newModTime
	serviceConfigMu.Unlock()
	return config, nil
}

var (
	//go:embed templates/404.html.template
	notFoundString   string
	notFoundTemplate = template.Must(template.New("notFound").Parse(notFoundString))
)

func notFound(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusNotFound)
	type args struct {
		Path string
	}
	if err := notFoundTemplate.Execute(w, args{Path: r.URL.Path}); err != nil {
		log.Printf("Warning: notFound: %s", err)
	}
}

var (
	//go:embed templates/outline.html.template
	outlineString   string
	outlineTemplate = template.Must(template.New("outline").Parse(outlineString))

	//go:embed templates/index.html.template
	indexString   string
	indexTemplate = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(indexString)).Lookup("outline")
)

func index(w http.ResponseWriter, _ *http.Request) {
	if err := indexTemplate.Execute(w, nil); err != nil {
		log.Printf("Warning: index: %s", err)
	}
}

var (
	//go:embed templates/pong.html.template
	pongString   string
	pongTemplate = template.Must(template.New("pong").Parse(pongString))
)

func pong(w http.ResponseWriter, _ *http.Request) {
	if err := pongTemplate.Execute(w, nil); err != nil {
		log.Printf("Warning: pong: %s", err)
	}
}

//go:embed static
var staticFiles embed.FS

func serveStaticFile(w http.ResponseWriter, r *http.Request, path string) {
	f, err := staticFiles.Open(path)
	if err != nil {
		notFound(w, r)
		return
	}
	defer f.Close()
	http.ServeContent(w, r, path, serverStartTime, f.(io.ReadSeeker))
}

func static(w http.ResponseWriter, r *http.Request) {
	serveStaticFile(w, r, strings.TrimPrefix(r.URL.Path, "/"))
}

func favicon(w http.ResponseWriter, r *http.Request) {
	serveStaticFile(w, r, "static/favicon.ico")
}

var (
	//go:embed templates/wormhole.html.template
	wormholeTemplateString string
	wormholeTemplate       = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(wormholeTemplateString)).Lookup("outline")
)

type wormholeTemplateArgs struct {
	Self string
	Hole string
}

func makeHole() string {
	const nWords = 10
	buf := make([]byte, 2*nWords)
	rand.Read(buf)
	words := make([]string, nWords)
	for i := range words {
		words[i] = wordlist.Words[binary.NativeEndian.Uint16(buf[2*i:])&0x1fff]
	}
	return strings.Join(words, "-")
}

func newWormhole(w http.ResponseWriter, r *http.Request) {
	http.Redirect(w, r, "/wormhole/"+makeHole()+"/upload", http.StatusSeeOther)
}

func wormhole(w http.ResponseWriter, r *http.Request) {
	if err := wormholeTemplate.Execute(w, wormholeTemplateArgs{Self: *selfURL, Hole: r.PathValue("hole")}); err != nil {
		log.Printf("Warning: wormhole: %s", err)
	}
}

type wormholeConn struct {
	done chan struct{}
	r    *multipart.Part
	w    http.ResponseWriter
}

var wormholeConnsMu sync.Mutex
var wormholeConns = make(map[string]wormholeConn)

func (c wormholeConn) wormholeCopy(ctx context.Context, hole string) error {
	wormholeConnsMu.Lock()
	prevConn, ok := wormholeConns[hole]
	if !ok {
		wormholeConns[hole] = c
		wormholeConnsMu.Unlock()
		defer func() {
			wormholeConnsMu.Lock()
			delete(wormholeConns, hole)
			wormholeConnsMu.Unlock()
		}()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-c.done:
			return nil
		}
	}
	wormholeConnsMu.Unlock()
	if c.w == nil {
		c.w = prevConn.w
	} else {
		c.r = prevConn.r
	}
	if c.w == nil || c.r == nil {
		return fmt.Errorf("connection already registered")
	}
	defer close(prevConn.done)
	c.w.Header().Set("Content-Disposition", "attachment; filename*=UTF-8''"+c.r.FileName())
	if _, err := io.Copy(c.w, c.r); err != nil {
		return fmt.Errorf("copy: %s", err)
	}
	return nil
}

var (
	//go:embed templates/wormhole-success.html.template
	wormholeSuccessTemplateString string
	wormholeSuccessTemplate       = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(wormholeSuccessTemplateString)).Lookup("outline")
)

func wormholeSend(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	reader, err := r.MultipartReader()
	if err != nil {
		http.Error(w, fmt.Sprintf("Not a multipart/form-data request: %s", err), http.StatusBadRequest)
		return
	}
	for {
		part, err := reader.NextPart()
		if err != nil {
			break
		}
		if part.FormName() != "file" {
			continue
		}
		hole := r.PathValue("hole")
		if err := (wormholeConn{done: make(chan struct{}), r: part}).wormholeCopy(ctx, hole); err != nil {
			http.Error(w, fmt.Sprintf("Error during copy: %s", err), http.StatusServiceUnavailable)
		}
		if err := wormholeSuccessTemplate.Execute(w, nil); err != nil {
			log.Printf("Warning: wormholeSend: %s", err)
		}
		return
	}
	http.Error(w, "file not found", http.StatusBadRequest)
}

func wormholeRecv(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	hole := r.PathValue("hole")
	if err := (wormholeConn{done: make(chan struct{}), w: w}).wormholeCopy(ctx, hole); err != nil {
		http.Error(w, fmt.Sprintf("Error during copy: %s", err), http.StatusServiceUnavailable)
	}
}

func wormholeQR(w http.ResponseWriter, r *http.Request) {
	qr, err := qrcode.Encode(*selfURL+"/wormhole/"+r.PathValue("hole"), qrcode.Medium, 200)
	if err != nil {
		http.Error(w, fmt.Sprintf("failed to encode QR code: %s", err), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "image/png")
	w.Write(qr)
}

func mac(msg []byte) []byte {
	mac := hmac.New(sha256.New, secretKey)
	mac.Write(msg)
	return mac.Sum(nil)
}

func sign(msg []byte) []byte {
	return append(msg, mac(msg)...)
}

func verify(msg []byte) ([]byte, bool) {
	if len(msg) < sha256.Size {
		return nil, false
	}
	msg, messageMAC := msg[:len(msg)-sha256.Size], msg[len(msg)-sha256.Size:]
	expectedMAC := mac(msg)
	if !hmac.Equal(messageMAC, expectedMAC) {
		return nil, false
	}
	return msg, true
}

func makeToken() (string, error) {
	b, err := time.Now().MarshalBinary()
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(sign(b)), nil
}

const cookieExpiration = 180 * 24 * time.Hour

func attachCookie(w http.ResponseWriter) error {
	token, err := makeToken()
	if err != nil {
		return err
	}
	http.SetCookie(w, &http.Cookie{
		Name:        "auth",
		Value:       token,
		Path:        "/",
		Expires:     time.Now().Add(cookieExpiration),
		Secure:      true,
		HttpOnly:    true,
		SameSite:    http.SameSiteStrictMode,
		Partitioned: true,
	})
	return nil
}

func cookieAuth(w http.ResponseWriter, r *http.Request) (string, bool) {
	const csrfTokenLen = 32
	var csrfToken string
	if csrfCookie, err := r.Cookie("csrf-token"); err == nil {
		csrfToken = csrfCookie.Value
	} else {
		buf := make([]byte, csrfTokenLen)
		rand.Read(buf)
		csrfToken = base64.RawURLEncoding.EncodeToString(buf)
		http.SetCookie(w, &http.Cookie{
			Name:        "csrf-token",
			Value:       csrfToken,
			Path:        "/",
			Secure:      true,
			HttpOnly:    true,
			SameSite:    http.SameSiteStrictMode,
			Partitioned: true,
		})
	}

	cookie, err := r.Cookie("auth")
	if err != nil {
		return csrfToken, false
	}
	authCookie, err := base64.RawURLEncoding.DecodeString(cookie.Value)
	if err != nil {
		return csrfToken, false
	}
	msg, ok := verify(authCookie)
	if !ok {
		return csrfToken, false
	}
	var t time.Time
	if err := t.UnmarshalBinary(msg); err != nil {
		return csrfToken, false
	}
	cookieAge := time.Since(t)
	if cookieAge > cookieExpiration {
		return csrfToken, false
	}
	if cookieAge > 24*time.Hour {
		attachCookie(w)
	}
	return csrfToken, true
}

func checkCSRFToken(r *http.Request) error {
	cookie, err := r.Cookie("csrf-token")
	if err != nil {
		return err
	}
	cookieHash := sha256.Sum256([]byte(cookie.Value))
	formValueHash := sha256.Sum256([]byte(r.FormValue("csrf-token")))
	if subtle.ConstantTimeCompare(cookieHash[:], formValueHash[:]) == 0 {
		return fmt.Errorf("bad CSRF token")
	}
	return nil
}

func checkPassword(password string) bool {
	expectedHash, salt := notepadPassword[:len(notepadPassword)-pwhash.SaltSize], notepadPassword[len(notepadPassword)-pwhash.SaltSize:]
	hash, err := pwhash.Hash(password, salt)
	return err == nil && subtle.ConstantTimeCompare(hash, expectedHash) != 0
}

var (
	//go:embed templates/login.html.template
	loginTemplateString string
	loginTemplate       = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(loginTemplateString)).Lookup("outline")
)

type loginTemplateArgs struct {
	CSRFToken string
	Error     bool
}

func executeLoginTemplate(w io.Writer, csrfToken string) {
	if err := loginTemplate.Execute(w, loginTemplateArgs{CSRFToken: csrfToken}); err != nil {
		log.Printf("Warning: login: %s", err)
	}
}

func login(w http.ResponseWriter, r *http.Request) {
	if err := checkCSRFToken(r); err != nil {
		http.Error(w, "bad CSRF token", http.StatusBadRequest)
		return
	}
	if !checkPassword(r.FormValue("password")) {
		if err := loginTemplate.Execute(w, loginTemplateArgs{
			Error:     true,
			CSRFToken: r.FormValue("csrf-token"),
		}); err != nil {
			log.Printf("Warning: login: %s", err)
		}
		return
	}
	attachCookie(w)
	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
}

var (
	//go:embed templates/note.html.template
	notepadString   string
	notepadTemplate = template.Must(template.Must(outlineTemplate.Clone()).New("body").Parse(notepadString)).Lookup("outline")

	notepadContentsMu sync.Mutex
	notepadContents   string
)

type notepadTemplateArgs struct {
	Content   string
	CSRFToken string
}

func notepad(w http.ResponseWriter, r *http.Request) {
	csrfToken, ok := cookieAuth(w, r)
	if !ok {
		executeLoginTemplate(w, csrfToken)
		return
	}
	notepadContentsMu.Lock()
	currentContent := notepadContents
	notepadContentsMu.Unlock()
	if err := notepadTemplate.Execute(w, notepadTemplateArgs{Content: currentContent, CSRFToken: csrfToken}); err != nil {
		log.Printf("Warning: notepad: %s", err)
	}
}

func saveNote(w http.ResponseWriter, r *http.Request) error {
	if err := checkCSRFToken(r); err != nil {
		return err
	}
	if _, ok := cookieAuth(w, r); !ok {
		return fmt.Errorf("not logged in")
	}
	newContent := r.FormValue("content")
	notepadContentsMu.Lock()
	notepadContents = newContent
	notepadContentsMu.Unlock()
	return nil
}

func autosave(w http.ResponseWriter, r *http.Request) {
	msg := "Saved ✓"
	if err := saveNote(w, r); err != nil {
		msg = fmt.Sprintf("Failed to save: %s", err)
	}
	if err := notepadTemplate.ExecuteTemplate(w, "saveIndicator", msg); err != nil {
		log.Printf("Warning: autosave: %s", err)
	}
}

func cmd(w http.ResponseWriter, r *http.Request) {
	serviceConfig, err := loadConfig()
	if err != nil {
		http.Error(w, fmt.Sprintf("load config: %s", err), http.StatusInternalServerError)
		return
	}
	source, ok := serviceConfig.commandRedirects[r.PathValue("cmd")]
	if !ok {
		notFound(w, r)
		return
	}
	if r.FormValue("go-get") == "1" {
		fmt.Fprintf(w, `<meta name="go-import" content="roseh.moe/cmd/%s %s">`, r.PathValue("cmd"), source)
	} else {
		http.Redirect(w, r, "https://pkg.go.dev/roseh.moe/cmd/"+r.PathValue("cmd"), http.StatusFound)
	}
}

func pkg(w http.ResponseWriter, r *http.Request) {
	serviceConfig, err := loadConfig()
	if err != nil {
		http.Error(w, fmt.Sprintf("load config: %s", err), http.StatusInternalServerError)
		return
	}
	source, ok := serviceConfig.packageRedirects[r.PathValue("pkg")]
	if !ok {
		notFound(w, r)
		return
	}
	if r.FormValue("go-get") == "1" {
		fmt.Fprintf(w, `<meta name="go-import" content="roseh.moe/pkg/%s %s">`, r.PathValue("pkg"), source)
	} else {
		http.Redirect(w, r, "https://pkg.go.dev/roseh.moe/pkg/"+r.PathValue("pkg"), http.StatusFound)
	}
}

func main() {
	flag.Parse()

	if err := loadSecrets(); err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("GET /pong", pong)
	http.HandleFunc("GET /wormhole", newWormhole)
	http.HandleFunc("GET /wormhole/{hole}/upload", wormhole)
	http.HandleFunc("POST /wormhole/{hole}/upload", wormholeSend)
	http.HandleFunc("GET /wormhole/{hole}", wormholeRecv)
	http.HandleFunc("GET /wormhole/{hole}/qr.png", wormholeQR)
	http.HandleFunc("POST /login", login)
	http.HandleFunc("GET /notepad", notepad)
	http.HandleFunc("POST /notepad", autosave)
	http.HandleFunc("GET /cmd/{cmd}", cmd)
	http.HandleFunc("GET /pkg/{pkg}", pkg)
	http.HandleFunc("GET /static/", static)
	http.HandleFunc("GET /favicon.ico", favicon)
	http.HandleFunc("GET /{$}", index)
	http.HandleFunc("GET /", notFound)

	addr := fmt.Sprintf(":%d", *port)
	log.Printf("Listening on %q", addr)
	log.Fatal(http.ListenAndServe(addr, nil))
}