import crypto from 'node:crypto'
import { sql } from 'drizzle-orm'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { installation, sshKey } from '../../server/db/schema'
import { clearDb, setDb, useDb } from '../../server/utils/db'
import { clearEncryptionKeyCache, decrypt } from '../../server/utils/encryption'
import { createTestDb } from '../utils/db'

const ORIGINAL_ENC_KEY = process.env.NUXT_ENCRYPTION_KEY

const createRecordMock = vi.fn<(input: { repo: string, collection: string, record: Record<string, unknown> }) => Promise<{ data: { uri: string, cid: string } }>>()
const deleteRecordMock = vi.fn<(input: { repo: string, collection: string, rkey: string }) => Promise<unknown>>()
const restoreMock = vi.fn<(did: string) => Promise<{ did: string }>>()

vi.mock('@atproto/api', () => ({
  Agent: class {
    com = {
      atproto: {
        repo: {
          createRecord: createRecordMock,
          deleteRecord: deleteRecordMock,
        },
      },
    }
  },
}))

vi.mock('../../server/utils/atproto-oauth', () => ({
  useOAuthClient: async () => ({ restore: restoreMock }),
}))

const { generateAndPublishKey, revokeKeyForInstallationDid, revokeKeysForInstallation, rotateKey } = await import('../../server/utils/tangled-pubkey')

function fakeOauthSession(did: string) {
  // The Agent mock above ignores its constructor argument, so we only need
  // a `.did` field for the helper itself.
  // eslint-disable-next-line ts/no-unsafe-type-assertion
  return { did } as unknown as Parameters<typeof generateAndPublishKey>[0]['oauthSession']
}

describe('generateAndPublishKey', () => {
  beforeEach(async () => {
    process.env.NUXT_ENCRYPTION_KEY = crypto.randomBytes(32).toString('base64')
    clearEncryptionKeyCache()

    setDb(await createTestDb())
    const db = useDb()
    await db.insert(installation).values({
      id: 1,
      accountLogin: 'alice',
      accountId: 100,
      accountType: 'User',
    })

    createRecordMock.mockReset()
    deleteRecordMock.mockReset()
    createRecordMock.mockResolvedValue({
      data: { uri: 'at://did:plc:abc/sh.tangled.publicKey/3kh2y4xq2lk2v', cid: 'bafy' },
    })
    deleteRecordMock.mockResolvedValue({})
  })

  afterEach(() => {
    if (ORIGINAL_ENC_KEY === undefined) delete process.env.NUXT_ENCRYPTION_KEY
    else process.env.NUXT_ENCRYPTION_KEY = ORIGINAL_ENC_KEY
    clearEncryptionKeyCache()
    clearDb()
  })

  it('generates a key, publishes to PDS, and stores the encrypted private half', async () => {
    const result = await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })

    expect(result.created).toBe(true)
    expect(createRecordMock).toHaveBeenCalledTimes(1)
    const call = createRecordMock.mock.calls[0]![0]
    expect(call.repo).toBe('did:plc:abc')
    expect(call.collection).toBe('sh.tangled.publicKey')
    expect(call.record.$type).toBe('sh.tangled.publicKey')
    expect(call.record.key).toMatch(/^ssh-ed25519 /)
    expect(call.record.name).toBe('synchub.to/1')

    const db = useDb()
    const rows = await db.select().from(sshKey)
      .where(sql`${sshKey.installationId} = 1 AND ${sshKey.did} = 'did:plc:abc'`)
    expect(rows).toHaveLength(1)
    const row = rows[0]!
    expect(row.publicKey).toMatch(/^ssh-ed25519 /)
    expect(row.tangledKeyRkey).toBe('3kh2y4xq2lk2v')

    const decrypted = decrypt(row.privateKeyCiphertext, row.privateKeyNonce)
    expect(decrypted).toMatch(/^-----BEGIN PRIVATE KEY-----/)
    expect(decrypted).toContain('-----END PRIVATE KEY-----')
  })

  it('no-ops if a key already exists for (installation, did)', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    expect(createRecordMock).toHaveBeenCalledTimes(1)

    const result = await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    expect(result.created).toBe(false)
    expect(createRecordMock).toHaveBeenCalledTimes(1) // not called again

    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(1)
  })

  it('does not write a row if the PDS publish fails', async () => {
    createRecordMock.mockRejectedValueOnce(new Error('pds is sad'))

    await expect(generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })).rejects.toThrow(/pds is sad/)

    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(0)
  })
})

