package main

import (
	"git.smesh.lol/smesh/web/common/crypto/aes256gcm"
	"crypto/sha256"
	"git.smesh.lol/smesh/web/common/helpers"
	"git.smesh.lol/smesh/web/common/jsbridge/ext"
	"git.smesh.lol/smesh/web/common/jsbridge/schnorr"
	"git.smesh.lol/smesh/web/common/jsbridge/subtle"
)

// Vault: Argon2id (v2) or PBKDF2 (v1, auto-migrated) + AES-256-GCM field
// encryption. Compatible with the Plebeian/Smesh signer format.
// Stored in browser.storage.local under vaultStorageKey.

const (
	vaultStorageKey = "smesh-vault"
	pbkdf2Salt      = "3e7cdebd-3b4c-4125-a18c-05750cad8ec3"
	pbkdf2Iters     = 1000
	argon2T         = 3
	// argon2MLegacy = original 64 MB. Existing vaults use this; new vaults use
	// argon2MNew. 64 MB allocates a 8MB block that can trip TinyGo GC stack
	// overflow on mobile WASM (limited linear memory, narrow GC stack).
	argon2MLegacy = 65536
	argon2MNew    = 4096 // 4 MB - mobile-safe, still ~4 GPU-min/guess
	argon2P       = 1
	argon2DKLen   = 32
)

type identity struct {
	Pubkey string // hex
	Seckey string // hex, plaintext in memory while unlocked
	Name   string
}

var (
	vaultKey     []byte
	vaultSalt    []byte // 32-byte Argon2id salt (v2+ only)
	vaultHash    string // hex SHA-256 of password, for fast unlock check
	vaultVersion int32    // 1, 2, or 3
	vaultArgon2M int32    // argon2 memory param this vault was encrypted with
	vaultOpen    bool
	vaultExists  bool
	vaultRawCache string // cached JSON from storage, set in loadVault callback

	// v2 compat: shared IV from stored vault, used only to decrypt v2 fields during migration.
	v2IV []byte

	identities []identity
	activeIdx  int32

	// lastUnlockErr is set by unlockVault on failure for diagnostic surface
	// via mgmtUnlockVault. Cleared on success. Values: "no-vault", "no-hash",
	// "wrong-password", "bad-iv", "kdf-failed", "decrypt-failed".
	lastUnlockErr string
)

func loadVault() {
	ext.StorageGet(vaultStorageKey, func(data string) {
		vaultRawCache = data
		vaultExists = data != ""
		if data == "" {
			ext.SignalReady()
			return
		}
		// Use heap-built key to avoid literal address relocation bug in wasm32.
		kVersion := string(append([]byte(nil), 'v', 'e', 'r', 's', 'i', 'o', 'n'))
		ver := helpers.JsonGetValue(data, kVersion)
		if len(ver) == 1 && ver[0] == '0' {
			restorePlaintextVault(data)
			ext.SignalReady()
			return
		}
		if ext.IsInPage() {
			tryRestoreSession()
		}
		ext.SignalReady()
	})
}

func activeIdentity() (p *identity) {
	if !vaultOpen || activeIdx < 0 || activeIdx >= len(identities) {
		return nil
	}
	return &identities[activeIdx]
}

func restoreActiveIdx(pk string) {
	if pk == "" {
		return
	}
	for i, id := range identities {
		if id.Pubkey == pk {
			activeIdx = i
			return
		}
	}
}

func createVault(password string) (ok bool) {
	hash := passwordHash(password)
	if hash == "" {
		return false
	}
	salt := []byte{:32}
	subtle.RandomBytes(salt)
	key := subtle.Argon2idDeriveKey(password, salt, argon2T, argon2MNew, argon2P, argon2DKLen)
	if len(key) == 0 {
		return false
	}
	vaultKey = key
	vaultSalt = salt
	vaultHash = hash
	vaultVersion = 3
	vaultArgon2M = argon2MNew
	vaultOpen = true
	identities = nil
	activeIdx = -1
	vaultExists = true
	saveVault()
	return true
}

