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"
)

// Vault: compatible with Plebeian/Smesh signer format.
// V2: Argon2id + AES-256-GCM, per-field encryption, JSON envelope.
// V1: PBKDF2 + AES-256-GCM, same structure (auto-migrated on unlock).
//
// Stored in browser.storage.local as JSON:
// {
//   "version": 2,
//   "iv": "base64(12 bytes)",
//   "salt": "base64(32 bytes)",       // v2 only
//   "vaultHash": "hex(SHA-256(pw))",
//   "identities": [{id,nick,createdAt,privkey}],  // each field AES-GCM encrypted base64
//   "selectedIdentityId": "encrypted" | null,
//   "permissions": [],
//   "relays": []
// }

const (
	vaultStorageKey = "smesh-vault"
	pbkdf2Salt      = "3e7cdebd-3b4c-4125-a18c-05750cad8ec3"
	pbkdf2Iters     = 1000
	argon2T         = 8
	argon2M         = 262144 // 256 MB
	argon2P         = 4
	argon2DKLen     = 32
)

type identity struct {
	Pubkey string // hex (derived from Seckey)
	Seckey string // hex (plaintext in memory when unlocked)
	Name   string
}

type encryptedID struct {
	privkey string
	nick    string
}

var (
	vaultKey      []byte // 32-byte derived key, nil when locked
	vaultIV       []byte // 12-byte shared IV
	vaultSalt     []byte // 32-byte salt (v2 only, nil for v1)
	vaultHash     string // hex SHA-256 of password
	vaultVersion  int    // 1 or 2
	vaultOpen     bool
	identities    []identity
	activeIdx     int
	vaultExists   bool
	vaultRawCache string // cached JSON from storage
)

func log(msg string) {
	ext.ConsoleLog(msg)
}

func loadVault() {
	ext.StorageGet(vaultStorageKey, func(data string) {
		vaultRawCache = data
		vaultExists = data != ""
		if data == "" {
			log("loadVault: no vault in storage")
		} else if len(data) > 40 {
			log("loadVault: loaded, first 40 chars: " + data[:40])
		} else {
			log("loadVault: loaded, data: " + data)
		}
	})
}

// passwordHash computes lowercase hex SHA-256 of the password via native crypto.subtle.
func passwordHash(pw string, fn func(string)) {
	subtle.SHA256Hex([]byte(pw), fn)
}

// createVault creates a new v2 vault with Argon2id.
func createVault(password string, done func(bool)) {
	passwordHash(password, func(hash string) {
		if hash == "" {
			done(false)
			return
		}

		// Generate random salt (32 bytes) and IV (12 bytes).
		salt := []byte{:32}
		iv := []byte{:12}
		subtle.RandomBytes(salt)
		subtle.RandomBytes(iv)

		subtle.Argon2idDeriveKey(password, salt, argon2T, argon2M, argon2P, argon2DKLen, func(key []byte) {
			if len(key) == 0 {
				done(false)
				return
			}
			vaultKey = key
			vaultIV = iv
			vaultSalt = salt
			vaultHash = hash
			vaultVersion = 2
			vaultOpen = true
			identities = nil
			activeIdx = -1
			vaultExists = true
			saveVault(func() {
				done(true)
			})
		})
	})
}