describe('rotateKey', () => {
  beforeEach(async () => {
    process.env.NUXT_ENCRYPTION_KEY = crypto.randomBytes(32).toString('base64')
    clearEncryptionKeyCache()

    setDb(await createTestDb())
    await useDb().insert(installation).values({
      id: 1, accountLogin: 'alice', accountId: 100, accountType: 'User',
    })

    createRecordMock.mockReset()
    deleteRecordMock.mockReset()
    let counter = 0
    createRecordMock.mockImplementation(async () => {
      counter += 1
      return { data: { uri: `at://did:plc:abc/sh.tangled.publicKey/rkey-${counter}`, cid: 'bafy' } }
    })
    deleteRecordMock.mockResolvedValue({})
  })

  afterEach(() => {
    if (ORIGINAL_ENC_KEY === undefined) delete process.env.NUXT_ENCRYPTION_KEY
    else process.env.NUXT_ENCRYPTION_KEY = ORIGINAL_ENC_KEY
    clearEncryptionKeyCache()
    clearDb()
  })

  it('deletes the old PDS record and publishes a fresh key', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    const db = useDb()
    const before = await db.select().from(sshKey).where(sql`${sshKey.installationId} = 1`)
    expect(before).toHaveLength(1)
    const oldPubKey = before[0]!.publicKey
    const oldRkey = before[0]!.tangledKeyRkey

    const result = await rotateKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    expect(result.created).toBe(true)

    expect(deleteRecordMock).toHaveBeenCalledTimes(1)
    const del = deleteRecordMock.mock.calls[0]![0]
    expect(del.repo).toBe('did:plc:abc')
    expect(del.collection).toBe('sh.tangled.publicKey')
    expect(del.rkey).toBe(oldRkey)

    const after = await db.select().from(sshKey).where(sql`${sshKey.installationId} = 1`)
    expect(after).toHaveLength(1)
    expect(after[0]!.publicKey).not.toBe(oldPubKey)
    expect(after[0]!.tangledKeyRkey).toBe('rkey-2')
  })

  it('proceeds when the PDS reports the record is already gone (404)', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })

    const notFound: Error & { status: number } = Object.assign(new Error('not found'), { status: 404 })
    deleteRecordMock.mockRejectedValueOnce(notFound)

    const result = await rotateKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    expect(result.created).toBe(true)

    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(1)
    expect(rows[0]!.tangledKeyRkey).toBe('rkey-2')
  })

  it('aborts the rotation if the PDS delete fails for a non-404 reason', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })

    deleteRecordMock.mockRejectedValueOnce(Object.assign(new Error('boom'), { status: 500 }))

    await expect(rotateKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })).rejects.toThrow(/boom/)

    // Old row still present, no second createRecord call.
    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(1)
    expect(rows[0]!.tangledKeyRkey).toBe('rkey-1')
    expect(createRecordMock).toHaveBeenCalledTimes(1)
  })

  it('mints a fresh key even if there is no existing row', async () => {
    const result = await rotateKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    expect(result.created).toBe(true)
    expect(deleteRecordMock).not.toHaveBeenCalled()
    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(1)
  })
})

describe('revokeKeysForInstallation', () => {
  beforeEach(async () => {
    process.env.NUXT_ENCRYPTION_KEY = crypto.randomBytes(32).toString('base64')
    clearEncryptionKeyCache()

    setDb(await createTestDb())
    await useDb().insert(installation).values({
      id: 1, accountLogin: 'alice', accountId: 100, accountType: 'User',
    })

    createRecordMock.mockReset()
    deleteRecordMock.mockReset()
    restoreMock.mockReset()
    createRecordMock.mockResolvedValue({
      data: { uri: 'at://did:plc:abc/sh.tangled.publicKey/3kh2y4xq2lk2v', cid: 'bafy' },
    })
    deleteRecordMock.mockResolvedValue({})
    restoreMock.mockImplementation(async (did: string) => ({ did }))
  })

  afterEach(() => {
    if (ORIGINAL_ENC_KEY === undefined) delete process.env.NUXT_ENCRYPTION_KEY
    else process.env.NUXT_ENCRYPTION_KEY = ORIGINAL_ENC_KEY
    clearEncryptionKeyCache()
    clearDb()
  })

  it('deletes the publicKey PDS record for the installation', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })

    await revokeKeysForInstallation(1)

    expect(restoreMock).toHaveBeenCalledWith('did:plc:abc')
    expect(deleteRecordMock).toHaveBeenCalledTimes(1)
    const del = deleteRecordMock.mock.calls[0]![0]
    expect(del.repo).toBe('did:plc:abc')
    expect(del.collection).toBe('sh.tangled.publicKey')
    expect(del.rkey).toBe('3kh2y4xq2lk2v')
  })

  it('no-ops when the installation has no keys', async () => {
    await revokeKeysForInstallation(1)
    expect(restoreMock).not.toHaveBeenCalled()
    expect(deleteRecordMock).not.toHaveBeenCalled()
  })

  it('swallows a 404 from the PDS delete', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    deleteRecordMock.mockRejectedValueOnce(Object.assign(new Error('not found'), { status: 404 }))

    await expect(revokeKeysForInstallation(1)).resolves.toBeUndefined()
  })

  it('continues when OAuth session restoration fails', async () => {
    await generateAndPublishKey({
      oauthSession: fakeOauthSession('did:plc:abc'),
      installationId: 1,
    })
    restoreMock.mockRejectedValueOnce(new Error('session gone'))

    await expect(revokeKeysForInstallation(1)).resolves.toBeUndefined()
    expect(deleteRecordMock).not.toHaveBeenCalled()
  })
})