func unlockVault(password string) (ok bool) {
	lastUnlockErr = ""
	data := vaultRawCache
	if data == "" {
		lastUnlockErr = "no-vault"
		return false
	}
	if len(data) > 0 && data[0] != '{' {
		return unlockLegacy(data, password)
	}
	kVaultHash := string(append([]byte(nil), 'v', 'a', 'u', 'l', 't', 'H', 'a', 's', 'h'))
	kIV := string(append([]byte(nil), 'i', 'v'))
	kSalt := string(append([]byte(nil), 's', 'a', 'l', 't'))
	kVersion := string(append([]byte(nil), 'v', 'e', 'r', 's', 'i', 'o', 'n'))
	kArgon2M := string(append([]byte(nil), 'a', 'r', 'g', 'o', 'n', '2', 'M'))
	storedHash := helpers.JsonGetString(data, kVaultHash)
	if storedHash == "" {
		lastUnlockErr = "no-hash"
		return false
	}
	computed := passwordHash(password)
	vaultHash = storedHash
	if computed != storedHash {
		// Include hash prefixes so the diagnostic distinguishes:
		// - typo (different prefixes, expected for a real typo)
		// - corrupted save (computed stable across attempts, stored differs)
		// - non-deterministic hash on WASM (computed differs across attempts)
		cp := computed
		if len(cp) > 12 {
			cp = cp[:12]
		}
		sp := storedHash
		if len(sp) > 12 {
			sp = sp[:12]
		}
		lastUnlockErr = "wrong-password computed=" | cp | " stored=" | sp | " pwlen=" | helpers.Itoa(int64(len(password)))
		return false
	}
	ver := helpers.JsonGetValue(data, kVersion)
	storedVersion := 2
	if len(ver) == 1 && ver[0] == '3' {
		storedVersion = 3
	}
	// v2 vaults have a shared IV; v3 has per-field IVs (no "iv" in JSON).
	var legacyIV []byte
	if storedVersion <= 2 {
		ivB64 := helpers.JsonGetString(data, kIV)
		legacyIV = helpers.Base64Decode(ivB64)
		if len(legacyIV) != 12 {
			lastUnlockErr = "bad-iv"
			return false
		}
	}
	// Read argon2M param. Missing field = legacy vault (m=65536).
	argMRaw := helpers.JsonGetValue(data, kArgon2M)
	argM := argon2MLegacy
	if argMRaw != "" {
		v := 0
		for i := 0; i < len(argMRaw); i++ {
			c := argMRaw[i]
			if c >= '0' && c <= '9' {
				v = v*10 + int32(c-'0')
			} else {
				v = 0
				break
			}
		}
		if v > 0 {
			argM = v
		}
	}
	saltB64 := helpers.JsonGetString(data, kSalt)
	if saltB64 != "" {
		salt := helpers.Base64Decode(saltB64)
		key := subtle.Argon2idDeriveKey(password, salt, argon2T, argM, argon2P, argon2DKLen)
		if len(key) == 0 {
			lastUnlockErr = "kdf-failed"
			return false
		}
		if !finishUnlock(data, key, legacyIV, salt, storedVersion) {
			lastUnlockErr = "decrypt-failed"
			return false
		}
		vaultArgon2M = argM
		return true
	}
	key := subtle.PBKDF2DeriveKey(password, []byte(pbkdf2Salt), pbkdf2Iters)
	if len(key) == 0 {
		lastUnlockErr = "kdf-failed"
		return false
	}
	if !finishUnlock(data, key, legacyIV, nil, 1) {
		lastUnlockErr = "decrypt-failed"
		return false
	}
	return true
}

