// Package bangcle implements the "Bangcle" envelope every BYD API request and
// response is wrapped in: white-box AES-CBC, base64, with an "F" prefix.
//
// The cipher has no key schedule. Its round functions are pre-computed lookup
// tables extracted from the app's libencrypt.so, shipped here as tables.bin and
// embedded in the binary. Ported from pybyd's _crypto/_bangcle_block.py.
package bangcle

import (
	_ "embed"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"strings"
)

//go:embed tables.bin
var tablesBin []byte

// tables is built at init because tables.bin is compiled into the binary: if it
// fails to parse the build is broken, and there is no runtime recovery.
var tables = mustLoadTables(tablesBin)

// Encode wraps a payload as "F" + base64(whitebox-AES-CBC(pkcs7(plain))).
func Encode(plain []byte) string {
	return "F" + base64.StdEncoding.EncodeToString(encryptCBC(tables, padPKCS7(plain)))
}

// Decode unwraps an envelope produced by the BYD API.
func Decode(envelope string) ([]byte, error) {
	s := strings.NewReplacer(" ", "", "\t", "", "\n", "", "\r", "", "-", "+", "_", "/").Replace(envelope)
	if s == "" {
		return nil, fmt.Errorf("bangcle: empty input")
	}
	if !strings.HasPrefix(s, "F") {
		return nil, fmt.Errorf("bangcle: envelope must start with %q", "F")
	}
	s = s[1:]
	if r := len(s) % 4; r != 0 {
		s += strings.Repeat("=", 4-r)
	}
	ct, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		return nil, fmt.Errorf("bangcle: invalid base64: %w", err)
	}
	if len(ct) == 0 || len(ct)%16 != 0 {
		return nil, fmt.Errorf("bangcle: ciphertext length %d is not a positive multiple of 16", len(ct))
	}
	return StripPKCS7(decryptCBC(tables, ct)), nil
}

// PadPKCS7 and StripPKCS7 are exported because the inner AES-CBC layer in
// package byd uses the same padding rules.
func PadPKCS7(data []byte) []byte { return padPKCS7(data) }

func padPKCS7(data []byte) []byte {
	n := 16 - len(data)%16 // 1..16, so an aligned input gets a whole extra block
	pad := make([]byte, n)
	for i := range pad {
		pad[i] = byte(n)
	}
	return append(data, pad...)
}

// StripPKCS7 returns data unchanged when the padding is not well formed. BYD
// occasionally returns unpadded payloads, and pybyd tolerates them too.
func StripPKCS7(data []byte) []byte {
	if len(data) == 0 {
		return data
	}
	n := int(data[len(data)-1])
	if n == 0 || n > 16 || n > len(data) {
		return data
	}
	for _, b := range data[len(data)-n:] {
		if int(b) != n {
			return data
		}
	}
	return data[:len(data)-n]
}

// encryptCBC and decryptCBC use a zero IV, matching the app.

func encryptCBC(t *lookupTables, data []byte) []byte {
	res := make([]byte, len(data))
	prev := make([]byte, 16)
	buf := make([]byte, 16)
	for off := 0; off < len(data); off += 16 {
		for i := 0; i < 16; i++ {
			buf[i] = data[off+i] ^ prev[i]
		}
		enc := encryptBlock(t, buf)
		copy(res[off:], enc)
		copy(prev, enc)
	}
	return res
}

func decryptCBC(t *lookupTables, data []byte) []byte {
	res := make([]byte, len(data))
	prev := make([]byte, 16)
	for off := 0; off < len(data); off += 16 {
		block := data[off : off+16]
		dec := decryptBlock(t, block)
		for i := 0; i < 16; i++ {
			res[off+i] = dec[i] ^ prev[i]
		}
		copy(prev, block)
	}
	return res
}

type lookupTables struct {
	invRound, invXor, invFirst []byte
	round, xor, final          []byte
	permDecrypt, permEncrypt   []byte
}

var tableSpecs = []struct {
	name string
	size int
}{
	{"inv_round", 0x28000},
	{"inv_xor", 0x3C000},
	{"inv_first", 0x1000},
	{"round", 0x28000},
	{"xor", 0x3C000},
	{"final", 0x1000},
	{"perm_decrypt", 8},
	{"perm_encrypt", 8},
}

func mustLoadTables(data []byte) *lookupTables {
	t, err := loadTables(data)
	if err != nil {
		panic("bangcle: " + err.Error())
	}
	return t
}

// loadTables parses the BGTB container: magic, uint16 version, uint16 count,
// then count (offset, length) uint32 pairs, then the raw table bytes.
func loadTables(data []byte) (*lookupTables, error) {
	const header, indexEntry, count = 8, 8, 8
	if len(data) < header+count*indexEntry {
		return nil, fmt.Errorf("table file too short (%d bytes)", len(data))
	}
	if string(data[:4]) != "BGTB" {
		return nil, fmt.Errorf("bad magic %q", data[:4])
	}
	if v := binary.LittleEndian.Uint16(data[4:]); v != 1 {
		return nil, fmt.Errorf("unsupported table version %d", v)
	}
	if n := binary.LittleEndian.Uint16(data[6:]); n != count {
		return nil, fmt.Errorf("expected %d tables, got %d", count, n)
	}

	out := make([][]byte, count)
	for i := 0; i < count; i++ {
		p := header + i*indexEntry
		off := binary.LittleEndian.Uint32(data[p:])
		length := binary.LittleEndian.Uint32(data[p+4:])
		spec := tableSpecs[i]
		if int(length) != spec.size {
			return nil, fmt.Errorf("table %s: expected %d bytes, got %d", spec.name, spec.size, length)
		}
		if int(off)+int(length) > len(data) {
			return nil, fmt.Errorf("table %s: extends beyond file", spec.name)
		}
		out[i] = data[off : off+length]
	}
	return &lookupTables{out[0], out[1], out[2], out[3], out[4], out[5], out[6], out[7]}, nil
}