import { toDelegateAddress, toStakeAddress } from '@arcblock/did-util';

import {
  getMigratedFromList,
  getDelegationAddressWithFallback,
  getStakingAddressWithFallback,
  getAppStakeAddressWithFallback,
  clearMigratedFromCache,
} from '../../src/libs/wallet-migration';

// Mock the auth module
jest.mock('../../src/libs/auth', () => ({
  wallet: {
    address: 'zNKcF7wyrvAvn4YaADjr4p9gpFUSHnW2aQJa', // Current app DID
  },
}));

// Mock logger
jest.mock('../../src/libs/logger', () => ({
  info: jest.fn(),
  warn: jest.fn(),
  error: jest.fn(),
}));

describe('wallet-migration', () => {
  // Mock OcapClient
  const createMockClient = (overrides: any = {}) => ({
    host: 'https://main.abtnetwork.io/api',
    getAccountState: jest.fn().mockResolvedValue({
      state: {
        migratedFrom: ['zNKq13Dr2TBHELpLDUJFGxepiGP7YHbAVxPn', 'zNKq6yG8AVbwdRBDJSNCDVJJUQD1AXdkK7Wp'],
        address: 'zNKcF7wyrvAvn4YaADjr4p9gpFUSHnW2aQJa',
      },
    }),
    getDelegateState: jest.fn().mockResolvedValue({ state: null }),
    getStakeState: jest.fn().mockResolvedValue({ state: null }),
    ...overrides,
  });

  beforeEach(() => {
    clearMigratedFromCache();
  });

  describe('getMigratedFromList', () => {
    it('should return migratedFrom list from chain', async () => {
      const mockClient = createMockClient();
      const result = await getMigratedFromList(mockClient as any);
      expect(result).toEqual(['zNKq13Dr2TBHELpLDUJFGxepiGP7YHbAVxPn', 'zNKq6yG8AVbwdRBDJSNCDVJJUQD1AXdkK7Wp']);
    });

    it('should cache the result and not query again', async () => {
      const mockClient = createMockClient();
      await getMigratedFromList(mockClient as any);
      await getMigratedFromList(mockClient as any);
      expect(mockClient.getAccountState).toHaveBeenCalledTimes(1);
    });

    it('should return empty array if no migratedFrom', async () => {
      const mockClient = createMockClient({
        getAccountState: jest.fn().mockResolvedValue({ state: {} }),
      });
      const result = await getMigratedFromList(mockClient as any);
      expect(result).toEqual([]);
    });

    it('should return empty array on error', async () => {
      const mockClient = createMockClient({
        getAccountState: jest.fn().mockRejectedValue(new Error('Network error')),
      });
      const result = await getMigratedFromList(mockClient as any);
      expect(result).toEqual([]);
    });
  });

  describe('getDelegationAddressWithFallback', () => {
    const delegator = 'zNKtCNqYWLYWYW3gWRA1vnRNdWLrFXfPxRKW';
    const currentWalletAddress = 'zNKcF7wyrvAvn4YaADjr4p9gpFUSHnW2aQJa';
    const oldAppDid = 'zNKq13Dr2TBHELpLDUJFGxepiGP7YHbAVxPn';

    it('should return stored address if available and valid', async () => {
      const mockClient = createMockClient({
        getDelegateState: jest.fn().mockImplementation(async ({ address }) => {
          if (address === 'stored_delegation_address') {
            return { state: { ops: [{ key: 'fg:x:transfer' }] } };
          }
          return { state: null };
        }),
      });

      const result = await getDelegationAddressWithFallback({
        storedAddress: 'stored_delegation_address',
        delegator,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: 'stored_delegation_address',
        needsBackfill: false,
        source: 'stored',
      });
      // Should verify stored address
      expect(mockClient.getDelegateState).toHaveBeenCalledWith({ address: 'stored_delegation_address' });
    });

    it('should fallback if stored address has no valid state', async () => {
      const oldAddress = toDelegateAddress(delegator, oldAppDid);
      const mockClient = createMockClient({
        getDelegateState: jest.fn().mockImplementation(async ({ address }) => {
          // stored address has no state, but migratedFrom address has
          if (address === oldAddress) {
            return { state: { ops: [{ key: 'fg:x:transfer' }] } };
          }
          return { state: null };
        }),
      });

      const result = await getDelegationAddressWithFallback({
        storedAddress: 'invalid_stored_address',
        delegator,
        client: mockClient as any,
      });

      // Should fallback to migratedFrom
      expect(result).toEqual({
        address: oldAddress,
        needsBackfill: true,
        source: 'migrated',
      });
    });

    it('should return current address if delegation exists on current wallet', async () => {
      const currentAddress = toDelegateAddress(delegator, currentWalletAddress);
      const mockClient = createMockClient({
        getDelegateState: jest.fn().mockImplementation(async ({ address }) => {
          if (address === currentAddress) {
            return { state: { ops: [{ key: 'fg:x:transfer' }] } };
          }
          return { state: null };
        }),
      });

      const result = await getDelegationAddressWithFallback({
        delegator,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: currentAddress,
        needsBackfill: true,
        source: 'current',
      });
    });

    it('should fallback to migratedFrom addresses', async () => {
      const oldAddress = toDelegateAddress(delegator, oldAppDid);
      const mockClient = createMockClient({
        getDelegateState: jest.fn().mockImplementation(async ({ address }) => {
          if (address === oldAddress) {
            return { state: { ops: [{ key: 'fg:x:transfer' }] } };
          }
          return { state: null };
        }),
      });

      const result = await getDelegationAddressWithFallback({
        delegator,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: oldAddress,
        needsBackfill: true,
        source: 'migrated',
      });
    });

    it('should return null if no delegation found', async () => {
      const mockClient = createMockClient();

      const result = await getDelegationAddressWithFallback({
        delegator,
        client: mockClient as any,
      });

      expect(result).toBeNull();
    });
  });

  describe('getStakingAddressWithFallback', () => {
    const userDid = 'zNKtCNqYWLYWYW3gWRA1vnRNdWLrFXfPxRKW';
    const nonce = 'sub_123';
    const currentWalletAddress = 'zNKcF7wyrvAvn4YaADjr4p9gpFUSHnW2aQJa';
    const oldAppDid = 'zNKq13Dr2TBHELpLDUJFGxepiGP7YHbAVxPn';

    it('should return stored address if available', async () => {
      const mockClient = createMockClient();

      const result = await getStakingAddressWithFallback({
        storedAddress: 'stored_stake_address',
        userDid,
        nonce,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: 'stored_stake_address',
        needsBackfill: false,
        source: 'stored',
      });
    });

    it('should return current address if stake exists on current wallet', async () => {
      const currentAddress = toStakeAddress(userDid, currentWalletAddress, nonce);
      const mockClient = createMockClient({
        getStakeState: jest.fn().mockImplementation(({ address }) => {
          if (address === currentAddress) {
            return { state: { tokens: [] } };
          }
          return { state: null };
        }),
      });

      const result = await getStakingAddressWithFallback({
        userDid,
        nonce,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: currentAddress,
        needsBackfill: true,
        source: 'current',
      });
    });

    it('should fallback to migratedFrom addresses', async () => {
      const oldAddress = toStakeAddress(userDid, oldAppDid, nonce);
      const mockClient = createMockClient({
        getStakeState: jest.fn().mockImplementation(({ address }) => {
          if (address === oldAddress) {
            return { state: { tokens: [] } };
          }
          return { state: null };
        }),
      });

      const result = await getStakingAddressWithFallback({
        userDid,
        nonce,
        client: mockClient as any,
      });

      expect(result).toEqual({
        address: oldAddress,
        needsBackfill: true,
        source: 'migrated',
      });
    });
  });

  describe('getAppStakeAddressWithFallback', () => {
    const currentWalletAddress = 'zNKcF7wyrvAvn4YaADjr4p9gpFUSHnW2aQJa';
    const oldAppDid = 'zNKq13Dr2TBHELpLDUJFGxepiGP7YHbAVxPn';

    it('should return current address if app stake exists', async () => {
      const currentAddress = toStakeAddress(currentWalletAddress, currentWalletAddress);
      const mockClient = createMockClient({
        getStakeState: jest.fn().mockImplementation(({ address }) => {
          if (address === currentAddress) {
            return { state: { tokens: [] } };
          }
          return { state: null };
        }),
      });

      const result = await getAppStakeAddressWithFallback({ client: mockClient as any });

      expect(result).toEqual({
        address: currentAddress,
        needsBackfill: false,
        source: 'current',
      });
    });

    it('should fallback to migratedFrom addresses for app stake', async () => {
      const oldAddress = toStakeAddress(oldAppDid, oldAppDid);
      const mockClient = createMockClient({
        getStakeState: jest.fn().mockImplementation(({ address }) => {
          if (address === oldAddress) {
            return { state: { tokens: [] } };
          }
          return { state: null };
        }),
      });

      const result = await getAppStakeAddressWithFallback({ client: mockClient as any });

      expect(result).toEqual({
        address: oldAddress,
        needsBackfill: false,
        source: 'migrated',
      });
    });
  });
});