// unlockVault unlocks an existing vault.
func unlockVault(password string, done func(bool)) {
	data := vaultRawCache
	log("unlockVault: data len=" + itoa(len(data)))
	if data == "" {
		log("unlockVault: no data")
		done(false)
		return
	}

	// Detect format: legacy "hex:hex" (AES-CBC) vs JSON envelope.
	if len(data) > 0 && data[0] != '{' {
		log("unlockVault: detected legacy format")
		unlockLegacy(data, password, done)
		return
	}

	// Quick password check via vaultHash (async — uses native crypto.subtle).
	log("unlockVault: JSON format detected")
	storedHash := helpers.JsonGetString(data, "vaultHash")
	if storedHash == "" {
		log("unlockVault: no vaultHash in data")
		done(false)
		return
	}
	passwordHash(password, func(computed string) {
		log("unlockVault: stored=" + storedHash + " computed=" + computed)
		if computed != storedHash {
			log("unlockVault: hash mismatch")
			done(false)
			return
		}
		log("unlockVault: hash OK")

		// Parse version and IV.
		ivB64 := helpers.JsonGetString(data, "iv")
		iv := helpers.Base64Decode(ivB64)
		if len(iv) != 12 {
			log("unlockVault: bad IV length")
			done(false)
			return
		}

		// Detect v2 by presence of salt.
		saltB64 := helpers.JsonGetString(data, "salt")
		isV2 := saltB64 != ""

		if isV2 {
			salt := helpers.Base64Decode(saltB64)
			if len(salt) == 0 {
				done(false)
				return
			}
			log("unlockVault: v2, deriving key with Argon2id...")
			subtle.Argon2idDeriveKey(password, salt, argon2T, argon2M, argon2P, argon2DKLen, func(key []byte) {
				if len(key) == 0 {
					log("unlockVault: Argon2id failed")
					done(false)
					return
				}
				log("unlockVault: key derived, decrypting...")
				finishUnlock(data, key, iv, salt, 2, password, done)
			})
		} else {
			log("unlockVault: v1, deriving key with PBKDF2...")
			subtle.PBKDF2DeriveKey(password, []byte(pbkdf2Salt), pbkdf2Iters, func(key []byte) {
				if len(key) == 0 {
					log("unlockVault: PBKDF2 failed")
					done(false)
					return
				}
				log("unlockVault: key derived, decrypting...")
				finishUnlock(data, key, iv, nil, 1, password, done)
			})
		}
	})
}

// unlockLegacy handles the old "hex(IV):hex(ciphertext)" AES-CBC format.
// Derives key via iterated SHA-256 and migrates to the new format on success.
func unlockLegacy(data, password string, done func(bool)) {
	sepIdx := -1
	for i := 0; i < len(data); i++ {
		if data[i] == ':' {
			sepIdx = i
			break
		}
	}
	if sepIdx < 1 {
		log("unlockLegacy: no separator found")
		done(false)
		return
	}
	log("unlockLegacy: sep at " + itoa(sepIdx) + ", iv hex len=" + itoa(sepIdx) + ", ct hex len=" + itoa(len(data)-sepIdx-1))
	iv := helpers.HexDecode(data[:sepIdx])
	ct := helpers.HexDecode(data[sepIdx+1:])
	if iv == nil {
		log("unlockLegacy: iv hex decode failed")
		done(false)
		return
	}
	if ct == nil {
		log("unlockLegacy: ct hex decode failed")
		done(false)
		return
	}
	log("unlockLegacy: iv len=" + itoa(len(iv)) + ", ct len=" + itoa(len(ct)))

	// Old key derivation: iterated SHA-256.
	log("unlockLegacy: deriving key (100k SHA-256 iterations)...")
	key := legacyDeriveKey(password)
	log("unlockLegacy: key derived, first 4 bytes: " + helpers.HexEncode(key[:4]))
	subtle.AESCBCDecrypt(key[:], iv, ct, func(pt []byte) {
		log("unlockLegacy: AESCBCDecrypt returned " + itoa(len(pt)) + " bytes")
		if len(pt) == 0 {
			log("unlockLegacy: decryption failed (empty result)")
			done(false)
			return
		}
		if len(pt) > 60 {
			log("unlockLegacy: plaintext first 60: " + string(pt[:60]))
		} else {
			log("unlockLegacy: plaintext: " + string(pt))
		}
		vaultOpen = true
		parseLegacyIdentities(string(pt))
		log("unlockLegacy: parsed " + itoa(len(identities)) + " identities, migrating...")

		// Migrate to new format.
		migrateV1ToV2(password, func(ok bool) {
			log("unlockLegacy: migration done, ok=" + boolStr(ok))
			done(true)
		})
	})
}

// legacyDeriveKey derives a 32-byte key using the old iterated SHA-256 scheme.
func legacyDeriveKey(password string) [32]byte {
	hSlice := schnorr.SHA256Sum([]byte("smesh-vault-salt:" + password))
	var h [32]byte
	copy(h[:], hSlice)
	for i := 0; i < 100000; i++ {
		hSlice = schnorr.SHA256Sum(h[:])
		copy(h[:], hSlice)
	}
	return h
}

