package secretstore

import (
	"bufio"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"syscall"

	"golang.org/x/crypto/argon2"
)

type Store struct {
	Salt [aes.BlockSize]byte
	Keys []KeyEntry

	unlocked bool
	key      [32]byte
}

type KeyEntry struct {
	Name   string
	Hash   [64]byte
	EncKey [32]byte
}

func New() (s *Store) {
	s = &Store{}
	syscall.Mlock(s.key[:])
	return
}

func Open(path string) (s *Store, err error) {
	f, err := os.Open(path)
	if err != nil {
		return
	}

	defer f.Close()

	s = New()
	_, err = s.ReadFrom(bufio.NewReader(f))

	return
}

func (s *Store) SaveTo(path string) (err error) {
	f, err := os.OpenFile(path, syscall.O_CREAT|syscall.O_TRUNC|syscall.O_WRONLY, 0600)
	if err != nil {
		return
	}

	defer f.Close()

	out := bufio.NewWriter(f)

	_, err = s.WriteTo(out)
	if err != nil {
		return
	}

	err = out.Flush()
	if err != nil {
		return
	}

	return
}

func (s *Store) Close() {
	memzero(s.key[:])
	syscall.Munlock(s.key[:])
	s.unlocked = false
}

func (s *Store) IsNew() bool {
	return len(s.Keys) == 0
}

func (s *Store) Unlocked() bool {
	return s.unlocked
}

func (s *Store) Init(name string, passphrase []byte) (err error) {
	err = randRead(s.key[:])
	if err != nil {
		return
	}
	err = randRead(s.Salt[:])
	if err != nil {
		return
	}

	s.AddKey(name, passphrase)

	s.unlocked = true

	return
}

var jsonFormatHdr = []byte("{json}")

func (s *Store) ReadFrom(in io.Reader) (n int64, err error) {
	memzero(s.key[:])
	s.unlocked = false

	defer func() {
		if err != nil {
			log.Output(2, fmt.Sprintf("failed after %d bytes", n))
		}
	}()

	readFull := func(ba []byte) {
		var nr int
		nr, err = io.ReadFull(in, ba)
		n += int64(nr)
	}

	// read the file's start (json header or start of salt)

	readFull(s.Salt[:len(jsonFormatHdr)])
	if err != nil {
		return
	}

	if !bytes.Equal(s.Salt[:len(jsonFormatHdr)], jsonFormatHdr) {
		// old key file

		// finish reading the salt
		readFull(s.Salt[len(jsonFormatHdr):])
		if err != nil {
			return
		}

		// read the (encrypted) keys
		s.Keys = make([]KeyEntry, 0)
		for {
			k := KeyEntry{Name: "key-" + strconv.Itoa(len(s.Keys))}
			readFull(k.Hash[:])
			if err != nil {
				if err == io.EOF {
					err = nil
				}
				return
			}
			readFull(k.EncKey[:])
			if err != nil {
				return
			}

			s.Keys = append(s.Keys, k)
		}
	}

	err = json.NewDecoder(in).Decode(s)
	return
}

func (s *Store) WriteTo(out io.Writer) (n int64, err error) {
	_, err = out.Write(jsonFormatHdr)
	if err != nil {
		return
	}
	err = json.NewEncoder(out).Encode(s)
	return
}

var ErrNoSuchKey = errors.New("no such key")

func (s *Store) Unlock(passphrase []byte) (ok bool) {
	key, hash := s.keyPairFromPassword(passphrase)
	memzero(passphrase)
	defer memzero(key[:])

	var idx = -1
	for i := range s.Keys {
		if hash == s.Keys[i].Hash {
			idx = i
			break
		}
	}

	if idx == -1 {
		return
	}

	s.decryptTo(s.key[:], s.Keys[idx].EncKey[:], &key)

	s.unlocked = true
	return true
}

func (s *Store) AddKey(name string, passphrase []byte) {
	key, hash := s.keyPairFromPassword(passphrase)
	memzero(passphrase)

	defer memzero(key[:])

	k := KeyEntry{Name: name, Hash: hash}

	encKey := s.encrypt(s.key[:], &key)
	copy(k.EncKey[:], encKey)

	s.Keys = append(s.Keys, k)
}

func (s *Store) keyPairFromPassword(password []byte) (key [32]byte, hash [64]byte) {
	keySlice := argon2.IDKey(password, s.Salt[:], 1, 64*1024, 4, 32)

	copy(key[:], keySlice)
	memzero(keySlice)

	hash = sha512.Sum512(key[:])

	return
}

func (s *Store) NewEncrypter(iv [aes.BlockSize]byte) cipher.Stream {
	if !s.unlocked {
		panic("not unlocked")
	}
	return newEncrypter(iv, &s.key)
}

func (s *Store) NewDecrypter(iv [aes.BlockSize]byte) cipher.Stream {
	if !s.unlocked {
		panic("not unlocked")
	}
	return newDecrypter(iv, &s.key)
}

func (s *Store) encrypt(src []byte, key *[32]byte) (dst []byte) {
	dst = make([]byte, len(src))
	newEncrypter(s.Salt, key).XORKeyStream(dst, src)
	return
}

func (s *Store) decryptTo(dst []byte, src []byte, key *[32]byte) {
	newDecrypter(s.Salt, key).XORKeyStream(dst, src)
}

func newEncrypter(iv [aes.BlockSize]byte, key *[32]byte) cipher.Stream {
	c, err := aes.NewCipher(key[:])
	if err != nil {
		panic(fmt.Errorf("failed to init AES: %w", err))
	}

	return cipher.NewCFBEncrypter(c, iv[:])
}

func newDecrypter(iv [aes.BlockSize]byte, key *[32]byte) cipher.Stream {
	c, err := aes.NewCipher(key[:])
	if err != nil {
		panic(fmt.Errorf("failed to init AES: %w", err))
	}

	return cipher.NewCFBDecrypter(c, iv[:])
}