// Package encoder implements the send-only YAMN v2 client encoder.
package encoder

import (
	"bufio"
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/mail"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/blake2s"
	"golang.org/x/crypto/nacl/box"
)

const (
	maxChainLength = 10
	headerBytes    = 256
	encHeadBytes   = 160
	bodyBytes      = 17920
	maxPlainBytes  = 17910
	messageBytes   = maxChainLength*headerBytes + bodyBytes
	armorVersion   = "0.2.7"
)

var (
	ErrInvalidRequest = errors.New("invalid YAMN request")
	ErrInvalidKeyring = errors.New("invalid YAMN public keyring")
)

type Kind string

const (
	Email  Kind = "email"
	Usenet Kind = "usenet"
)

// Request contains one plaintext message and an entry-to-exit chain.
// PublicKeyring is a local pubring.mix path and is never fetched by Encode.
type Request struct {
	Kind          Kind
	PublicKeyring string
	Entry         string
	Chain         []string
	From          string
	ReplyTo       string
	To            string
	Subject       string
	Newsgroup     string
	Body          string
	References    string
}

type Result struct {
	EntryAddress string
	Envelope     []byte
}

type remailerKey struct {
	name, address    string
	keyID, publicKey []byte
}

// Encode creates an armored YAMN v2 packet. It performs no network I/O.
func Encode(r Request) (Result, error) {
	if err := Validate(r); err != nil {
		return Result{}, err
	}
	keys, err := loadKeyring(r.PublicKeyring)
	if err != nil {
		return Result{}, err
	}
	chain := make([]remailerKey, len(r.Chain))
	for i, name := range r.Chain {
		key, ok := keys[name]
		if !ok {
			return Result{}, fmt.Errorf("remailer %q not found in keyring", name)
		}
		chain[i] = key
	}
	plain, err := composeMessage(r)
	if err != nil {
		return Result{}, err
	}
	packet, err := encodePacket(plain, chain)
	if err != nil {
		return Result{}, err
	}
	return Result{EntryAddress: chain[0].address, Envelope: armor(packet)}, nil
}

func Validate(r Request) error {
	if r.Kind != Email && r.Kind != Usenet {
		return fmt.Errorf("%w: unsupported message kind", ErrInvalidRequest)
	}
	if strings.TrimSpace(r.PublicKeyring) == "" {
		return fmt.Errorf("%w: missing public keyring", ErrInvalidRequest)
	}
	if len(r.Chain) == 0 || len(r.Chain) > maxChainLength {
		return fmt.Errorf("%w: invalid chain length", ErrInvalidRequest)
	}
	for _, hop := range r.Chain {
		if !isRemailerName(hop) {
			return fmt.Errorf("%w: invalid remailer name", ErrInvalidRequest)
		}
	}
	if strings.TrimSpace(r.Entry) != "" && strings.TrimSpace(r.Entry) != r.Chain[0] {
		return fmt.Errorf("%w: entry does not match chain", ErrInvalidRequest)
	}
	if strings.TrimSpace(r.Body) == "" || len([]byte(r.Body)) > maxPlainBytes {
		return fmt.Errorf("%w: body is empty or too large", ErrInvalidRequest)
	}
	if r.Kind == Email {
		if _, err := mail.ParseAddress(r.To); err != nil {
			return fmt.Errorf("%w: invalid recipient address", ErrInvalidRequest)
		}
	} else {
		if strings.TrimSpace(r.Newsgroup) == "" || !validHeaderValue(r.Newsgroup) {
			return fmt.Errorf("%w: missing or invalid newsgroup", ErrInvalidRequest)
		}
		if _, err := mail.ParseAddress(r.To); err != nil {
			return fmt.Errorf("%w: invalid Usenet gateway recipient", ErrInvalidRequest)
		}
	}
	for _, value := range []string{r.From, r.ReplyTo, r.To, r.Subject, r.Newsgroup, r.References} {
		if !validHeaderValue(value) {
			return fmt.Errorf("%w: invalid header value", ErrInvalidRequest)
		}
	}
	return nil
}

