/**
 * Wallet state management using Zustand.
 */

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import {
  createWallet,
  generateMnemonic,
  validateMnemonic,
  encrypt,
  decrypt,
  bytesToHex,
  hexToBytes,
  type Keypair,
} from '../lib/crypto';
import { getClient, type Balance, type Transaction } from '../lib/rpc';

export type Network = 'mainnet' | 'testnet' | 'devnet';

interface EncryptedData {
  ciphertext: string;
  iv: string;
  salt: string;
}

interface WalletState {
  // Persisted state
  hasWallet: boolean;
  encryptedSeed: EncryptedData | null;
  address: string | null;
  network: Network;
  rpcEndpoint: string;

  // Session state (not persisted)
  isUnlocked: boolean;
  keypair: Keypair | null;
  balance: Balance | null;
  transactions: Transaction[];
  isLoading: boolean;
  error: string | null;

  // Actions
  createNewWallet: (password: string) => Promise<string>;
  recoverWallet: (mnemonic: string, password: string) => Promise<void>;
  unlock: (password: string) => Promise<void>;
  lock: () => void;
  deleteWallet: () => void;
  refreshBalance: () => Promise<void>;
  refreshTransactions: () => Promise<void>;
  setNetwork: (network: Network) => void;
  setRpcEndpoint: (endpoint: string) => void;
  clearError: () => void;
}

export const useWalletStore = create<WalletState>()(
  persist(
    (set, get) => ({
      // Persisted state
      hasWallet: false,
      encryptedSeed: null,
      address: null,
      network: 'testnet',
      rpcEndpoint: 'http://localhost:16110',

      // Session state
      isUnlocked: false,
      keypair: null,
      balance: null,
      transactions: [],
      isLoading: false,
      error: null,

      createNewWallet: async (password: string): Promise<string> => {
        set({ isLoading: true, error: null });

        try {
          const mnemonic = generateMnemonic(24);
          const wallet = await createWallet(mnemonic, '', get().network);

          // Encrypt the seed
          const encrypted = await encrypt(wallet.seed, password);

          set({
            hasWallet: true,
            encryptedSeed: {
              ciphertext: bytesToHex(encrypted.ciphertext),
              iv: bytesToHex(encrypted.iv),
              salt: bytesToHex(encrypted.salt),
            },
            address: wallet.address,
            isUnlocked: true,
            keypair: wallet.keypair,
            isLoading: false,
          });

          return mnemonic;
        } catch (error) {
          const message = error instanceof Error ? error.message : 'Failed to create wallet';
          set({ error: message, isLoading: false });
          throw error;
        }
      },

      recoverWallet: async (mnemonic: string, password: string): Promise<void> => {
        set({ isLoading: true, error: null });

        try {
          if (!validateMnemonic(mnemonic)) {
            throw new Error('Invalid mnemonic phrase');
          }

          const wallet = await createWallet(mnemonic, '', get().network);

          // Encrypt the seed
          const encrypted = await encrypt(wallet.seed, password);

          set({
            hasWallet: true,
            encryptedSeed: {
              ciphertext: bytesToHex(encrypted.ciphertext),
              iv: bytesToHex(encrypted.iv),
              salt: bytesToHex(encrypted.salt),
            },
            address: wallet.address,
            isUnlocked: true,
            keypair: wallet.keypair,
            isLoading: false,
          });
        } catch (error) {
          const message = error instanceof Error ? error.message : 'Failed to recover wallet';
          set({ error: message, isLoading: false });
          throw error;
        }
      },

      unlock: async (password: string): Promise<void> => {
        const { encryptedSeed, network } = get();

        if (!encryptedSeed) {
          throw new Error('No wallet found');
        }

        set({ isLoading: true, error: null });

        try {
          const seed = await decrypt(
            hexToBytes(encryptedSeed.ciphertext),
            hexToBytes(encryptedSeed.iv),
            hexToBytes(encryptedSeed.salt),
            password
          );

          // Re-derive keypair from decrypted seed
          const { deriveKeypair, publicKeyToAddress } = await import('../lib/crypto');
          const keypair = await deriveKeypair(seed);
          const address = publicKeyToAddress(keypair.publicKey, network);

          set({
            isUnlocked: true,
            keypair,
            address,
            isLoading: false,
          });

          // Refresh balance after unlock
          get().refreshBalance();
        } catch (error) {
          const message =
            error instanceof Error
              ? error.message.includes('decrypt')
                ? 'Incorrect password'
                : error.message
              : 'Failed to unlock wallet';
          set({ error: message, isLoading: false });
          throw error;
        }
      },

      lock: () => {
        set({
          isUnlocked: false,
          keypair: null,
          balance: null,
          transactions: [],
        });
      },

      deleteWallet: () => {
        set({
          hasWallet: false,
          encryptedSeed: null,
          address: null,
          isUnlocked: false,
          keypair: null,
          balance: null,
          transactions: [],
        });
      },

      refreshBalance: async (): Promise<void> => {
        const { address, isUnlocked } = get();

        if (!isUnlocked || !address) {
          return;
        }

        try {
          const client = getClient();
          const balance = await client.getBalance(address);
          set({ balance });
        } catch (error) {
          console.error('Failed to refresh balance:', error);
        }
      },

      refreshTransactions: async (): Promise<void> => {
        const { address, isUnlocked } = get();

        if (!isUnlocked || !address) {
          return;
        }

        try {
          const client = getClient();
          const transactions = await client.getAddressTransactions(address, 50);
          set({ transactions });
        } catch (error) {
          console.error('Failed to refresh transactions:', error);
        }
      },

      setNetwork: (network: Network) => {
        set({ network });
        // Note: Address format might change based on network
        // User would need to re-unlock to update address
      },

      setRpcEndpoint: (endpoint: string) => {
        set({ rpcEndpoint: endpoint });
      },

      clearError: () => {
        set({ error: null });
      },
    }),
    {
      name: 'synor-wallet',
      // Only persist these fields
      partialize: (state) => ({
        hasWallet: state.hasWallet,
        encryptedSeed: state.encryptedSeed,
        address: state.address,
        network: state.network,
        rpcEndpoint: state.rpcEndpoint,
      }),
    }
  )
);
