import type { FunctionBinding } from '@astrale-os/kernel-api/routed'
import type { FunctionSchema } from '@astrale-os/kernel-core/domain'

import { deriveAllowedAlgorithms } from '@astrale-os/kernel-core'
import { functionBindingSchema, functionOutputSchema } from '@astrale-os/kernel-core/domain'
import { importJWK, jwtVerify } from 'jose'
import { z } from 'zod'

import type { RemoteIdentityConfig } from '../auth/identity.js'

import { signCredential } from '../auth/sign.js'

export const SERVICE_FUNCTIONS_PATH = '/_astrale/functions'
export const SERVICE_FUNCTIONS_PURPOSE = 'astrale.service-functions.v1'

const JsonSchema = z.record(z.string(), z.unknown())

export const ServiceFunctionManifestEntrySchema = z.object({
  slug: z.string().regex(/^[a-z][a-z0-9-]*$/),
  ref: z.string().regex(/^function\.[a-z][a-z0-9-]*$/),
  inputSchema: JsonSchema,
  outputSchema: JsonSchema,
  output: functionOutputSchema.default('value'),
  binding: functionBindingSchema.unwrap(),
})

export const ServiceFunctionManifestSchema = z.object({
  version: z.literal(1),
  service: z.object({ issuer: z.string().min(1), subject: z.string().min(1) }),
  functions: z.array(ServiceFunctionManifestEntrySchema).max(500),
})

export const ServiceFunctionDiscoveryRequestSchema = z.object({
  kernelIssuer: z.string().url(),
  nonce: z.string().min(16).max(256),
})

export const ServiceFunctionDiscoveryResponseSchema = ServiceFunctionManifestSchema.extend({
  identity: z.object({ credential: z.string().min(1) }),
})

export type ServiceFunctionManifestEntry = z.infer<typeof ServiceFunctionManifestEntrySchema>
export type ServiceFunctionManifest = z.infer<typeof ServiceFunctionManifestSchema>
export type ServiceFunctionDiscoveryRequest = z.infer<typeof ServiceFunctionDiscoveryRequestSchema>
export type ServiceFunctionDiscoveryResponse = z.infer<
  typeof ServiceFunctionDiscoveryResponseSchema
>

export function toServiceFunctionManifestEntries(
  schemas: readonly FunctionSchema[],
): ServiceFunctionManifestEntry[] {
  return schemas
    .map((schema) => {
      const slug = schema.ref.slice('function.'.length)
      return ServiceFunctionManifestEntrySchema.parse({
        slug,
        ref: schema.ref,
        inputSchema: schema.inputSchema,
        outputSchema: schema.outputSchema,
        output: schema.output ?? 'value',
        binding: schema.binding,
      })
    })
    .sort((a, b) => a.slug.localeCompare(b.slug))
}

export async function signServiceFunctionManifest(
  manifest: ServiceFunctionManifest,
  request: ServiceFunctionDiscoveryRequest,
  identity: RemoteIdentityConfig,
): Promise<ServiceFunctionDiscoveryResponse> {
  const parsedManifest = ServiceFunctionManifestSchema.parse(manifest)
  const manifestHash = await hashServiceFunctionManifest(parsedManifest)
  const credential = await signCredential(
    { purpose: SERVICE_FUNCTIONS_PURPOSE, nonce: request.nonce, manifestHash },
    {
      issuer: identity.issuer,
      subject: identity.subject,
      audience: request.kernelIssuer,
      privateKey: identity.privateKey,
      ttl: '2m',
    },
  )
  return { ...parsedManifest, identity: { credential } }
}

export async function verifyServiceFunctionManifest(
  input: unknown,
  expected: {
    kernelIssuer: string
    nonce: string
    serviceIssuer: string
    serviceSubject: string
    publicKey: JsonWebKey
  },
): Promise<ServiceFunctionDiscoveryResponse> {
  const response = ServiceFunctionDiscoveryResponseSchema.parse(input)
  if (
    response.service.issuer !== expected.serviceIssuer ||
    response.service.subject !== expected.serviceSubject
  ) {
    throw new Error('service function manifest identity does not match the deployed Service')
  }

  const algorithms = deriveAllowedAlgorithms(expected.publicKey)
  const algorithm = algorithms[0]
  if (!algorithm) throw new Error('service function manifest public key has no supported algorithm')
  const key = await importJWK(expected.publicKey, algorithm)
  const verified = await jwtVerify(response.identity.credential, key, {
    algorithms,
    issuer: expected.serviceIssuer,
    subject: expected.serviceSubject,
    audience: expected.kernelIssuer,
  })
  const expectedHash = await hashServiceFunctionManifest({
    version: response.version,
    service: response.service,
    functions: response.functions,
  })
  if (verified.payload.purpose !== SERVICE_FUNCTIONS_PURPOSE) {
    throw new Error('service function manifest credential has the wrong purpose')
  }
  if (verified.payload.nonce !== expected.nonce) {
    throw new Error('service function manifest credential has the wrong nonce')
  }
  if (verified.payload.manifestHash !== expectedHash) {
    throw new Error('service function manifest credential does not cover the returned manifest')
  }
  return response
}

export async function hashServiceFunctionManifest(
  manifest: ServiceFunctionManifest,
): Promise<string> {
  const bytes = new TextEncoder().encode(
    canonicalJson(ServiceFunctionManifestSchema.parse(manifest)),
  )
  const digest = await crypto.subtle.digest('SHA-256', bytes)
  return `sha256:${hex(new Uint8Array(digest))}`
}

export function functionBindingForManifest(
  binding: FunctionBinding,
  auth: 'required' | 'optional' | 'public' | undefined,
): FunctionBinding {
  return { ...binding, auth: auth ?? 'required' }
}

function canonicalJson(value: unknown): string {
  if (value === null || typeof value !== 'object') return JSON.stringify(value)
  if (Array.isArray(value)) return `[${value.map(canonicalJson).join(',')}]`
  const entries = Object.entries(value as Record<string, unknown>)
    .filter(([, item]) => item !== undefined)
    .sort(([a], [b]) => a.localeCompare(b))
  return `{${entries
    .map(([key, item]) => `${JSON.stringify(key)}:${canonicalJson(item)}`)
    .join(',')}}`
}

function hex(bytes: Uint8Array): string {
  let value = ''
  for (const byte of bytes) value += byte.toString(16).padStart(2, '0')
  return value
}