func finishUnlock(data string, key, legacyIV, salt []byte, version int32) (ok bool) {
	vaultKey = key
	vaultSalt = salt
	vaultVersion = version
	v2IV = legacyIV
	vaultOpen = true
	// Restore vaultHash and vaultArgon2M from stored JSON so a subsequent
	// saveVault doesn't see len(vaultHash)==0 and wipe storage with "{}".
	// On a fresh page load, globals reset to zero and only data fields are
	// rehydrated by callers - these two were missed before.
	kVaultHashF := string(append([]byte(nil), 'v', 'a', 'u', 'l', 't', 'H', 'a', 's', 'h'))
	if h := helpers.JsonGetString(data, kVaultHashF); h != "" {
		vaultHash = h
	}
	kArgon2MF := string(append([]byte(nil), 'a', 'r', 'g', 'o', 'n', '2', 'M'))
	if vaultArgon2M == 0 {
		raw := helpers.JsonGetValue(data, kArgon2MF)
		v := 0
		for i := 0; i < len(raw); i++ {
			c := raw[i]
			if c < '0' || c > '9' { v = 0; break }
			v = v*10 + int32(c-'0')
		}
		if v > 0 {
			vaultArgon2M = v
		} else {
			vaultArgon2M = argon2MLegacy
		}
	}
	kIdentities := string(append([]byte(nil), 'i', 'd', 'e', 'n', 't', 'i', 't', 'i', 'e', 's'))
	idsRaw := helpers.JsonGetValue(data, kIdentities)
	identities = nil
	activeIdx = -1
	if idsRaw != "" {
		kPubkey := string(append([]byte(nil), 'p', 'u', 'b', 'k', 'e', 'y'))
		kSeckey := string(append([]byte(nil), 's', 'e', 'c', 'k', 'e', 'y'))
		kName := string(append([]byte(nil), 'n', 'a', 'm', 'e'))
		i := 0
		for i < len(idsRaw) && idsRaw[i] != '[' {
			i++
		}
		i++
		for i < len(idsRaw) {
			for i < len(idsRaw) && idsRaw[i] != '{' && idsRaw[i] != ']' {
				i++
			}
			if i >= len(idsRaw) || idsRaw[i] == ']' {
				break
			}
			end := i + 1
			depth := 1
			for end < len(idsRaw) && depth > 0 {
				if idsRaw[end] == '{' {
					depth++
				} else if idsRaw[end] == '}' {
					depth--
				} else if idsRaw[end] == '"' {
					end++
					for end < len(idsRaw) && idsRaw[end] != '"' {
						if idsRaw[end] == '\\' {
							end++
						}
						end++
					}
				}
				end++
			}
			obj := idsRaw[i:end]
			pk := helpers.JsonGetString(obj, kPubkey)
			encSk := helpers.JsonGetString(obj, kSeckey)
			name := helpers.JsonGetString(obj, kName)
			if pk != "" {
				sk := encSk
				decryptOK := true
				if version >= 2 && encSk != "" {
					if plain, ok := decryptField(encSk); ok {
						sk = plain
					} else {
						decryptOK = false
					}
				}
				if !decryptOK {
					// Bail. Caller surfaces this as decrypt-failed.
					lockVault()
					return false
				}
				identities = append(identities, identity{Pubkey: pk, Seckey: sk, Name: name})
			}
			i = end
		}
	}
	if len(identities) > 0 {
		activeIdx = 0
	}
	// Migrate v2 -> v3: re-encrypt with per-field IVs and save.
	if version <= 2 && len(identities) > 0 {
		vaultVersion = 3
		v2IV = nil
		saveVault()
	}
	return true
}

func lockVault() {
	for i := range vaultKey {
		vaultKey[i] = 0
	}
	vaultKey = nil
	vaultSalt = nil
	v2IV = nil
	vaultOpen = false
	identities = nil
}

func saveVault() {
	kVault := string(append([]byte(nil), 's', 'm', 'e', 's', 'h', '-', 'v', 'a', 'u', 'l', 't'))
	// Refuse to save if state is incomplete. Writing "{}" here would destroy
	// the on-disk vault when called from a refresh-then-modify flow where
	// vaultHash hadn't been rehydrated yet.
	if !vaultExists || len(vaultHash) == 0 {
		return
	}
	// Persist argon2M used at create time, so unlock uses the right param even
	// if defaults change later. Legacy vaults (no field) default to argon2MLegacy.
	mField := vaultArgon2M
	if mField == 0 {
		mField = argon2MLegacy
	}
	b := []byte{:512}[:0]
	b = append(b, '{', '"', 'v', 'e', 'r', 's', 'i', 'o', 'n', '"', ':')
	b = b | helpers.Itoa(int64(vaultVersion))
	b = append(b, ',', '"', 'a', 'r', 'g', 'o', 'n', '2', 'M', '"', ':')
	b = b | helpers.Itoa(int64(mField))
	b = append(b, ',', '"', 'v', 'a', 'u', 'l', 't', 'H', 'a', 's', 'h', '"', ':')
	b = b | helpers.JsonString(vaultHash)
	b = append(b, ',', '"', 's', 'a', 'l', 't', '"', ':')
	b = b | helpers.JsonString(helpers.Base64Encode(vaultSalt))
	b = append(b, ',', '"', 'i', 'd', 'e', 'n', 't', 'i', 't', 'i', 'e', 's', '"', ':', '[')
	for i, id := range identities {
		if i > 0 {
			b = append(b, ',')
		}
		encSk, ok := encryptField(id.Seckey)
		if !ok {
			encSk = id.Seckey
		}
		b = append(b, '{', '"', 'p', 'u', 'b', 'k', 'e', 'y', '"', ':')
		b = b | helpers.JsonString(id.Pubkey)
		b = append(b, ',', '"', 's', 'e', 'c', 'k', 'e', 'y', '"', ':')
		b = b | helpers.JsonString(encSk)
		b = append(b, ',', '"', 'n', 'a', 'm', 'e', '"', ':')
		b = b | helpers.JsonString(id.Name)
		b = append(b, '}')
	}
	b = append(b, ']', '}')
	vaultRawCache = string(b)
	ext.StorageSet(kVault, vaultRawCache)
	if vaultOpen && ext.IsInPage() {
		cacheSession()
	}
}

