package main

import (
	"bufio"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/pbkdf2"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"slices"
	"strings"
	"syscall"

	"golang.org/x/term"

	_ "modernc.org/sqlite"
)

var (
	serverAddress string
	username      string
	password      string
	dataDir       string
	dataPath      string
	cacheDir      string
	cachePath     string
	scanner       bufio.Scanner
)

type Note struct {
	Body  string `json:"body"`
	Color int64  `json:"color"`
	Date  string `json:"date"`
	ID    string `json:"id"`
	Title string `json:"title"`
}

func check(err error, format string, a ...any) bool {
	if err != nil {
		if writeErr := os.WriteFile("error.log", []byte(fmt.Sprintf(format+": %s", a, err)), 0600); writeErr != nil {
			log.Fatalf("Error writing error '%s' to error.log: %s\n", err, writeErr)
		}
		log.Fatalf("\033[31m"+format+"\033[0m: %s\n", a, err)
		return false
	}
	return true
}

func readPassword() (string, error) {
	stdin := int(syscall.Stdin)
	oldState, err := term.GetState(stdin)
	if err != nil {
		return "", err
	}
	defer term.Restore(stdin, oldState)

	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt)

	go func() {
		<-sigch
		term.Restore(stdin, oldState)
		os.Exit(1)
	}()

	password, err := term.ReadPassword(stdin)
	if err != nil {
		return "", err
	}

	return string(password), nil
}

func pkcs7strip(data []byte, blockSize int) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("pkcs7: Data is empty")
	}
	if length%blockSize != 0 {
		return nil, errors.New("pkcs7: Data is not block-aligned")
	}
	padLen := int(data[length-1])
	ref := bytes.Repeat([]byte{byte(padLen)}, padLen)
	if padLen > blockSize || padLen == 0 || !bytes.HasSuffix(data, ref) {
		return nil, errors.New("pkcs7: Invalid padding")
	}
	return data[:length-padLen], nil
}

func main() {
	scanner = *bufio.NewScanner(os.Stdin)
	home := os.Getenv("HOME")
	dataPath = os.Getenv("XDG_DATA_HOME")
	cachePath = os.Getenv("XDG_CACHE_HOME")
	if len(dataPath) < 1 {
		dataPath = home + "/.local/share"
	}
	if len(cachePath) < 1 {
		cachePath = home + "/.cache"
	}
	cacheDir = cachePath + "/notes"
	cachePath += "/notes/notes.db"
	dataDir = dataPath + "/notes"
	dataPath += "/notes/userinfo.json"

	err := os.Mkdir(cacheDir, 0700)
	if !errors.Is(err, os.ErrExist) {
		check(err, "Error creating cache directory")
	}
	db, err := sql.Open("sqlite", cachePath)
	check(err, "Error opening cache database at '%s'", cachePath)
	check(db.Ping(), "Ping of cache database failed!")
	db.Exec(`
CREATE TABLE notes (
	id TEXT PRIMARY KEY,
	data TEXT
);
	`)

	args := os.Args[1:]
	var command string
	if len(args) < 1 {
		command = ""
	} else {
		command = args[0]
	}
	switch command {
	case "ls":
		userDataBytes, err := os.ReadFile(dataPath)
		if errors.Is(err, fs.ErrNotExist) {
			login()
		}
		check(err, "Error reading user data at '%s'", dataPath)
		var userData []string
		err = json.Unmarshal(userDataBytes, &userData)
		check(err, "Error unmarshalling user data")

		serverAddress, username, password = userData[0], userData[1], userData[2]

		fzf := exec.Command("fzf")
		stdin, err := fzf.StdinPipe()
		check(err, "Error getting stdin pipe for fzf")
		check(fzf.Start(), "Error starting fzf")

		cacheJsonBody := load(db)
		listAll(cacheJsonBody, stdin)
		fetchJsonBody := fetch()

		for i := range fetchJsonBody {
			if !slices.Contains(cacheJsonBody, fetchJsonBody[i]) {
				go list(fetchJsonBody[i][1], stdin, Note{})
			}
		}

		cache(db, fetchJsonBody)
		fzf.Wait()
	default:
		login()
	}
}

func cache(db *sql.DB, jsonData [][2]string) {
	for i := range jsonData {
		var id string
		err := db.QueryRow("SELECT id FROM notes WHERE id = ?", jsonData[i][0]).Scan(&id)
		if err != sql.ErrNoRows {
			check(err, "Error scanning cache database at %s for duplicates", cachePath)
			continue
		}

		stmt, err := db.Prepare("INSERT INTO notes VALUES(?, ?)")
		check(err, "Error preparing to insert note into cache database at %s", cachePath)
		defer stmt.Close()

		_, err = stmt.Exec(jsonData[i][0], jsonData[i][1])
		check(err, "Error executing insert note into cache database at %s", cachePath)
	}
}