// parseLegacyIdentities parses the old plaintext format: [{pubkey,seckey,name},...]
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]
		pk := helpers.JsonGetString(obj, "pubkey")
		sk := helpers.JsonGetString(obj, "seckey")
		nm := helpers.JsonGetString(obj, "name")
		if pk != "" && sk != "" {
			identities = append(identities, identity{Pubkey: pk, Seckey: sk, Name: nm})
		}
		i = end
	}
	if len(identities) > 0 {
		activeIdx = 0
	}
}

// finishUnlock decrypts vault fields and populates memory state.
// If v1, triggers auto-migration to v2.
func finishUnlock(data string, key, iv, salt []byte, version int, password string, done func(bool)) {
	vaultKey = key
	vaultIV = iv
	vaultSalt = salt
	vaultHash = helpers.JsonGetString(data, "vaultHash")
	vaultVersion = version
	vaultOpen = true
	identities = nil
	activeIdx = -1

	// Parse and decrypt identities.
	idList := helpers.JsonGetValue(data, "identities")
	decryptIdentities(idList, key, iv, func(ok bool) {
		if !ok {
			vaultOpen = false
			vaultKey = nil
			done(false)
			return
		}

		// Load HD fields (mnemonic, hdNextAccount) if present.
		loadHDFields(data, key, iv, func() {
			// Decrypt selectedIdentityId if present.
			selEnc := helpers.JsonGetString(data, "selectedIdentityId")
			if selEnc != "" {
				decryptField(selEnc, key, iv, func(selID string) {
					selectIdentityByID(selID)
					if version == 1 {
						migrateV1ToV2(password, done)
					} else {
						done(true)
					}
				})
			} else {
				if len(identities) > 0 {
					activeIdx = 0
				}
				if version == 1 {
					migrateV1ToV2(password, done)
				} else {
					done(true)
				}
			}
		})
	})
}

func selectIdentityByID(id string) {
	_ = id
	if len(identities) > 0 && activeIdx < 0 {
		activeIdx = 0
	}
}

// migrateV1ToV2 re-encrypts vault from PBKDF2 to Argon2id.
func migrateV1ToV2(password string, done func(bool)) {
	salt := []byte{:32}
	iv := []byte{:12}
	subtle.RandomBytes(salt)
	subtle.RandomBytes(iv)

	subtle.Argon2idDeriveKey(password, salt, argon2T, argon2M, argon2P, argon2DKLen, func(key []byte) {
		if len(key) == 0 {
			// Migration failed, but v1 unlock succeeded.
			done(true)
			return
		}
		vaultKey = key
		vaultIV = iv
		vaultSalt = salt
		vaultVersion = 2
		saveVault(func() {
			done(true)
		})
	})
}

func lockVault() {
	vaultKey = nil
	vaultIV = nil
	vaultSalt = nil
	vaultOpen = false
	identities = nil
	activeIdx = -1
	lockHD()
}

// --- Per-field AES-GCM encryption ---

func encryptField(plaintext string, key, iv []byte, fn func(string)) {
	if plaintext == "" {
		fn("")
		return
	}
	subtle.AESGCMEncrypt(key, iv, []byte(plaintext), func(ct []byte) {
		if len(ct) == 0 {
			fn("")
			return
		}
		fn(helpers.Base64Encode(ct))
	})
}

func decryptField(b64 string, key, iv []byte, fn func(string)) {
	if b64 == "" {
		fn("")
		return
	}
	ct := helpers.Base64Decode(b64)
	if ct == nil {
		fn("")
		return
	}
	subtle.AESGCMDecrypt(key, iv, ct, func(pt []byte) {
		fn(string(pt))
	})
}

// --- Identity encryption/decryption ---

