package main

import (
	"smesh.lol/web/common/helpers"
	"smesh.lol/web/common/jsbridge/ext"
	"smesh.lol/web/common/jsbridge/schnorr"
	"smesh.lol/web/common/jsbridge/subtle"
)

// HD keychain state (populated when an HD vault is unlocked).
var (
	hdMnemonic    string // decrypted mnemonic phrase
	hdNextAccount int    // next account index to derive
)

const hardenedBit = 0x80000000

// --- BIP39 mnemonic generation ---

// generateMnemonic creates a 12-word BIP39 mnemonic from 128 bits of entropy.
func generateMnemonic() string {
	entropy := []byte{:16}
	subtle.RandomBytes(entropy)
	return entropyToMnemonic(entropy)
}

// entropyToMnemonic converts 16 bytes of entropy to a 12-word mnemonic.
func entropyToMnemonic(entropy []byte) string {
	// SHA-256 checksum: first 4 bits (128/32 = 4 bits).
	hash := schnorr.SHA256Sum(entropy)

	// 17 bytes: 16 entropy + 1 checksum (only top 4 bits used).
	all := []byte{:17}
	copy(all, entropy)
	all[16] = hash[0]

	// Extract 12 groups of 11 bits from the 132-bit big-endian stream.
	var words string
	for i := 0; i < 12; i++ {
		bitPos := i * 11
		byteIdx := bitPos / 8
		bitOff := uint(bitPos % 8)

		// Pull up to 3 bytes to cover the 11-bit window.
		var val uint32
		val = uint32(all[byteIdx]) << 16
		if byteIdx+1 < 17 {
			val |= uint32(all[byteIdx+1]) << 8
		}
		if byteIdx+2 < 17 {
			val |= uint32(all[byteIdx+2])
		}

		// Shift left by bitOff, then take top 11 bits of the 24-bit window.
		idx := (val << bitOff >> 13) & 0x7FF

		if i > 0 {
			words += " "
		}
		words += bip39Words[idx]
	}
	return words
}

// validateMnemonic checks if a mnemonic is valid BIP39 (12 or 24 words).
func validateMnemonic(mnemonic string) bool {
	words := splitWords(mnemonic)
	if len(words) != 12 && len(words) != 24 {
		return false
	}

	// Convert words to 11-bit indices.
	var indices []int
	for _, w := range words {
		idx := wordIndex(w)
		if idx < 0 {
			return false
		}
		indices = append(indices, idx)
	}

	// Reconstruct entropy + checksum from the bit stream.
	totalBits := len(words) * 11
	checksumBits := totalBits / 33 // 4 for 12 words, 8 for 24
	entropyBits := totalBits - checksumBits
	entropyBytes := entropyBits / 8

	entropy := []byte{:entropyBytes}
	csFromMnemonic := 0
	pos := 0
	for _, idx := range indices {
		for b := 10; b >= 0; b-- {
			bit := (idx >> uint(b)) & 1
			if pos < entropyBits {
				entropy[pos/8] |= byte(bit) << uint(7-(pos%8))
			} else {
				csFromMnemonic = (csFromMnemonic << 1) | bit
			}
			pos++
		}
	}

	// Verify checksum against SHA-256 of entropy.
	hash := schnorr.SHA256Sum(entropy)
	csFromHash := int(hash[0]) >> uint(8-checksumBits)

	return csFromMnemonic == csFromHash
}

// wordIndex returns the BIP39 word list index, or -1 if not found.
func wordIndex(word string) int {
	for i := 0; i < 2048; i++ {
		if bip39Words[i] == word {
			return i
		}
	}
	return -1
}

// splitWords splits a space-separated string into words.
func splitWords(s string) []string {
	var words []string
	start := -1
	for i := 0; i < len(s); i++ {
		if s[i] == ' ' || s[i] == '\t' || s[i] == '\n' {
			if start >= 0 {
				words = append(words, s[start:i])
				start = -1
			}
		} else if start < 0 {
			start = i
		}
	}
	if start >= 0 {
		words = append(words, s[start:])
	}
	return words
}

// --- BIP39 seed derivation ---

// mnemonicToSeed converts a mnemonic to a 64-byte seed via PBKDF2-HMAC-SHA512.
func mnemonicToSeed(mnemonic, passphrase string, fn func([]byte)) {
	salt := []byte("mnemonic" + passphrase)
	subtle.PBKDF2SHA512(mnemonic, salt, 2048, 64, fn)
}