func composeMessage(r Request) ([]byte, error) {
	var b strings.Builder
	b.WriteString("Content-Type: text/plain; charset=utf-8\nContent-Transfer-Encoding: 8bit\nMIME-Version: 1.0\n")
	if r.From != "" {
		b.WriteString("From: " + r.From + "\n")
	}
	if r.ReplyTo != "" {
		b.WriteString("Reply-To: " + r.ReplyTo + "\n")
	}
	if r.To != "" {
		b.WriteString("To: " + r.To + "\n")
	}
	if r.Subject != "" {
		b.WriteString("Subject: " + r.Subject + "\n")
	}
	if r.Newsgroup != "" {
		b.WriteString("Newsgroups: " + r.Newsgroup + "\n")
	}
	if r.References != "" {
		b.WriteString("References: " + r.References + "\n")
	}
	b.WriteString("\n")
	b.WriteString(r.Body)
	return []byte(b.String()), nil
}

func validHeaderValue(value string) bool { return !strings.ContainsAny(value, "\r\n\x00") }
func isRemailerName(s string) bool {
	if s == "" {
		return false
	}
	for _, c := range s {
		if !(c == '-' || c == '_' || c >= 'a' && c <= 'z' || c >= '0' && c <= '9') {
			return false
		}
	}
	return true
}

func loadKeyring(path string) (map[string]remailerKey, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open public keyring: %w", err)
	}
	defer f.Close()
	keys := make(map[string]remailerKey)
	s := bufio.NewScanner(f)
	var current remailerKey
	var phase int
	for s.Scan() {
		line := strings.TrimSpace(s.Text())
		switch phase {
		case 0:
			parts := strings.Fields(line)
			if len(parts) != 7 {
				continue
			}
			id, e := hex.DecodeString(parts[2])
			if e != nil || len(id) != 16 {
				continue
			}
			from, e1 := time.Parse("2006-01-02", parts[5])
			until, e2 := time.Parse("2006-01-02", parts[6])
			if e1 != nil || e2 != nil {
				continue
			}
			now := time.Now().UTC()
			if now.Before(from) || now.After(until.Add(24*time.Hour)) {
				continue
			}
			current = remailerKey{name: parts[0], address: parts[1], keyID: id}
			phase = 1
		case 1:
			if line == "-----Begin Mix Key-----" {
				phase = 2
			}
		case 2:
			id, e := hex.DecodeString(line)
			if e != nil || !bytes.Equal(id, current.keyID) {
				phase = 0
			} else {
				phase = 3
			}
		case 3:
			key, e := hex.DecodeString(line)
			if e != nil || len(key) != 32 {
				phase = 0
			} else {
				current.publicKey = key
				phase = 4
			}
		case 4:
			if line == "-----End Mix Key-----" {
				keys[current.name] = current
			}
			phase = 0
		}
	}
	if err := s.Err(); err != nil {
		return nil, fmt.Errorf("read public keyring: %w", err)
	}
	if len(keys) == 0 {
		return nil, ErrInvalidKeyring
	}
	return keys, nil
}

func encodePacket(plain []byte, chain []remailerKey) ([]byte, error) {
	if len(chain) == 0 || len(chain) > maxChainLength || len(plain) > maxPlainBytes {
		return nil, ErrInvalidRequest
	}
	payload := make([]byte, messageBytes)
	if _, err := io.ReadFull(rand.Reader, payload); err != nil {
		return nil, err
	}
	copy(payload[maxChainLength*headerBytes:], plain)
	keys := make([][]byte, len(chain)-1)
	ivs := make([][]byte, len(chain)-1)
	for i := range keys {
		keys[i] = randomBytes(32)
		ivs[i] = randomBytes(12)
	}
	finalIV := randomBytes(16)
	finalID := randomBytes(16)
	exitAES := randomBytes(32)
	final := make([]byte, 64)
	copy(final, finalIV)
	final[16], final[17] = 1, 1
	copy(final[18:], finalID)
	binary.LittleEndian.PutUint32(final[34:38], uint32(len(plain)))
	copy(payload[maxChainLength*headerBytes:], aesCTR(payload[maxChainLength*headerBytes:], exitAES, finalIV))
	shiftHeaders(payload)
	if len(chain) > 1 {
		deterministic(payload, keys, ivs, len(chain), 0)
	}
	exitData := slotData(1, finalID, exitAES, final, antiTag(payload))
	copy(payload[:headerBytes], encryptedHeader(chain[len(chain)-1], exitData))
	for hop := 0; hop < len(chain)-1; hop++ {
		partial := ivs[hop]
		next := chain[len(chain)-hop-1].address
		info := make([]byte, 64)
		copy(info, partial)
		copy(info[12:], []byte(next))
		encryptAll(payload, keys[hop], ivs[hop])
		shiftHeaders(payload)
		deterministic(payload, keys, ivs, len(chain), hop+1)
		data := slotData(0, randomBytes(16), keys[hop], info, antiTag(payload))
		copy(payload[:headerBytes], encryptedHeader(chain[len(chain)-hop-2], data))
	}
	return payload, nil
}