func restorePlaintextVault(data string) {
	kIdentities := string(append([]byte(nil), 'i', 'd', 'e', 'n', 't', 'i', 't', 'i', 'e', 's'))
	ids := helpers.JsonGetValue(data, kIdentities)
	if ids == "" {
		return
	}
	kPubkey := string(append([]byte(nil), 'p', 'u', 'b', 'k', 'e', 'y'))
	kSeckey := string(append([]byte(nil), 's', 'e', 'c', 'k', 'e', 'y'))
	pk := helpers.JsonGetString(ids, kPubkey)
	sk := helpers.JsonGetString(ids, kSeckey)
	if pk == "" || sk == "" {
		return
	}
	vaultOpen = true
	vaultExists = true
	vaultVersion = 0
	identities = []identity{{Pubkey: pk, Seckey: sk, Name: ""}}
	activeIdx = 0
}

const sessionCacheKey = "vault-session"

func cacheSession() {
	if !ext.IsInPage() || !vaultOpen || len(vaultKey) == 0 {
		return
	}
	b := []byte(nil) | helpers.HexEncode(vaultKey)
	b = append(b, ':')
	b = b | helpers.HexEncode(vaultSalt)
	b = append(b, ':')
	b = b | helpers.Itoa(int64(vaultVersion))
	b = append(b, ':')
	aid := activeIdentity()
	if aid != nil {
		b = b | aid.Pubkey
	}
	kSession := string(append([]byte(nil), 'v', 'a', 'u', 'l', 't', '-', 's', 'e', 's', 's', 'i', 'o', 'n'))
	ext.SessionSet(kSession, string(b))
}

func clearSessionCache() {
	kSession := string(append([]byte(nil), 'v', 'a', 'u', 'l', 't', '-', 's', 'e', 's', 's', 'i', 'o', 'n'))
	ext.SessionSet(kSession, "")
}

func tryRestoreSession() {
	kSession := string(append([]byte(nil), 'v', 'a', 'u', 'l', 't', '-', 's', 'e', 's', 's', 'i', 'o', 'n'))
	ext.SessionGet(kSession, func(val string) {
		if val == "" {
			return
		}
		var parts [6]string
		n := 0
		start := 0
		for i := 0; i < len(val) && n < 5; i++ {
			if val[i] == ':' {
				parts[n] = val[start:i]
				n++
				start = i + 1
			}
		}
		parts[n] = val[start:]
		n++
		data := vaultRawCache
		if data == "" {
			clearSessionCache()
			return
		}
		// Detect format by part count and whether parts[2] is a 1-char
		// version digit (v3) or a long hex salt (v2).
		// v3: key:salt:version[:activePubkey] - parts[2] is 1 char
		// v2: key:iv:salt:version[:activePubkey] - parts[2] is hex salt
		isV3 := n >= 3 && len(parts[2]) == 1
		if isV3 {
			key := helpers.HexDecode(parts[0])
			salt := helpers.HexDecode(parts[1])
			if len(key) != 32 {
				clearSessionCache()
				return
			}
			ver := parseSessionVer(parts[2])
			if !finishUnlock(data, key, nil, salt, ver) {
				clearSessionCache()
				return
			}
			if n >= 4 {
				restoreActiveIdx(parts[3])
			}
		} else if n >= 4 {
			key := helpers.HexDecode(parts[0])
			iv := helpers.HexDecode(parts[1])
			salt := helpers.HexDecode(parts[2])
			if len(key) != 32 || len(iv) != 12 {
				clearSessionCache()
				return
			}
			ver := parseSessionVer(parts[3])
			if !finishUnlock(data, key, iv, salt, ver) {
				clearSessionCache()
				return
			}
			if n >= 5 {
				restoreActiveIdx(parts[4])
			}
		} else {
			clearSessionCache()
		}
	})
}