// --- BIP32 HD key derivation ---

// bip32MasterKey derives the master key and chain code from a seed.
// HMAC-SHA512 with key="Bitcoin seed".
func bip32MasterKey(seed []byte, fn func(key, chain []byte)) {
	subtle.HMACSHA512([]byte("Bitcoin seed"), seed, func(result []byte) {
		if len(result) != 64 {
			fn(nil, nil)
			return
		}
		key := []byte{:32}
		chain := []byte{:32}
		copy(key, result[:32])
		copy(chain, result[32:])
		fn(key, chain)
	})
}

// bip32DeriveChild derives a child key at the given index.
// For hardened derivation, index must include hardenedBit.
func bip32DeriveChild(key, chain []byte, index uint32, fn func(childKey, childChain []byte)) {
	var data []byte
	if index >= hardenedBit {
		// Hardened: 0x00 || key || ser32(index)
		data = []byte{:37}
		data[0] = 0x00
		copy(data[1:33], key)
	} else {
		// Normal: compressedPubkey(key) || ser32(index)
		compressed, ok := schnorr.CompressedPubKey(key)
		if !ok {
			fn(nil, nil)
			return
		}
		data = []byte{:37}
		copy(data[:33], compressed)
	}
	data[33] = byte(index >> 24)
	data[34] = byte(index >> 16)
	data[35] = byte(index >> 8)
	data[36] = byte(index)

	subtle.HMACSHA512(chain, data, func(result []byte) {
		if len(result) != 64 {
			fn(nil, nil)
			return
		}
		childKey, ok := schnorr.ScalarAddModN(result[:32], key)
		if !ok {
			fn(nil, nil)
			return
		}
		childChain := []byte{:32}
		copy(childChain, result[32:])
		fn(childKey, childChain)
	})
}

// --- NIP-06 derivation: m/44'/1237'/account'/0/0 ---

func deriveNIP06(seed []byte, account int, fn func(seckey []byte)) {
	bip32MasterKey(seed, func(mk, mc []byte) {
		if mk == nil {
			fn(nil)
			return
		}
		// m/44'
		bip32DeriveChild(mk, mc, 44|hardenedBit, func(k1, c1 []byte) {
			if k1 == nil {
				fn(nil)
				return
			}
			// m/44'/1237'
			bip32DeriveChild(k1, c1, 1237|hardenedBit, func(k2, c2 []byte) {
				if k2 == nil {
					fn(nil)
					return
				}
				// m/44'/1237'/account'
				bip32DeriveChild(k2, c2, uint32(account)|hardenedBit, func(k3, c3 []byte) {
					if k3 == nil {
						fn(nil)
						return
					}
					// m/44'/1237'/account'/0
					bip32DeriveChild(k3, c3, 0, func(k4, c4 []byte) {
						if k4 == nil {
							fn(nil)
							return
						}
						// m/44'/1237'/account'/0/0
						bip32DeriveChild(k4, c4, 0, func(k5, _ []byte) {
							fn(k5)
						})
					})
				})
			})
		})
	})
}

// --- HD vault operations ---

// hdCreateVault generates a mnemonic, creates the vault, and derives the first identity.
func hdCreateVault(password, name string, done func(mnemonic string)) {
	mnemonic := generateMnemonic()
	hdRestoreVault(password, mnemonic, name, func(ok bool) {
		if ok {
			done(mnemonic)
		} else {
			done("")
		}
	})
}

// hdRestoreVault creates a vault from an existing mnemonic and derives the first identity.
func hdRestoreVault(password, mnemonic, name string, done func(bool)) {
	if !validateMnemonic(mnemonic) {
		done(false)
		return
	}

	createVault(password, func(ok bool) {
		if !ok {
			done(false)
			return
		}
		hdMnemonic = mnemonic
		hdNextAccount = 0
		hdDeriveNext(name, func(pubkey string) {
			if pubkey == "" {
				done(false)
				return
			}
			saveHDVault(func() {
				done(true)
			})
		})
	})
}

