import { numberToHex, zeroAddress, type Address } from 'viem'
import { describe, expect, test } from 'vitest'

import * as Projection from './Projection.js'

const alice = '0x0000000000000000000000000000000000000001'
const bob = '0x0000000000000000000000000000000000000002'

describe('apply', () => {
  test('qualifies deposits under both caps and never qualifies transfers', () => {
    const result = Projection.apply({
      accounts: [],
      eligible: [
        { recipient: alice, registrationOrder: 1n },
        { recipient: bob, registrationOrder: 2n },
      ],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, alice, bob, 25n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 60n,
      totalPrincipalCapAssets: 60n,
    })
    expect(
      result.accounts.map((account) => ({
        allocated: account.allocatedPrincipalAssets,
        public: account.publicEarnShares,
        qualified: account.qualifiedEarnShares,
        recipient: account.recipient,
      })),
    ).toMatchInlineSnapshot(`
      [
        {
          "allocated": 60n,
          "public": 75n,
          "qualified": 35n,
          "recipient": "0x0000000000000000000000000000000000000001",
        },
        {
          "allocated": 0n,
          "public": 25n,
          "qualified": 0n,
          "recipient": "0x0000000000000000000000000000000000000002",
        },
      ]
    `)
  })
})

describe('bootstrap', () => {
  test('allocates the global cap only to surviving historical deposit lots', () => {
    const result = Projection.bootstrap({
      eligible: [
        { recipient: alice, registrationOrder: 1n },
        { recipient: bob, registrationOrder: 2n },
      ],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, alice, zeroAddress, 100n),
        transfer(4, zeroAddress, bob, 100n),
        deposit(5, bob, 100n, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      totalPrincipalCapAssets: 100n,
    })
    expect(
      result.accounts.map((account) => ({
        allocated: account.allocatedPrincipalAssets,
        public: account.publicEarnShares,
        qualified: account.qualifiedEarnShares,
        recipient: account.recipient,
      })),
    ).toStrictEqual([
      { allocated: 0n, public: 0n, qualified: 0n, recipient: alice },
      { allocated: 100n, public: 100n, qualified: 100n, recipient: bob },
    ])
  })

  test('materializes a registered wallet without historical activity', () => {
    expect(
      Projection.bootstrap({
        eligible: [{ recipient: alice, registrationOrder: 1n }],
        events: [],
        excluded: [],
        perUserPrincipalCapAssets: 100n,
        totalPrincipalCapAssets: 100n,
      }).accounts,
    ).toHaveLength(1)
  })

  test('preserves qualified principal across historical and live self-transfers', () => {
    const bootstrapped = Projection.bootstrap({
      eligible: [{ recipient: alice, registrationOrder: 1n }],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, alice, alice, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      totalPrincipalCapAssets: 100n,
    })
    const calculated = Projection.calculate({
      accounts: bootstrapped.accounts,
      annualRateBps: 700,
      assetPerEarnShareWad: 10n ** 18n,
      endsAt: 300,
      events: [transfer(4, alice, alice, 100n, 150)],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      startsAt: 0,
      totalPrincipalCapAssets: 100n,
    })

    expect(
      calculated.accounts.map((account) => ({
        lots: account.lots.map((lot) => lot.remainingEarnShares),
        public: account.publicEarnShares,
        qualified: account.qualifiedEarnShares,
      })),
    ).toStrictEqual([{ lots: [], public: 100n, qualified: 100n }])
  })

  test('replays late registrations once while preserving registration-order cap priority', () => {
    const result = Projection.bootstrap({
      eligible: [
        { recipient: alice, registrationOrder: 2n },
        { recipient: bob, registrationOrder: 1n },
      ],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, zeroAddress, bob, 100n),
        deposit(4, bob, 100n, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'registration',
      totalPrincipalCapAssets: 100n,
    })

    expect(
      result.accounts.map((account) => ({
        allocated: account.allocatedPrincipalAssets,
        recipient: account.recipient,
      })),
    ).toStrictEqual([
      { allocated: 100n, recipient: bob },
      { allocated: 0n, recipient: alice },
    ])
  })

  test('allocates the initial campaign cap in canonical deposit order', () => {
    const result = Projection.bootstrap({
      eligible: [
        { recipient: alice, registrationOrder: 2n },
        { recipient: bob, registrationOrder: 1n },
      ],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, zeroAddress, bob, 100n),
        deposit(4, bob, 100n, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'deposit',
      totalPrincipalCapAssets: 100n,
    })

    expect(
      Object.fromEntries(
        result.accounts.map((account) => [account.recipient, account.allocatedPrincipalAssets]),
      ),
    ).toStrictEqual({ [alice]: 100n, [bob]: 0n })
  })
})