func parseSessionVer(s string) (n int32) {
	if len(s) == 1 {
		if s[0] == '1' {
			return 1
		}
		if s[0] == '2' {
			return 2
		}
	}
	return 3
}

// encryptFieldWith encrypts using explicit key (for export with different password).
// Generates a fresh 12-byte random IV per call, prepended to ciphertext.
// Output: base64(iv[12] || ciphertext || tag[16]).
func encryptFieldWith(plain string, key []byte) (string, bool) {
	if len(key) == 0 {
		return "", false
	}
	iv := []byte{:12}
	subtle.RandomBytes(iv)
	ct := aes256gcm.Seal(key, iv, []byte(plain), nil)
	if ct == nil {
		return "", false
	}
	out := []byte{:0:12 + len(ct)}
	out = out | iv
	out = out | ct
	return helpers.Base64Encode(out), true
}

// encryptField encrypts a plaintext string using vaultKey via AES-256-GCM.
// Fresh random IV per call - no nonce reuse.
func encryptField(plain string) (string, bool) {
	if !vaultOpen || len(vaultKey) == 0 {
		return "", false
	}
	return encryptFieldWith(plain, vaultKey)
}

// decryptField decrypts a base64-encoded AES-256-GCM ciphertext.
// v3 format: base64(iv[12] || ciphertext || tag[16]).
// v2 compat: if v2IV is set, tries old format (no prepended IV) as fallback.
func decryptField(enc string) (string, bool) {
	if !vaultOpen || len(vaultKey) == 0 {
		return "", false
	}
	raw := helpers.Base64Decode(enc)
	if raw == nil {
		return "", false
	}
	// v3: first 12 bytes are per-field IV.
	if len(raw) >= 12+16 {
		iv := raw[:12]
		ct := raw[12:]
		pt, ok := aes256gcm.Open(vaultKey, iv, ct, nil)
		if ok {
			return string(pt), true
		}
	}
	// v2 fallback: shared IV from vault JSON.
	if len(v2IV) == 12 {
		pt, ok := aes256gcm.Open(vaultKey, v2IV, raw, nil)
		if ok {
			return string(pt), true
		}
	}
	return "", false
}

func passwordHash(pw string) (s string) {
	return subtle.SHA256Hex([]byte(pw))
}

func unlockLegacy(data, password string) (ok bool) {
	sepIdx := -1
	for i := 0; i < len(data); i++ {
		if data[i] == ':' {
			sepIdx = i
			break
		}
	}
	if sepIdx < 1 {
		return false
	}
	iv := helpers.HexDecode(data[:sepIdx])
	ct := helpers.HexDecode(data[sepIdx+1:])
	if iv == nil || ct == nil {
		return false
	}
	key := legacyDeriveKey(password)
	pt := subtle.AESCBCDecrypt(key[:], iv, ct)
	if len(pt) == 0 {
		return false
	}
	parseLegacyIdentities(string(pt))
	newSalt := []byte{:32}
	subtle.RandomBytes(newSalt)
	hash := passwordHash(password)
	if hash == "" {
		return true
	}
	newKey := subtle.Argon2idDeriveKey(password, newSalt, argon2T, argon2MNew, argon2P, argon2DKLen)
	if len(newKey) == 0 {
		return true
	}
	vaultKey = newKey
	vaultSalt = newSalt
	vaultHash = hash
	vaultVersion = 3
	vaultArgon2M = argon2MNew
	vaultOpen = true
	vaultExists = true
	saveVault()
	return true
}

