package main

import (
	"bytes"
	"crypto/rand"
	"crypto/subtle"
	"embed"
	"encoding/base64"
	"encoding/binary"
	"encoding/gob"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitlab.com/rhogenson/roseh.moe/internal/api"
	"gitlab.com/rhogenson/roseh.moe/internal/cryptoutil"
)

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

	serverStartTime = time.Now()
)

var (
	notepadPassword cryptoutil.PasswordHash
	secretKey       cryptoutil.HMACKey

	encryptionKeyMu sync.Mutex
	encryptionKey   cryptoutil.EncryptionKey
)

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 {
			if hex.DecodedLen(len(pw)) != cryptoutil.PasswordHashSize {
				return fmt.Errorf("invalid PBKDF2-SHA256 hash")
			}
			notepadPassword = make(cryptoutil.PasswordHash, hex.DecodedLen(len(pw)))
			if _, err := hex.Decode(notepadPassword, pw); err != nil {
				return err
			}
		} else if key, ok := bytes.CutPrefix(line, []byte("secret-key=")); ok {
			if hex.DecodedLen(len(key)) != cryptoutil.HMACKeySize {
				return fmt.Errorf("invalid HMAC-SHA256 key")
			}
			secretKey = make(cryptoutil.HMACKey, hex.DecodedLen(len(key)))
			if _, err := hex.Decode(secretKey, 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")
}

type authToken struct {
	Issued time.Time
}

func makeToken() (string, error) {
	buf := new(bytes.Buffer)
	if err := gob.NewEncoder(buf).Encode(authToken{Issued: time.Now()}); err != nil {
		return "", err
	}
	return base64.RawStdEncoding.EncodeToString(secretKey.Sign(buf.Bytes(), "auth")), nil
}

func checkToken(token string) bool {
	authCookie, err := base64.RawStdEncoding.DecodeString(token)
	if err != nil {
		return false
	}
	msg, ok := secretKey.Verify(authCookie, "auth")
	if !ok {
		return false
	}
	var t authToken
	if err := gob.NewDecoder(bytes.NewReader(msg)).Decode(&t); err != nil {
		return false
	}
	return time.Since(t.Issued) < cookieExpiration
}

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:        "/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
	}
	if !checkToken(cookie.Value) {
		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
}

func setEncryptionKey(rawKey cryptoutil.RawKey) error {
	encryptionKeyMu.Lock()
	currentKey := encryptionKey
	encryptionKeyMu.Unlock()
	if currentKey != nil {
		return nil
	}
	key, err := rawKey.EncryptionKey()
	if err != nil {
		return err
	}
	encryptionKeyMu.Lock()
	if encryptionKey == nil {
		encryptionKey = key
	}
	encryptionKeyMu.Unlock()
	return nil
}

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) {
	rawKey, err := notepadPassword.CheckPassword(r.FormValue("password"))
	if err != nil {
		if err := loginTemplate.Execute(w, loginTemplateArgs{Error: true}); err != nil {
			log.Printf("Warning: login: %s", err)
		}
		return
	}
	if err := setEncryptionKey(rawKey); err != nil {
		http.Error(w, fmt.Sprintf("Error: %s", err), http.StatusInternalServerError)
		return
	}
	attachCookie(w)
	http.Redirect(w, r, "/notepad", http.StatusSeeOther)
}

func readNotepad(key cryptoutil.EncryptionKey) (string, error) {
	f, err := os.Open(*notepadDir + "/notepad")
	if err != nil {
		return "", err
	}
	defer f.Close()
	notepad, err := io.ReadAll(key.NewReader(f, []byte("notepad")))
	if err != nil {
		return "", err
	}
	return string(notepad), 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 errors.New("not logged in")
	}
	f, err := os.CreateTemp(*notepadDir, "notepad")
	if err != nil {
		return err
	}
	defer f.Close()
	encryptingWriter := key.NewWriter(f, []byte("notepad"))
	if _, err := encryptingWriter.Write([]byte(r.FormValue("content"))); err != nil {
		os.Remove(f.Name())
		return err
	}
	if err := encryptingWriter.Close(); err != nil {
		os.Remove(f.Name())
		return err
	}
	if err := f.Close(); err != nil {
		os.Remove(f.Name())
		return err
	}
	return os.Rename(f.Name(), *notepadDir+"/notepad")
}

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 writeGob(w io.Writer, resp any) {
	if err := gob.NewEncoder(w).Encode(&api.Response{Ok: resp}); err != nil {
		log.Printf("Warning: write gob response: %s", err)
	}
}