describe('reconcileRegistrations', () => {
  test('rebuilds refreshed accounts once and preserves accrued rewards', () => {
    const previous = {
      ...Projection.bootstrap({
        eligible: [{ recipient: alice, registrationOrder: 1n }],
        events: [],
        excluded: [],
        perUserPrincipalCapAssets: 100n,
        totalPrincipalCapAssets: 150n,
      }).accounts[0]!,
      accrualRemainder: 3n,
      cumulativeEntitlement: 11n,
      cumulativePaid: 7n,
      eligibilityRegisteredAt: '2026-09-01T12:00:00.000Z',
      pendingRewardAssets: 5n,
    }
    const options = {
      accounts: [previous],
      eligible: [
        {
          recipient: alice,
          registeredAt: '2026-09-01T13:00:00.000Z',
          registrationOrder: 1n,
        },
        {
          recipient: bob,
          registeredAt: '2026-09-01T13:00:01.000Z',
          registrationOrder: 2n,
        },
      ],
      events: [
        transfer(1, zeroAddress, alice, 100n),
        deposit(2, alice, 100n, 100n),
        transfer(3, zeroAddress, bob, 100n),
        deposit(4, bob, 100n, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'registration',
      totalPrincipalCapAssets: 150n,
    } satisfies Projection.reconcileRegistrations.Options

    const first = Projection.reconcileRegistrations(options)
    const replay = Projection.reconcileRegistrations({ ...options, accounts: first.accounts })

    expect({
      accounts: first.accounts.map((account) => ({
        allocated: account.allocatedPrincipalAssets,
        entitlement: account.cumulativeEntitlement,
        paid: account.cumulativePaid,
        pending: account.pendingRewardAssets,
        recipient: account.recipient,
        registeredAt: account.eligibilityRegisteredAt,
        remainder: account.accrualRemainder,
      })),
      replayed: replay.reconciled,
    }).toMatchInlineSnapshot(`
      {
        "accounts": [
          {
            "allocated": 100n,
            "entitlement": 11n,
            "paid": 7n,
            "pending": 5n,
            "recipient": "0x0000000000000000000000000000000000000001",
            "registeredAt": "2026-09-01T13:00:00.000Z",
            "remainder": 3n,
          },
          {
            "allocated": 50n,
            "entitlement": 0n,
            "paid": 0n,
            "pending": 0n,
            "recipient": "0x0000000000000000000000000000000000000002",
            "registeredAt": "2026-09-01T13:00:01.000Z",
            "remainder": 0n,
          },
        ],
        "replayed": false,
      }
    `)
  })

  test('preserves lifetime allocations when withdrawn registrations are refreshed', () => {
    const previous = {
      ...Projection.bootstrap({
        eligible: [{ recipient: alice, registrationOrder: 2n }],
        events: [transfer(1, zeroAddress, alice, 100n), deposit(2, alice, 100n, 100n)],
        excluded: [],
        perUserPrincipalCapAssets: 100n,
        totalPrincipalCapAssets: 100n,
      }).accounts[0]!,
      eligibilityRegisteredAt: '2026-09-01T12:00:00.000Z',
      lots: [],
      publicEarnShares: 0n,
      qualifiedEarnShares: 0n,
    }

    const result = Projection.reconcileRegistrations({
      accounts: [previous],
      eligible: [
        {
          recipient: bob,
          registeredAt: '2026-09-01T13:00:00.000Z',
          registrationOrder: 1n,
        },
        {
          recipient: alice,
          registeredAt: '2026-09-01T13:00:01.000Z',
          registrationOrder: 2n,
        },
      ],
      events: [
        transfer(1, zeroAddress, bob, 100n),
        deposit(2, bob, 100n, 100n),
        transfer(3, zeroAddress, alice, 100n),
        deposit(4, alice, 100n, 100n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'registration',
      totalPrincipalCapAssets: 100n,
    })

    expect(
      result.accounts.map((account) => ({
        allocated: account.allocatedPrincipalAssets,
        recipient: account.recipient,
        registeredAt: account.eligibilityRegisteredAt,
      })),
    ).toStrictEqual([
      {
        allocated: 0n,
        recipient: bob,
        registeredAt: '2026-09-01T13:00:00.000Z',
      },
      {
        allocated: 100n,
        recipient: alice,
        registeredAt: '2026-09-01T13:00:01.000Z',
      },
    ])
  })
})

describe('reconcilePages', () => {
  test('resumes registration traversal after empty recipients', async () => {
    const recipients = Array.from({ length: 1_000 }, (_, index) =>
      numberToHex(index + 1, { size: 20 }),
    )
    const options: Omit<Projection.reconcileRegistrations.Options, 'events'> = {
      accounts: [],
      eligible: recipients.map((recipient, index) => ({
        recipient,
        registeredAt: '2026-09-01T12:00:00.000Z',
        registrationOrder: BigInt(index + 1),
      })),
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'registration',
      totalPrincipalCapAssets: 100n,
    }
    const traversed: Address[] = []
    let checkpoint: Projection.reconcilePages.Checkpoint | undefined
    await expect(
      Projection.reconcilePages({
        ...options,
        checkpointState: (state) => {
          if (!state.allocationRecipient) return
          traversed.push(state.allocationRecipient)
          checkpoint = state
          if (traversed.length === 500) throw new Error('interrupted')
        },
        pages: (pass) => (pass === 'survival' ? projectionPages([]) : recipientPages(recipients)),
      }),
    ).rejects.toThrow('interrupted')

    const resumed = await Projection.reconcilePages({
      ...options,
      checkpoint,
      checkpointState: (state) => {
        if (state.allocationRecipient) traversed.push(state.allocationRecipient)
      },
      pages: (pass) =>
        pass === 'survival'
          ? projectionPages([])
          : recipientPages(recipients, checkpoint?.allocationRecipient),
    })

    expect(resumed).toStrictEqual(Projection.reconcileRegistrations({ ...options, events: [] }))
    expect(traversed).toStrictEqual(recipients)
    expect(checkpoint?.cursor).toBeUndefined()
  })

  test('bounds checkpoints across many differently priced FIFO deposit lots', async () => {
    const deposits = Array.from({ length: 1_000 }, (_, index) =>
      deposit(index + 1, alice, BigInt((index % 3) + 1), 1n),
    )
    const events = [...deposits, transfer(2_000, alice, zeroAddress, 750n)]
    const options: Omit<Projection.reconcileRegistrations.Options, 'events'> = {
      accounts: [],
      eligible: [
        { recipient: alice, registeredAt: '2026-09-01T12:00:00.000Z', registrationOrder: 1n },
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'deposit',
      totalPrincipalCapAssets: 100n,
    }
    const checkpoints: Projection.reconcilePages.Checkpoint[] = []
    let checkpoint: Projection.reconcilePages.Checkpoint | undefined
    await expect(
      Projection.reconcilePages({
        ...options,
        checkpointState: (state) => {
          checkpoints.push(state)
          checkpoint = state
          if (state.pass === 'allocation') throw new Error('interrupted')
        },
        pages: () =>
          projectionPages(
            Array.from({ length: 11 }, (_, index) => events.slice(index * 100, (index + 1) * 100)),
          ),
      }),
    ).rejects.toThrow('interrupted')
    const result = await Projection.reconcilePages({
      ...options,
      checkpoint,
      checkpointState: (state) => {
        checkpoints.push(state)
      },
      pages: (_, after) => {
        const remaining = after
          ? events.filter((event) => event.cursor.logIndex > after.logIndex)
          : events
        return projectionPages(
          Array.from({ length: Math.ceil(remaining.length / 100) }, (_, index) =>
            remaining.slice(index * 100, (index + 1) * 100),
          ),
        )
      },
    })

    expect(result).toStrictEqual(Projection.reconcileRegistrations({ ...options, events }))
    expect(checkpoints.every((checkpoint) => checkpoint.accounts[0]?.lots.length === 0)).toBe(true)
    expect(
      Math.max(...checkpoints.map((checkpoint) => JSON.stringify(checkpoint).length)),
    ).toBeLessThan(1_500)
  })

  test('resumes historical replay without changing cap allocation', async () => {
    const events = [
      transfer(1, zeroAddress, alice, 100n),
      deposit(2, alice, 100n, 100n),
      transfer(3, zeroAddress, bob, 100n),
      deposit(4, bob, 100n, 100n),
    ]
    const options: Omit<Projection.reconcileRegistrations.Options, 'events'> = {
      accounts: [],
      eligible: [
        { recipient: alice, registeredAt: '2026-09-01T12:00:00.000Z', registrationOrder: 2n },
        { recipient: bob, registeredAt: '2026-09-01T12:00:01.000Z', registrationOrder: 1n },
      ],
      excluded: [],
      perUserPrincipalCapAssets: 100n,
      principalAllocationOrder: 'registration',
      totalPrincipalCapAssets: 100n,
    }
    let checkpoint: Projection.reconcilePages.Checkpoint | undefined
    await expect(
      Projection.reconcilePages({
        ...options,
        checkpointState: (state) => {
          checkpoint = state
          throw new Error('interrupted')
        },
        pages: () => projectionPages([events.slice(0, 2)]),
      }),
    ).rejects.toThrow('interrupted')
    expect(() => JSON.stringify(checkpoint)).not.toThrow()

    const resumed = await Projection.reconcilePages({
      ...options,
      checkpoint,
      pages: (pass) =>
        projectionPages([pass === 'survival' ? events.slice(2) : [events[3]!, events[1]!]]),
    })
    const oneShot = Projection.reconcileRegistrations({ ...options, events })

    expect(resumed).toStrictEqual(oneShot)
  })
})

describe('calculate', () => {
  test('uses time-weighted qualified ownership and carries remainders', () => {
    const account = Projection.apply({
      accounts: [],
      eligible: [{ recipient: alice, registrationOrder: 1n }],
      events: [
        transfer(1, zeroAddress, alice, 10_000_000n),
        deposit(2, alice, 10_000_000n, 10_000_000n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 10_000_000n,
      totalPrincipalCapAssets: 10_000_000n,
    }).accounts
    const result = Projection.calculate({
      accounts: account,
      annualRateBps: 700,
      assetPerEarnShareWad: 10n ** 18n,
      endsAt: 300,
      events: [transfer(3, alice, bob, 5_000_000n, 150)],
      excluded: [],
      perUserPrincipalCapAssets: 10_000_000n,
      startsAt: 0,
      totalPrincipalCapAssets: 10_000_000n,
    })
    expect({
      alice: result.accounts[0]?.pendingRewardAssets,
      rewards: [...result.rewards],
    }).toMatchInlineSnapshot(`
      {
        "alice": 4n,
        "rewards": [
          [
            "0x0000000000000000000000000000000000000001",
            4n,
          ],
        ],
      }
    `)
  })

  test('stops accrual for an existing account added to exclusions', () => {
    const account = Projection.apply({
      accounts: [],
      eligible: [{ recipient: alice, registrationOrder: 1n }],
      events: [
        transfer(1, zeroAddress, alice, 10_000_000n),
        deposit(2, alice, 10_000_000n, 10_000_000n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 10_000_000n,
      totalPrincipalCapAssets: 10_000_000n,
    }).accounts
    const result = Projection.calculate({
      accounts: account,
      annualRateBps: 700,
      assetPerEarnShareWad: 10n ** 18n,
      endsAt: 300,
      events: [],
      excluded: [alice],
      perUserPrincipalCapAssets: 10_000_000n,
      startsAt: 0,
      totalPrincipalCapAssets: 10_000_000n,
    })
    expect({
      pending: result.accounts[0]?.pendingRewardAssets,
      rewards: [...result.rewards],
    }).toStrictEqual({ pending: 0n, rewards: [] })
  })

  test('calculates 10,000 active accounts with ordered transfers', () => {
    const accounts = Array.from({ length: 10_000 }, (_, index) => ({
      accrualRemainder: 0n,
      allocatedPrincipalAssets: 1n,
      cumulativeEntitlement: 0n,
      cumulativePaid: 0n,
      lots: [],
      pendingRewardAssets: 0n,
      publicEarnShares: 10n ** 18n,
      qualifiedEarnShares: 10n ** 18n,
      recipient: numberToHex(index + 1, { size: 20 }) as Address,
      registrationOrder: BigInt(index + 1),
    }))
    const events = accounts
      .slice(0, -1)
      .map((account, index) =>
        transfer(index, account.recipient, accounts[index + 1]!.recipient, 1n, 150),
      )

    const result = Projection.calculate({
      accounts,
      annualRateBps: 700,
      assetPerEarnShareWad: 10n ** 18n,
      endsAt: 300,
      events,
      excluded: [],
      perUserPrincipalCapAssets: 10_000n,
      startsAt: 0,
      totalPrincipalCapAssets: 100_000_000n,
    })

    expect(result.accounts).toHaveLength(10_000)
    expect(result.rewards.size).toBe(10_000)
  })
})

describe('calculatePages', () => {
  test('resumes time weighting without changing interval rewards', async () => {
    const accounts = Projection.apply({
      accounts: [],
      eligible: [{ recipient: alice, registrationOrder: 1n }],
      events: [
        transfer(1, zeroAddress, alice, 10_000_000n),
        deposit(2, alice, 10_000_000n, 10_000_000n),
      ],
      excluded: [],
      perUserPrincipalCapAssets: 10_000_000n,
      totalPrincipalCapAssets: 10_000_000n,
    }).accounts
    const events = [
      transfer(3, alice, zeroAddress, 2_000_000n, 100),
      transfer(4, zeroAddress, alice, 1_000_000n, 200),
    ]
    const options = {
      accounts,
      annualRateBps: 700,
      assetPerEarnShareWad: 10n ** 18n,
      endsAt: 300,
      excluded: [],
      perUserPrincipalCapAssets: 10_000_000n,
      startsAt: 0,
      totalPrincipalCapAssets: 10_000_000n,
    }
    let checkpoint: Projection.calculatePages.Checkpoint | undefined
    await Projection.calculatePages({
      ...options,
      checkpointState: (state) => {
        checkpoint = state
      },
      pages: pages([[events[0]!]]),
    })
    expect(() => JSON.stringify(checkpoint)).not.toThrow()

    const resumed = await Projection.calculatePages({
      ...options,
      checkpoint,
      pages: pages([[events[1]!]]),
    })
    const oneShot = Projection.calculate({ ...options, events })

    expect(resumed).toStrictEqual(oneShot)
  })
})

async function* pages(
  values: readonly (readonly Projection.Event[])[],
): AsyncGenerator<readonly Projection.Event[]> {
  yield* values
}

async function* projectionPages(
  values: readonly (readonly Projection.Event[])[],
): AsyncGenerator<Projection.reconcilePages.Page> {
  for (const events of values) yield { events }
}

async function* recipientPages(
  recipients: readonly Address[],
  after?: Address,
): AsyncGenerator<Projection.reconcilePages.Page> {
  const start = after ? recipients.indexOf(after) + 1 : 0
  for (const allocationRecipient of recipients.slice(start))
    yield { allocationRecipient, events: [] }
}

function deposit(
  logIndex: number,
  recipient: Address,
  assets: bigint,
  earnShares: bigint,
): Projection.Event {
  return {
    assets,
    cursor: { blockNumber: 1, logIndex, transactionIndex: 0 },
    earnShares,
    kind: 'deposit',
    recipient,
    timestamp: 0,
  }
}

function transfer(
  logIndex: number,
  from: Address,
  to: Address,
  amount: bigint,
  timestamp = 0,
): Projection.Event {
  return {
    amount,
    cursor: { blockNumber: 1, logIndex, transactionIndex: 0 },
    from,
    kind: 'transfer',
    timestamp,
    to,
  }
}