import assert from 'node:assert/strict'
import { createSign, generateKeyPairSync } from 'node:crypto'
import { test } from 'node:test'

import { OAuthClient, SessionRenewalPendingError } from './oauth.mjs'
import { SESSION_COOKIE, SESSION_RENEWAL_PENDING } from './sessions.mjs'

const identifier = 'a'.repeat(43)
const replacement = 'b'.repeat(43)
const now = Math.floor(Date.now() / 1000)
const { privateKey, publicKey } = generateKeyPairSync('rsa', {
  modulusLength: 2048,
})
const publicJwk = publicKey.export({ format: 'jwk' })
publicJwk.kid = 'test-key'
publicJwk.alg = 'RS256'
publicJwk.use = 'sig'

const config = {
  publicOrigin: 'https://remora.example.com',
  oidcAuthority: 'https://identity.example.com',
  oidcClientId: 'remora client',
  oidcClientSecret: 'secret:value',
  oidcScope: 'openid profile',
  secureCookie: true,
}

test('completes authorization code flow without exposing tokens to the browser', async () => {
  const store = new MemoryStore()
  const requests = []
  const oauth = new OAuthClient(config, store, async (url, init = {}) => {
    requests.push({ url: String(url), init })
    if (String(url).endsWith('/.well-known/openid-configuration'))
      return Response.json(metadata())
    if (String(url) === metadata().token_endpoint) {
      assert.equal(
        init.headers.Authorization,
        `Basic ${Buffer.from('remora+client:secret%3Avalue').toString('base64')}`,
      )
      const parameters = new URLSearchParams(init.body)
      assert.equal(parameters.get('code_verifier'), store.value.verifier)
      return Response.json({
        access_token: 'server-only-access-token',
        refresh_token: 'server-only-refresh-token',
        id_token: idToken(store.value.nonce),
        expires_in: 300,
      })
    }
    if (String(url) === metadata().jwks_uri)
      return Response.json({ keys: [publicJwk] })
    throw new Error(`unexpected request: ${url}`)
  })

  const login = new TestResponse()
  await oauth.login({ headers: {} }, login)
  const authorization = new URL(login.headers.Location)
  assert.equal(authorization.searchParams.get('code_challenge_method'), 'S256')
  assert.equal(authorization.searchParams.get('state'), store.value.state)
  assert.equal(authorization.searchParams.get('nonce'), store.value.nonce)
  assert.match(login.headers['Set-Cookie'], /HttpOnly; Secure; SameSite=Strict/)

  const callback = new TestResponse()
  await oauth.callback(
    { headers: { cookie: `${SESSION_COOKIE}=${identifier}` } },
    callback,
    new URL(
      `/auth/callback?code=authorization-code&state=${store.value.state}`,
      config.publicOrigin,
    ),
  )

  assert.equal(callback.status, 302)
  assert.equal(callback.headers.Location, config.publicOrigin)
  assert.match(callback.headers['Set-Cookie'], new RegExp(replacement))
  assert.equal(store.value.kind, 'authenticated')
  assert.equal(store.value.identity, `${config.oidcAuthority}\u0000subject-123`)
  assert.equal(store.value.accessToken, 'server-only-access-token')
  assert.equal(requests.length, 3)

  const session = new TestResponse()
  await oauth.session(
    { headers: { cookie: `${SESSION_COOKIE}=${replacement}` } },
    session,
  )
  assert.deepEqual(JSON.parse(session.body), {
    kind: 'authenticated',
    identity: `${config.oidcAuthority}\u0000subject-123`,
    csrfToken: store.value.csrfToken,
  })
  assert.doesNotMatch(session.body, /access|refresh|id_token/i)
})

test('destroys pending authorization when nonce validation fails', async () => {
  const store = new MemoryStore()
  const oauth = new OAuthClient(config, store, async (url) => {
    if (String(url).endsWith('/.well-known/openid-configuration'))
      return Response.json(metadata())
    if (String(url) === metadata().token_endpoint)
      return Response.json({
        access_token: 'access',
        id_token: idToken('wrong-nonce'),
        expires_in: 300,
      })
    if (String(url) === metadata().jwks_uri)
      return Response.json({ keys: [publicJwk] })
    throw new Error(`unexpected request: ${url}`)
  })
  await oauth.login({ headers: {} }, new TestResponse())
  const response = new TestResponse()

  await oauth.callback(
    { headers: { cookie: `${SESSION_COOKIE}=${identifier}` } },
    response,
    new URL(
      `/auth/callback?code=authorization-code&state=${store.value.state}`,
      config.publicOrigin,
    ),
  )

  assert.equal(response.status, 302)
  assert.equal(
    response.headers.Location,
    `${config.publicOrigin}/?authentication=failed`,
  )
  assert.equal(store.value, null)
})

test('reauthentication destroys the preceding browser session', async () => {
  const store = new MemoryStore()
  store.value = validSession()
  const oauth = new OAuthClient(config, store, async () =>
    Response.json(metadata()),
  )

  await oauth.login(
    { headers: { cookie: `${SESSION_COOKIE}=${replacement}` } },
    new TestResponse(),
  )

  assert.deepEqual(store.destroyed, [replacement])
  assert.equal(store.value.kind, 'authorization')
})