func slotData(kind byte, packetID, aesKey, info, tag []byte) []byte {
	b := make([]byte, encHeadBytes)
	b[0], b[1], b[2] = 2, kind, 0
	copy(b[3:], packetID)
	copy(b[19:], aesKey)
	binary.LittleEndian.PutUint16(b[51:53], uint16(time.Now().UTC().Unix()/86400))
	copy(b[53:], info)
	copy(b[117:], tag)
	return b
}

func encryptedHeader(key remailerKey, data []byte) []byte {
	var recipient, sender [32]byte
	copy(recipient[:], key.publicKey)
	if _, err := io.ReadFull(rand.Reader, sender[:]); err != nil {
		panic(err)
	}
	var nonce [24]byte
	if _, err := io.ReadFull(rand.Reader, nonce[:]); err != nil {
		panic(err)
	}
	sealed := box.Seal(nil, data, &nonce, &recipient, &sender)
	b := make([]byte, headerBytes)
	copy(b, key.keyID)
	copy(b[16:], sender[:])
	copy(b[48:], nonce[:])
	copy(b[72:], sealed)
	return b
}

func aesCTR(input, key, iv []byte) []byte {
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	output := make([]byte, len(input))
	cipher.NewCTR(block, iv).XORKeyStream(output, input)
	return output
}
func randomBytes(n int) []byte {
	b := make([]byte, n)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		panic(err)
	}
	return b
}
func shiftHeaders(p []byte) {
	copy(p[headerBytes:maxChainLength*headerBytes], p[:(maxChainLength-1)*headerBytes])
}
func antiTag(p []byte) []byte {
	h, _ := blake2s.New256(nil)
	h.Write(p[headerBytes:])
	return h.Sum(nil)
}
func seqIV(partial []byte, slot int) []byte {
	iv := make([]byte, 16)
	copy(iv[:4], partial[:4])
	binary.LittleEndian.PutUint32(iv[4:8], uint32(slot))
	copy(iv[8:], partial[4:])
	return iv
}
func encryptAll(p, key, partial []byte) {
	for slot := 0; slot <= maxChainLength; slot++ {
		start := slot * headerBytes
		end := start + headerBytes
		if slot == maxChainLength {
			start = maxChainLength * headerBytes
			end = len(p)
		}
		copy(p[start:end], aesCTR(p[start:end], key, seqIV(partial, slot)))
	}
}
func deterministic(p []byte, keys, ivs [][]byte, chainLen, hop int) {
	bottom := maxChainLength - 1
	top := bottom - (chainLen - hop - 2)
	for slot := top; slot <= bottom; slot++ {
		fake := make([]byte, headerBytes)
		use := bottom
		for i := bottom - slot + hop; i >= 0; i-- {
			copy(fake, aesCTR(fake, keys[i], seqIV(ivs[i], use)))
			use--
		}
		copy(p[slot*headerBytes:(slot+1)*headerBytes], fake)
	}
}

func armor(payload []byte) []byte {
	h, _ := blake2s.New256(nil)
	h.Write(payload)
	var b bytes.Buffer
	b.WriteString("::\nRemailer-Type: yamn-" + armorVersion + "\n\n-----BEGIN REMAILER MESSAGE-----\n")
	b.WriteString(strconv.Itoa(len(payload)) + "\n")
	b.WriteString(hex.EncodeToString(h.Sum(nil)) + "\n")
	encoded := base64.StdEncoding.EncodeToString(payload)
	for len(encoded) > 0 {
		n := 64
		if len(encoded) < n {
			n = len(encoded)
		}
		b.WriteString(encoded[:n])
		b.WriteByte('\n')
		encoded = encoded[n:]
	}
	b.WriteString("\n-----END REMAILER MESSAGE-----\n")
	return b.Bytes()
}