describe('revokeKeyForInstallationDid', () => {
  beforeEach(async () => {
    process.env.NUXT_ENCRYPTION_KEY = crypto.randomBytes(32).toString('base64')
    clearEncryptionKeyCache()

    setDb(await createTestDb())
    await useDb().insert(installation).values({
      id: 1, accountLogin: 'alice', accountId: 100, accountType: 'User',
    })

    createRecordMock.mockReset()
    deleteRecordMock.mockReset()
    restoreMock.mockReset()
    let counter = 0
    createRecordMock.mockImplementation(async () => {
      counter += 1
      return { data: { uri: `at://did:plc:${counter}/sh.tangled.publicKey/rkey-${counter}`, cid: 'bafy' } }
    })
    deleteRecordMock.mockResolvedValue({})
    restoreMock.mockImplementation(async (did: string) => ({ did }))
  })

  afterEach(() => {
    if (ORIGINAL_ENC_KEY === undefined) delete process.env.NUXT_ENCRYPTION_KEY
    else process.env.NUXT_ENCRYPTION_KEY = ORIGINAL_ENC_KEY
    clearEncryptionKeyCache()
    clearDb()
  })

  it('revokes the PDS record and drops the local row for one (install, did)', async () => {
    await generateAndPublishKey({ oauthSession: fakeOauthSession('did:plc:1'), installationId: 1 })
    const db = useDb()
    expect(await db.select().from(sshKey)).toHaveLength(1)

    await revokeKeyForInstallationDid(1, 'did:plc:1')

    expect(restoreMock).toHaveBeenCalledWith('did:plc:1')
    const del = deleteRecordMock.mock.calls[0]![0]
    expect(del.repo).toBe('did:plc:1')
    expect(del.collection).toBe('sh.tangled.publicKey')
    expect(del.rkey).toBe('rkey-1')
    expect(await db.select().from(sshKey)).toHaveLength(0)
  })

  it('leaves other dids on the same installation untouched', async () => {
    await generateAndPublishKey({ oauthSession: fakeOauthSession('did:plc:1'), installationId: 1 })
    await generateAndPublishKey({ oauthSession: fakeOauthSession('did:plc:2'), installationId: 1 })

    await revokeKeyForInstallationDid(1, 'did:plc:1')

    const db = useDb()
    const rows = await db.select().from(sshKey)
    expect(rows).toHaveLength(1)
    expect(rows[0]!.did).toBe('did:plc:2')
  })

  it('no-ops when no key exists for the pair', async () => {
    await revokeKeyForInstallationDid(1, 'did:plc:none')
    expect(restoreMock).not.toHaveBeenCalled()
    expect(deleteRecordMock).not.toHaveBeenCalled()
  })

  it('drops the local row even when the PDS delete fails', async () => {
    await generateAndPublishKey({ oauthSession: fakeOauthSession('did:plc:1'), installationId: 1 })
    deleteRecordMock.mockRejectedValueOnce(Object.assign(new Error('boom'), { status: 500 }))

    await expect(revokeKeyForInstallationDid(1, 'did:plc:1')).resolves.toBeUndefined()

    const db = useDb()
    expect(await db.select().from(sshKey)).toHaveLength(0)
  })

  it('drops the local row even when session restoration fails', async () => {
    await generateAndPublishKey({ oauthSession: fakeOauthSession('did:plc:1'), installationId: 1 })
    restoreMock.mockRejectedValueOnce(new Error('session gone'))

    await expect(revokeKeyForInstallationDid(1, 'did:plc:1')).resolves.toBeUndefined()

    const db = useDb()
    expect(await db.select().from(sshKey)).toHaveLength(0)
  })
})