// decryptIdentities parses the JSON identity array and decrypts each field.
// Calls done(true) when all identities are decrypted.
func decryptIdentities(listJSON string, key, iv []byte, done func(bool)) {
	if listJSON == "" || listJSON == "[]" {
		done(true)
		return
	}

	var enc []encryptedID

	i := 0
	for i < len(listJSON) && listJSON[i] != '[' {
		i++
	}
	i++
	for i < len(listJSON) {
		for i < len(listJSON) && listJSON[i] != '{' && listJSON[i] != ']' {
			i++
		}
		if i >= len(listJSON) || listJSON[i] == ']' {
			break
		}
		end := i + 1
		depth := 1
		for end < len(listJSON) && depth > 0 {
			if listJSON[end] == '{' {
				depth++
			} else if listJSON[end] == '}' {
				depth--
			} else if listJSON[end] == '"' {
				end++
				for end < len(listJSON) && listJSON[end] != '"' {
					if listJSON[end] == '\\' {
						end++
					}
					end++
				}
			}
			end++
		}
		obj := listJSON[i:end]
		enc = append(enc, encryptedID{
			privkey: helpers.JsonGetString(obj, "privkey"),
			nick:    helpers.JsonGetString(obj, "nick"),
		})
		i = end
	}

	if len(enc) == 0 {
		done(true)
		return
	}

	// Decrypt sequentially (async chain).
	decryptNext(enc, 0, key, iv, done)
}

func decryptNext(enc []encryptedID, idx int, key, iv []byte, done func(bool)) {
	if idx >= len(enc) {
		done(true)
		return
	}
	e := enc[idx]
	decryptField(e.privkey, key, iv, func(skHex string) {
		if skHex == "" {
			// Skip invalid entries.
			decryptNext(enc, idx+1, key, iv, done)
			return
		}
		decryptField(e.nick, key, iv, func(nick string) {
			// Derive pubkey from secret key.
			skBytes := helpers.HexDecode(skHex)
			if skBytes == nil {
				decryptNext(enc, idx+1, key, iv, done)
				return
			}
			// Use schnorr to get pubkey.
			pk, ok := schnorrPubFromSec(skBytes)
			if !ok {
				decryptNext(enc, idx+1, key, iv, done)
				return
			}
			identities = append(identities, identity{
				Pubkey: helpers.HexEncode(pk),
				Seckey: skHex,
				Name:   nick,
			})
			decryptNext(enc, idx+1, key, iv, done)
		})
	})
}

// schnorrPubFromSec wraps the schnorr bridge to derive pubkey from seckey.
func schnorrPubFromSec(sk []byte) ([]byte, bool) {
	return schnorr.PubKeyFromSecKey(sk)
}

// --- Vault serialization ---

// saveVault encrypts and saves the vault to storage.
func saveVault(done func()) {
	if !vaultOpen {
		if done != nil {
			done()
		}
		return
	}

	// Encrypt all identities, then build JSON.
	encryptAllIdentities(vaultKey, vaultIV, func(idJSON string) {
		// Build vault JSON.
		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 += ",\"identities\":" + idJSON
		s += ",\"permissions\":[]"
		s += ",\"relays\":[]"
		s += ",\"selectedIdentityId\":null"
		s += "}"

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

func encryptAllIdentities(key, iv []byte, fn func(string)) {
	if len(identities) == 0 {
		fn("[]")
		return
	}
	encryptIDAt(key, iv, 0, "[", fn)
}

func encryptIDAt(key, iv []byte, idx int, acc string, fn func(string)) {
	if idx >= len(identities) {
		fn(acc + "]")
		return
	}
	id := identities[idx]
	if idx > 0 {
		acc += ","
	}
	// Encrypt privkey and nick.
	encryptField(id.Seckey, key, iv, func(encSK string) {
		encryptField(id.Name, key, iv, func(encNick string) {
			// Generate a stable ID from pubkey.
			encryptField(id.Pubkey, key, iv, func(encID string) {
				obj := "{\"id\":" + helpers.JsonString(encID) +
					",\"nick\":" + helpers.JsonString(encNick) +
					",\"createdAt\":" + helpers.JsonString("") +
					",\"privkey\":" + helpers.JsonString(encSK) + "}"
				encryptIDAt(key, iv, idx+1, acc+obj, fn)
			})
		})
	})
}

func boolStr(b bool) string {
	if b {
		return "true"
	}
	return "false"
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	s := ""
	for n > 0 {
		s = string(rune('0'+n%10)) + s
		n /= 10
	}
	return s
}

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