package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"crypto/subtle"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"gitlab.com/rhogenson/roseh.moe/internal/pwhash"
)

var (
	port        = flag.Int("port", 42069, "port to listen on")
	secretsFile = flag.String("secrets", "secrets", "path to the secrets file")
	notepadFile = flag.String("notepad", "notepad", "path to save user notes")

	serverStartTime = time.Now()
)

var (
	notepadPassword     []byte
	notepadPasswordSalt []byte
	privateKey          []byte

	encryptionKeyMu sync.Mutex
	encryptionKey   []byte
)

func loadSecrets() error {
	secrets, err := os.ReadFile(*secretsFile)
	if err != nil {
		return err
	}
	for _, line := range bytes.Split(bytes.TrimSuffix(secrets, []byte("\n")), []byte("\n")) {
		if pw, ok := bytes.CutPrefix(line, []byte("notepad-password=")); ok {
			buf := make([]byte, hex.DecodedLen(len(pw)))
			if _, err := hex.Decode(buf, pw); err != nil {
				return err
			}
			if len(buf) < pwhash.SaltLen {
				return fmt.Errorf("bad password hash")
			}
			notepadPasswordSalt, notepadPassword = buf[:pwhash.SaltLen], buf[pwhash.SaltLen:]
		} else if key, ok := bytes.CutPrefix(line, []byte("secret-key=")); ok {
			if hex.DecodedLen(len(key)) != sha512.Size256 {
				return fmt.Errorf("invalid HMAC-SHA512/256 key")
			}
			privateKey = make([]byte, hex.DecodedLen(len(key)))
			if _, err := hex.Decode(privateKey, key); err != nil {
				return err
			}
		}
	}
	return 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")
}

const cookieExpiration = 400 * 24 * time.Hour

func attachCookie(w http.ResponseWriter) error {
	nowBytes, err := time.Now().MarshalBinary()
	if err != nil {
		return err
	}
	mac := hmac.New(sha512.New512_256, privateKey)
	mac.Write(nowBytes)
	http.SetCookie(w, &http.Cookie{
		Name:        "auth",
		Value:       base64.RawStdEncoding.EncodeToString(append(nowBytes, mac.Sum(nil)...)),
		Path:        "/notepad",
		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

	cookie, err := r.Cookie("auth")
	if err != nil {
		return "", false
	}
	bytes, err := base64.RawStdEncoding.DecodeString(cookie.Value)
	if err != nil {
		return "", false
	}
	mac := hmac.New(sha512.New512_256, privateKey)
	macSize := mac.Size()
	if len(bytes) < macSize {
		return "", false
	}
	msg, sig := bytes[:len(bytes)-macSize], bytes[len(bytes)-macSize:]
	mac.Write(msg)
	if !hmac.Equal(sig, mac.Sum(nil)) {
		return "", false
	}
	var t time.Time
	if err := t.UnmarshalBinary(msg); err != nil {
		return "", false
	}
	if time.Since(t) > cookieExpiration {
		return "", false
	}
	attachCookie(w)

	if csrfToken, err := r.Cookie("csrf-token"); err == nil {
		return csrfToken.Value, true
	}

	buf := make([]byte, csrfTokenLen)
	rand.Read(buf)
	csrfToken := base64.RawStdEncoding.EncodeToString(buf)
	http.SetCookie(w, &http.Cookie{
		Name:        "csrf-token",
		Value:       csrfToken,
		Path:        "/notepad",
		Secure:      true,
		HttpOnly:    true,
		SameSite:    http.SameSiteStrictMode,
		Partitioned: true,
	})
	return csrfToken, true
}

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

type loginTemplateArgs struct {
	Error bool
}

func login(w http.ResponseWriter, r *http.Request) {
	key, pwHash, err := pwhash.Hash(r.FormValue("password"), notepadPasswordSalt)
	if err != nil {
		http.Error(w, fmt.Sprintf("Unable to hash password: %s", err), http.StatusInternalServerError)
		return
	}
	if subtle.ConstantTimeCompare(pwHash[:], notepadPassword) == 0 {
		if err := loginTemplate.Execute(w, loginTemplateArgs{Error: true}); err != nil {
			log.Printf("Warning: login: %s", err)
		}
		return
	}
	encryptionKeyMu.Lock()
	if encryptionKey == nil {
		encryptionKey = key
	}
	encryptionKeyMu.Unlock()
	attachCookie(w)
	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
}

func readNotepad(key []byte) (string, error) {
	encrypted, err := os.ReadFile(*notepadFile)
	if err != nil {
		return "", err
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	aead, err := cipher.NewGCMWithRandomNonce(block)
	if err != nil {
		return "", err
	}
	decrypted, err := aead.Open(nil, nil, encrypted, nil)
	if err != nil {
		return "", err
	}
	return string(decrypted), nil
}

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

type notepadTemplateArgs struct {
	Content   string
	CSRFToken string
}

func notepad(w http.ResponseWriter, r *http.Request) {
	csrfToken, ok := cookieAuth(w, r)
	if !ok {
		if err := loginTemplate.Execute(w, loginTemplateArgs{}); err != nil {
			log.Printf("Warning: login: %s", err)
		}
		return
	}
	encryptionKeyMu.Lock()
	key := encryptionKey
	encryptionKeyMu.Unlock()
	if key == nil {
		if err := loginTemplate.Execute(w, loginTemplateArgs{}); err != nil {
			log.Printf("Warning: login: %s", err)
		}
		return
	}
	currentContent, err := readNotepad(key)
	if err != nil {
		currentContent = fmt.Sprintf("Error reading notepad file: %s", err)
	}
	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 csrfToken, err := r.Cookie("csrf-token"); err != nil || subtle.ConstantTimeCompare([]byte(csrfToken.Value), []byte(r.FormValue("csrf-token"))) == 0 {
		return fmt.Errorf("bad CSRF token")
	}
	if _, ok := cookieAuth(w, r); !ok {
		return fmt.Errorf("not logged in")
	}
	encryptionKeyMu.Lock()
	key := encryptionKey
	encryptionKeyMu.Unlock()
	if key == nil {
		return fmt.Errorf("not logged in")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	aead, err := cipher.NewGCMWithRandomNonce(block)
	if err != nil {
		return err
	}
	f, err := os.CreateTemp(filepath.Dir(*notepadFile), "notepad")
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err = f.Write(aead.Seal(nil, nil, []byte(r.FormValue("content")), nil)); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), *notepadFile)
}

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 main() {
	flag.Parse()

	if err := loadSecrets(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	http.HandleFunc("GET /pong", pong)
	http.HandleFunc("POST /login", login)
	http.HandleFunc("GET /notepad", notepad)
	http.HandleFunc("POST /notepad/autosave", autosave)
	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))
}