// hdDeriveNext derives the next identity from the HD mnemonic.
func hdDeriveNext(name string, done func(pubkey string)) {
	if hdMnemonic == "" {
		done("")
		return
	}

	account := hdNextAccount
	mnemonicToSeed(hdMnemonic, "", func(seed []byte) {
		if len(seed) == 0 {
			done("")
			return
		}
		deriveNIP06(seed, account, func(sk []byte) {
			if sk == nil {
				done("")
				return
			}
			pk, ok := schnorr.PubKeyFromSecKey(sk)
			if !ok {
				done("")
				return
			}
			pkHex := helpers.HexEncode(pk)
			skHex := helpers.HexEncode(sk)

			// Check for duplicate.
			for _, id := range identities {
				if id.Pubkey == pkHex {
					done("")
					return
				}
			}

			identities = append(identities, identity{Pubkey: pkHex, Seckey: skHex, Name: name})
			if activeIdx < 0 {
				activeIdx = 0
			}
			hdNextAccount = account + 1
			saveHDVault(func() {
				done(pkHex)
			})
		})
	})
}

// --- HD vault persistence ---

// saveHDVault saves the vault with the HD mnemonic field.
func saveHDVault(done func()) {
	if !vaultOpen {
		if done != nil {
			done()
		}
		return
	}

	if hdMnemonic == "" {
		// Non-HD vault, use normal save.
		saveVault(done)
		return
	}

	// Encrypt mnemonic, then build vault JSON with HD fields.
	encryptField(hdMnemonic, vaultKey, vaultIV, func(encMnemonic string) {
		encryptAllIdentities(vaultKey, vaultIV, func(idJSON string) {
			s := "{\"version\":" + itoa(vaultVersion) +
				",\"iv\":" + helpers.JsonString(helpers.Base64Encode(vaultIV)) +
				",\"vaultHash\":" + helpers.JsonString(vaultHash)

			if vaultVersion >= 2 && vaultSalt != nil {
				s += ",\"salt\":" + helpers.JsonString(helpers.Base64Encode(vaultSalt))
			}

			s += ",\"mnemonic\":" + helpers.JsonString(encMnemonic)
			s += ",\"hdNextAccount\":" + itoa(hdNextAccount)
			s += ",\"identities\":" + idJSON
			s += ",\"permissions\":[]"
			s += ",\"relays\":[]"
			s += ",\"selectedIdentityId\":null"
			s += "}"

			vaultRawCache = s
			ext.StorageSet(vaultStorageKey, s)
			if done != nil {
				done()
			}
		})
	})
}

// loadHDFields decrypts the HD-specific fields after vault unlock.
// Called from finishUnlock when mnemonic field is present.
func loadHDFields(data string, key, iv []byte, done func()) {
	encMnemonic := helpers.JsonGetString(data, "mnemonic")
	if encMnemonic == "" {
		hdMnemonic = ""
		hdNextAccount = 0
		done()
		return
	}
	decryptField(encMnemonic, key, iv, func(m string) {
		hdMnemonic = m
		// Parse hdNextAccount.
		naStr := helpers.JsonGetValue(data, "hdNextAccount")
		hdNextAccount = parseSimpleInt(naStr)
		done()
	})
}

// parseSimpleInt parses a non-negative integer from a string.
func parseSimpleInt(s string) int {
	n := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		if c >= '0' && c <= '9' {
			n = n*10 + int(c-'0')
		} else {
			break
		}
	}
	return n
}

// probeHDAccount derives the key at a given account index without adding to vault.
// Returns the hex pubkey, or empty string on failure.
func probeHDAccount(account int, fn func(string)) {
	if hdMnemonic == "" {
		fn("")
		return
	}
	mnemonicToSeed(hdMnemonic, "", func(seed []byte) {
		if len(seed) == 0 {
			fn("")
			return
		}
		deriveNIP06(seed, account, func(sk []byte) {
			if sk == nil {
				fn("")
				return
			}
			pk, ok := schnorr.PubKeyFromSecKey(sk)
			if !ok {
				fn("")
				return
			}
			fn(helpers.HexEncode(pk))
		})
	})
}

// lockHD clears HD state.
func lockHD() {
	hdMnemonic = ""
	hdNextAccount = 0
}

// exportHDFields encrypts mnemonic with the given key/iv for vault export.
// Returns extra JSON fields (with leading comma) or empty string if not HD.
func exportHDFields(key, iv []byte, fn func(string)) {
	if hdMnemonic == "" {
		fn("")
		return
	}
	encryptField(hdMnemonic, key, iv, func(encMnemonic string) {
		fn(",\"mnemonic\":" + helpers.JsonString(encMnemonic) +
			",\"hdNextAccount\":" + itoa(hdNextAccount))
	})
}