test('expired and malformed authenticated sessions are destroyed', async () => {
  const store = new MemoryStore()
  const oauth = new OAuthClient(config, store)
  store.value = { kind: 'authenticated', accessToken: 'incomplete' }
  assert.equal(
    await oauth.authenticated(
      { headers: { cookie: `${SESSION_COOKIE}=${identifier}` } },
      false,
    ),
    null,
  )
  store.value = { ...validSession(), accessExpiresAt: now - 1 }
  assert.equal(
    await oauth.authenticated(
      { headers: { cookie: `${SESSION_COOKIE}=${identifier}` } },
      false,
    ),
    null,
  )
  assert.deepEqual(store.destroyed, [identifier, identifier])
})

test('token renewal preserves the session identifier and CSRF value', async () => {
  const store = new MemoryStore()
  store.value = {
    ...validSession(),
    refreshToken: 'refresh-token',
    accessExpiresAt: now + 20,
  }
  const oauth = new OAuthClient(config, store)
  oauth.refresh = async (session) => ({
    ...session,
    accessToken: 'renewed-access-token',
    accessExpiresAt: now + 300,
  })

  const loaded = await oauth.authenticated({
    headers: { cookie: `${SESSION_COOKIE}=${identifier}` },
  })

  assert.equal(loaded.identifier, identifier)
  assert.equal(loaded.session.accessToken, 'renewed-access-token')
  assert.equal(loaded.session.csrfToken, 'csrf-value')
  assert.equal(store.replaced, 0)
})

test('concurrent renewal uses only an unexpired access token', async () => {
  const store = new MemoryStore()
  store.pending = true
  store.value = {
    ...validSession(),
    refreshToken: 'refresh-token',
    accessExpiresAt: now + 20,
  }
  const oauth = new OAuthClient(config, store)
  const request = {
    headers: { cookie: `${SESSION_COOKIE}=${identifier}` },
  }

  assert.equal(
    (await oauth.authenticated(request)).session.accessToken,
    'access-token',
  )
  store.value.accessExpiresAt = now - 1
  await assert.rejects(
    oauth.authenticated(request, false),
    SessionRenewalPendingError,
  )
  assert.deepEqual(store.destroyed, [])
})

test('logout destroys an expired session during owned renewal', async () => {
  const store = new MemoryStore()
  store.pending = true
  store.value = {
    ...validSession(),
    accessExpiresAt: now - 1,
  }
  const oauth = new OAuthClient(config, store)
  oauth.metadata = async () => ({})
  const response = new TestResponse()

  await oauth.logout(
    {
      headers: {
        cookie: `${SESSION_COOKIE}=${identifier}`,
        origin: config.publicOrigin,
        'sec-fetch-site': 'same-origin',
        'x-remora-csrf': 'csrf-value',
      },
    },
    response,
  )

  assert.equal(response.status, 200)
  assert.deepEqual(store.destroyed, [identifier])
  assert.match(response.headers['Set-Cookie'], /Max-Age=0/)
})

class MemoryStore {
  value = null
  destroyed = []
  replaced = 0
  pending = false

  async create(value) {
    this.value = value
    return identifier
  }

  async read() {
    return this.value
  }

  async replace(_identifier, value) {
    this.replaced += 1
    this.value = value
    return replacement
  }

  async renew(identifier, shouldRenew, renewal) {
    if (!shouldRenew(this.value)) return this.value
    if (this.pending) return SESSION_RENEWAL_PENDING
    const renewed = await renewal(this.value)
    if (renewed === null) {
      await this.destroy(identifier)
      return null
    }
    this.value = renewed
    return this.value
  }

  async destroy(identifier) {
    this.destroyed.push(identifier)
    this.value = null
  }
}

class TestResponse {
  status = 0
  headers = {}
  body = ''

  writeHead(status, headers) {
    this.status = status
    this.headers = headers
  }

  end(body = '') {
    this.body = body
  }
}

function metadata() {
  return {
    issuer: config.oidcAuthority,
    authorization_endpoint: `${config.oidcAuthority}/oauth/authorize`,
    token_endpoint: `${config.oidcAuthority}/oauth/token`,
    jwks_uri: `${config.oidcAuthority}/oauth/keys`,
  }
}

function validSession() {
  return {
    kind: 'authenticated',
    identity: `${config.oidcAuthority}\u0000subject-123`,
    csrfToken: 'csrf-value',
    accessToken: 'access-token',
    refreshToken: null,
    idToken: 'id-token',
    accessExpiresAt: now + 300,
  }
}

function idToken(nonce) {
  const header = encode({ alg: 'RS256', kid: publicJwk.kid })
  const payload = encode({
    iss: config.oidcAuthority,
    aud: config.oidcClientId,
    sub: 'subject-123',
    nonce,
    iat: now,
    exp: now + 300,
  })
  const signer = createSign('RSA-SHA256')
  signer.update(`${header}.${payload}`)
  signer.end()
  return `${header}.${payload}.${signer.sign(privateKey).toString('base64url')}`
}

function encode(value) {
  return Buffer.from(JSON.stringify(value)).toString('base64url')
}