func gobError(w io.Writer, err error) {
	var code api.ErrorCode
	if !errors.As(err, &code) {
		code = api.Internal
	}
	if err := gob.NewEncoder(w).Encode(&api.Response{Status: code, Err: err.Error()}); err != nil {
		log.Printf("Warning: write gob response: %s", err)
	}
}

func apiLogin(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/gob")
	rawKey, err := notepadPassword.CheckPassword(r.Header.Get("Roseh-Password"))
	if err != nil {
		gobError(w, fmt.Errorf("%w: incorrect password", api.PermissionDenied))
		return
	}
	if err := setEncryptionKey(rawKey); err != nil {
		gobError(w, err)
		return
	}
	token, err := makeToken()
	if err != nil {
		gobError(w, err)
		return
	}
	writeGob(w, &api.LoginResponse{Token: token})
}

type encoder struct {
	enc *gob.Encoder
}

func (e *encoder) send(resp any) error {
	return e.enc.Encode(&api.Response{Ok: resp})
}

func listNotes(req *api.ListNotesRequest) (*api.ListNotesResponse, error) {
	notes, err := os.ReadDir(*notepadDir + "/notes")
	if err != nil {
		return nil, err
	}
	names := make([]string, len(notes))
	for i, note := range notes {
		names[i] = note.Name()
	}
	return &api.ListNotesResponse{Notes: names}, nil
}

func dumpFile(w io.Writer, key cryptoutil.EncryptionKey, name string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = io.Copy(w, key.NewReader(f, []byte("chunks/"+filepath.Base(name))))
	return err
}

func chunkName(prefix string, i int64) string {
	return fmt.Sprintf("%s%013s", prefix, strconv.FormatUint(uint64(i), 36))
}

var (
	//go:embed wordlist.txt
	wordListString string
	wordList       = strings.Split(strings.TrimSuffix(wordListString, "\n"), "\n")
)

func createNote(req *api.CreateNoteRequest) (*api.CreateNoteResponse, error) {
	type continuationToken struct {
		Prefix string
		I      int64
	}

	encryptionKeyMu.Lock()
	key := encryptionKey
	encryptionKeyMu.Unlock()
	if key == nil {
		return nil, fmt.Errorf("%w: need login", api.PermissionDenied)
	}

	var contToken continuationToken
	if len(req.ContinuationToken) > 0 {
		msg, ok := secretKey.Verify(req.ContinuationToken, "continuation-token")
		if !ok {
			return nil, fmt.Errorf("%w: invalid continuation token", api.InvalidArgument)
		}
		if err := gob.NewDecoder(bytes.NewReader(msg)).Decode(&contToken); err != nil {
			return nil, fmt.Errorf("bad continuation token")
		}
	} else {
		token := make([]byte, 64)
		rand.Read(token)
		contToken.Prefix = base64.RawURLEncoding.EncodeToString(token)
		contToken.I = 0
	}

	if req.More {
		if err := os.MkdirAll(*notepadDir+"/chunks", 0700); err != nil {
			return nil, err
		}
		name := chunkName(contToken.Prefix, contToken.I)
		f, err := os.OpenFile(*notepadDir+"/chunks/"+name, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
		if err != nil {
			if errors.Is(err, os.ErrExist) {
				return nil, fmt.Errorf("%w: continuation token expired", api.InvalidArgument)
			}
			return nil, err
		}
		defer f.Close()
		if contToken.I > 0 {
			_, err := os.Stat(*notepadDir + "/chunks/" + chunkName(contToken.Prefix, 0))
			if err != nil {
				if errors.Is(err, os.ErrNotExist) {
					return nil, fmt.Errorf("%w: continuation token expired", api.InvalidArgument)
				}
				return nil, err
			}
		}
		encryptingWriter := key.NewWriter(f, []byte("chunks/"+name))
		if _, err := encryptingWriter.Write(req.Chunk); err != nil {
			return nil, err
		}
		if err := encryptingWriter.Close(); err != nil {
			return nil, err
		}
		if err := f.Close(); err != nil {
			return nil, err
		}
		contToken.I++
		buf := new(bytes.Buffer)
		if err := gob.NewEncoder(buf).Encode(contToken); err != nil {
			return nil, err
		}
		return &api.CreateNoteResponse{ContinuationToken: secretKey.Sign(buf.Bytes(), "continuation-token")}, nil
	}
	if err := os.MkdirAll(*notepadDir+"/notes", 0700); err != nil {
		return nil, err
	}
	var finalNote *os.File
	for n := 1; ; n++ {
		buf := make([]byte, 2*n)
		rand.Read(buf)
		words := make([]string, n)
		for i := range words {
			words[i] = wordList[binary.NativeEndian.Uint16(buf[2*i:])&0x1fff]
		}
		name := strings.Join(words, "-")
		var err error
		finalNote, err = os.OpenFile(*notepadDir+"/notes/"+name, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
		if err != nil {
			if errors.Is(err, os.ErrExist) {
				continue
			}
			return nil, fmt.Errorf("create note: create note file: %s", err)
		}
		defer finalNote.Close()
		break
	}
	writer := key.NewWriter(finalNote, []byte("notes/"+filepath.Base(finalNote.Name())))
	for i := range contToken.I {
		if err := dumpFile(writer, key, *notepadDir+"/chunks/"+chunkName(contToken.Prefix, i)); err != nil {
			return nil, err
		}
	}
	if _, err := writer.Write(req.Chunk); err != nil {
		return nil, err
	}
	if err := writer.Close(); err != nil {
		return nil, err
	}
	if err := finalNote.Close(); err != nil {
		return nil, err
	}
	for i := range contToken.I {
		if err := os.Remove(*notepadDir + "/chunks/" + chunkName(contToken.Prefix, i)); err != nil {
			log.Printf("Warning: createNote: failed to delete chunk: %s", err)
		}
	}
	return &api.CreateNoteResponse{Name: filepath.Base(finalNote.Name())}, nil
}

type readNoteResponseWriter struct {
	stream *encoder
}

func (w *readNoteResponseWriter) Write(buf []byte) (int, error) {
	if err := w.stream.send(&api.ReadNoteResponseStream{Chunk: buf}); err != nil {
		return 0, err
	}
	return len(buf), nil
}

func readNote(stream *encoder, req *api.ReadNoteRequest) error {
	f, err := os.OpenInRoot(*notepadDir+"/notes", req.Note)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return fmt.Errorf("%w: note %q", api.NotFound, req.Note)
		}
		return err
	}
	defer f.Close()
	encryptionKeyMu.Lock()
	key := encryptionKey
	encryptionKeyMu.Unlock()
	if key == nil {
		return fmt.Errorf("%w: need login", api.PermissionDenied)
	}
	_, err = io.Copy(&readNoteResponseWriter{stream: stream}, key.NewReader(f, []byte("notes/"+req.Note)))
	return err
}