func login() {
	fmt.Println("Welcome to Notes! To get started, choose your server and user.")
	fmt.Print("Server Address [notes.miningtcup.me]: ")
	_, err := fmt.Scanln(&serverAddress)
	if err != nil {
		serverAddress = "notes.miningtcup.me"
	}

	fmt.Print("Username: ")
	scanner.Scan()
	username = url.QueryEscape(scanner.Text())

	fmt.Print("Password: ")
	passwordBytes, err := readPassword()
	check(err, "Error reading user input (password)")
	password = url.QueryEscape(string(passwordBytes))

	jsonData, err := json.Marshal([]string{serverAddress, username, password})
	check(err, "Error marshalling server address and credentials")
	err = os.Mkdir(dataDir, 0700)
	if !errors.Is(err, fs.ErrExist) {
		check(err, "Error creating data directory at '%s'", dataDir)
	}
	err = os.WriteFile(dataPath, jsonData, 0600)
	check(err, "Error writing to '%s'", dataPath)

	fmt.Println("\nAwesome! To get started, run 'notes add' or 'notes list'")
}

func decrypt(data string, password string) string {
	dataBytes, err := base64.StdEncoding.DecodeString(data)
	check(err, "Error decoding base64 data")

	if len(dataBytes) < 32 {
		return ""
	}

	iv := dataBytes[0:16]
	salt := dataBytes[16:32]
	ciphertext := dataBytes[32:]

	if len(ciphertext)%aes.BlockSize != 0 {
		fmt.Printf("Ciphertext length %d is not a multiple of block size %d\n", len(ciphertext), aes.BlockSize)
		return ""
	}

	key, err := pbkdf2.Key(sha256.New, password, salt, 65536, 32)
	check(err, "Error deriving pbkdf2-sha256 key from password")
	block, err := aes.NewCipher(key)
	check(err, "Error creating AES key")
	mode := cipher.NewCBCDecrypter(block, iv)

	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	strippedPlaintext, err := pkcs7strip(plaintext, 16)
	if err != nil {
		return ""
	}
	return string(strippedPlaintext)
}

func fetch() [][2]string {
	response, err := http.Get(fmt.Sprintf("https://%s/get?user=%s", serverAddress, url.QueryEscape(username)))
	check(err, "Error getting notes from '%s' as '%s'", serverAddress, username)
	defer response.Body.Close()
	body, err := io.ReadAll(response.Body)
	check(err, "Error reading response body while getting notes")

	var jsonBody [][2]string
	err = json.Unmarshal(body, &jsonBody)
	check(err, "Error unmarshalling response body")

	return jsonBody
}

func load(db *sql.DB) [][2]string {
	var rowCount int
	err := db.QueryRow("SELECT COUNT(*) FROM notes").Scan(&rowCount)
	if err != nil {
		log.Fatal(err)
	}
	const limit = 100
	if rowCount > limit { // change this if there's ever an option to increase limit
		rowCount = limit // change this if there's ever an option to increase limit
	}

	rows, err := db.Query("SELECT * FROM notes LIMIT ?", rowCount)
	check(err, "Error querying cache database", err)
	defer rows.Close()

	notes := make([][2]string, rowCount)

	for i := 0; rows.Next(); i++ {
		var currentRow [2]string
		check(rows.Scan(&currentRow[0], &currentRow[1]), "Error scanning cache database rows")
		notes[i] = currentRow
	}

	return notes
}

func listAll(jsonBody [][2]string, fzfStdin io.WriteCloser) {
	decryptedBody := make([]Note, len(jsonBody))
	for i := range jsonBody {
		go list(jsonBody[i][1], fzfStdin, decryptedBody[i])
	}
}

func list(jsonBody string, fzfStdin io.WriteCloser, decryptedBody Note) {
	err := json.Unmarshal([]byte(decrypt(jsonBody, password)), &decryptedBody)
	if err != nil {
		return
	}
	if len(decryptedBody.Title) > 0 && len(decryptedBody.Body) > 0 {
		fmt.Fprintf(fzfStdin, "%s | %s\n", decryptedBody.Title, strings.ReplaceAll(decryptedBody.Body, "\n", " "))
	}
}