// legacyDeriveKey derives a 32-byte key via 100000 rounds of SHA-256 over "smesh-vault-salt:" | password.
func legacyDeriveKey(password string) (v [32]byte) {
	saltPrefix := append([]byte(nil), 's', 'm', 'e', 's', 'h', '-', 'v', 'a', 'u', 'l', 't', '-', 's', 'a', 'l', 't', ':')
	saltPrefix = saltPrefix | password
	h := sha256.Sum(saltPrefix)
	for i := 0; i < 100000; i++ {
		h = sha256.Sum(h[:])
	}
	return h
}

// parseLegacyIdentities parses plaintext JSON array of identities from old format.
func parseLegacyIdentities(s string) {
	identities = nil
	activeIdx = -1
	i := 0
	for i < len(s) && s[i] != '[' {
		i++
	}
	i++
	for i < len(s) {
		for i < len(s) && s[i] != '{' && s[i] != ']' {
			i++
		}
		if i >= len(s) || s[i] == ']' {
			break
		}
		end := i + 1
		depth := 1
		for end < len(s) && depth > 0 {
			if s[end] == '{' {
				depth++
			} else if s[end] == '}' {
				depth--
			} else if s[end] == '"' {
				end++
				for end < len(s) && s[end] != '"' {
					if s[end] == '\\' {
						end++
					}
					end++
				}
			}
			end++
		}
		obj := s[i:end]
		kPubkey := string(append([]byte(nil), 'p', 'u', 'b', 'k', 'e', 'y'))
		kPublicKey := string(append([]byte(nil), 'p', 'u', 'b', 'l', 'i', 'c', 'K', 'e', 'y'))
		kPub := string(append([]byte(nil), 'p', 'u', 'b'))
		kSeckey := string(append([]byte(nil), 's', 'e', 'c', 'k', 'e', 'y'))
		kPrivkey := string(append([]byte(nil), 'p', 'r', 'i', 'v', 'k', 'e', 'y'))
		kSecretKey := string(append([]byte(nil), 's', 'e', 'c', 'r', 'e', 't', 'K', 'e', 'y'))
		kSec := string(append([]byte(nil), 's', 'e', 'c'))
		kKey := string(append([]byte(nil), 'k', 'e', 'y'))
		kName := string(append([]byte(nil), 'n', 'a', 'm', 'e'))
		kNick := string(append([]byte(nil), 'n', 'i', 'c', 'k'))
		kLabel := string(append([]byte(nil), 'l', 'a', 'b', 'e', 'l'))
		pk := helpers.JsonGetString(obj, kPubkey)
		if pk == "" {
			pk = helpers.JsonGetString(obj, kPublicKey)
		}
		if pk == "" {
			pk = helpers.JsonGetString(obj, kPub)
		}
		sk := helpers.JsonGetString(obj, kSeckey)
		if sk == "" {
			sk = helpers.JsonGetString(obj, kPrivkey)
		}
		if sk == "" {
			sk = helpers.JsonGetString(obj, kSecretKey)
		}
		if sk == "" {
			sk = helpers.JsonGetString(obj, kSec)
		}
		if sk == "" {
			sk = helpers.JsonGetString(obj, kKey)
		}
		nm := helpers.JsonGetString(obj, kName)
		if nm == "" {
			nm = helpers.JsonGetString(obj, kNick)
		}
		if nm == "" {
			nm = helpers.JsonGetString(obj, kLabel)
		}
		// nsec decode
		kNsec1 := string(append([]byte(nil), 'n', 's', 'e', 'c', '1'))
		if len(sk) >= 5 && sk[:5] == kNsec1 {
			skBytes := helpers.DecodeNsec(sk)
			if skBytes != nil {
				sk = helpers.HexEncode(skBytes)
			}
		}
		// derive pubkey if absent
		if pk == "" && sk != "" {
			skBytes := helpers.HexDecode(sk)
			if skBytes != nil {
				dpk, ok := schnorr.PubKeyFromSecKey(skBytes)
				if ok {
					pk = helpers.HexEncode(dpk)
				}
			}
		}
		if pk != "" && sk != "" {
			identities = append(identities, identity{Pubkey: pk, Seckey: sk, Name: nm})
		}
		i = end
	}
	if len(identities) > 0 {
		activeIdx = 0
	}
}