func tokenAuth(w http.ResponseWriter, r *http.Request) bool {
	token := r.Header.Get("Roseh-Token")
	if token == "" {
		gobError(w, api.Unauthenticated)
		return false
	}
	if !checkToken(token) {
		gobError(w, fmt.Errorf("%w: bad token", api.PermissionDenied))
		return false
	}
	return true
}

func gobReqRespMiddleware[Request, Response any](next func(*Request) (*Response, error)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/gob")
		if !tokenAuth(w, r) {
			return
		}
		req := new(Request)
		if err := gob.NewDecoder(r.Body).Decode(req); err != nil {
			gobError(w, fmt.Errorf("%w: read body: %s", api.BadRequest, err))
			return
		}
		resp, err := next(req)
		if err != nil {
			gobError(w, err)
			return
		}
		writeGob(w, resp)
	}
}

func gobRespStreamMiddleware[Request any](next func(*encoder, *Request) error) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/gob")
		if !tokenAuth(w, r) {
			return
		}
		req := new(Request)
		if err := gob.NewDecoder(r.Body).Decode(req); err != nil {
			gobError(w, fmt.Errorf("%w: read body: %s", api.BadRequest, err))
			return
		}
		gobEncoder := gob.NewEncoder(w)
		if err := next(&encoder{gobEncoder}, req); err != nil {
			gobError(w, err)
			return
		}
	}
}

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("POST /api/login", apiLogin)
	http.HandleFunc("GET /api/list-notes", gobReqRespMiddleware(listNotes))
	http.HandleFunc("POST /api/create-note", gobReqRespMiddleware(createNote))
	http.HandleFunc("GET /api/read-note", gobRespStreamMiddleware(readNote))
	http.HandleFunc("GET /static/", static)
	http.HandleFunc("GET /favicon.ico", favicon)
	http.HandleFunc("GET /{$}", index)
	http.HandleFunc("GET /", notFound)

	addr := fmt.Sprintf(":%d", *port)
	server := &http.Server{
		Addr:         addr,
		ReadTimeout:  10 * time.Minute,
		WriteTimeout: time.Minute,
	}
	log.Printf("Listening on %q", addr)
	log.Fatal(server.ListenAndServe())
}