import { NatsConnection, QueuedIterator } from '@nats-io/nats-core'
import { jetstream } from '@nats-io/jetstream'
import { KvEntry, Kvm, KvOptions, KvStatus, KvWatchEntry } from '@nats-io/kv'
import { assign, sendParent, setup } from 'xstate'
import {
  beginSubscriptionSync,
  detachSubscriptions,
  finishSubscriptionSync,
  hasPendingSync,
  updateSubscriptionTarget,
} from './subscriptions'
import { KvSubscriptionKey, KvSubscriptionConfig, kvConsolidateState } from '../actions/kv'
import {
  byteLength,
  createEmptyTrafficMetrics,
  NatsTrafficMetrics,
  NatsTrafficRecordEvent,
  NatsTrafficResetEvent,
  recordTrafficMetric,
  resetTrafficAll,
  resetTrafficDownload,
  resetTrafficUpload,
} from '../traffic'

// internal events and events from nats connection
type InternalEvents = { type: 'ERROR'; error: Error } | NatsTrafficRecordEvent

// events which can be sent to the machine from the user
export type ExternalEvents =
  | { type: 'KV.CONNECT'; connection: NatsConnection }
  | { type: 'KV.CONNECTED' }
  | { type: 'KV.DISCONNECTED' }
  | {
      type: 'KV.BUCKET_LIST'
      bucket?: string
      onResult: (result: KvStatus[] | string[] | { error: Error }) => void
    }
  | {
      type: 'KV.BUCKET_CREATE'
      bucket: string
      onResult: (result: { ok: true } | { ok: false } | { error: Error }) => void
    }
  | {
      type: 'KV.BUCKET_DELETE'
      bucket: string
      onResult: (result: { ok: true } | { ok: false } | { error: Error }) => void
    }
  | {
      type: 'KV.GET'
      bucket: string
      key: string
      onResult: (result: KvEntry | null | { error: Error }) => void
    }
  | {
      type: 'KV.PUT'
      bucket: string
      key: string
      value: any
      onResult: (result: { ok: true } | { error: Error }) => void
    }
  | {
      type: 'KV.DELETE'
      bucket: string
      key: string
      onResult: (result: { ok: true } | { error: Error }) => void
    }
  | { type: 'KV.SUBSCRIBE'; config: KvSubscriptionConfig }
  | { type: 'KV.UNSUBSCRIBE'; bucket: string; key: string }
  | { type: 'KV.UNSUBSCRIBE_ALL' }
  | NatsTrafficResetEvent

export type Events = InternalEvents | ExternalEvents

export interface Context {
  uid: string
  cachedConnection: NatsConnection | null
  cachedKvm: Kvm | null
  kvmOpts?: KvOptions
  subscriptions: Map<string, QueuedIterator<KvWatchEntry>>
  subscriptionConfigs: Map<string, KvSubscriptionConfig>
  syncRequired: number
  trafficMetrics: NatsTrafficMetrics
  error?: Error
}
export const kvManagerLogic = setup({
  types: {
    context: {} as Context,
    events: {} as Events,
  },
  actors: {
    kvConsolidateState: kvConsolidateState,
  },
  guards: {
    hasPendingSync,
  },
}).createMachine({
  /** @xstate-layout N4IgpgJg5mDOIC5QAoC2BDAxgCwJYDswBKAOlgFcAjAKzEwBcB9XCAGzAGIBlAVQCEAUgFEAwgBVGASQByksZICCAGUlchAbQAMAXUSgADgHtYuerkP49IAB6IATJs0kAnAFYAzAEY7rgDQgAT3tXTxJNAA53KOiYrwBfOP80LDxCUgoaOiYCU1x0VlwTfCgOa1h6dHowEnQAMyqAJ2RXRyIOZJwCYjIqWgZmfFz8woIoLV0kECMTMwsrWwQANk0AFhJPNy8ffyCEHzsSV0jY2M8EpIxOtJ7M-swLQgZIbn5hcUZePi4RACVJPg0Ois01yc0mC08jnCJAA7O4Vu4YX5AvYVjDDnZ4YjXOcQB1Ut0Mn0mPd8I8qhAXoJRBIeNJPt8-gDxsDjKDLODEEiDu5FjDPOEkTtEJ4Yc5cfiuulellGKTyc9Pm8JCIlEIFD8PvxGf9ARMDGzZhzQAs4a4SHY7ItXIttiiEN5oTjcfhDBA4FZJWlWTNzMabIgALSLYUIYMSy4E6W3bJsMA+9nzRArOyhy3QiInU4RlJSm7EgZDApFKAJo1JhArTzmzTV23I3bpsLHLNRM6JPGRvNE2XyrKQMt+iswpEkRaChv2eGHHNXQkyu4WWCGAoQSpgRjldeDsEmkWY0IrFYebyTvZomcdr3zmOMMANBqGBo7-0LFaLC0CoX26KHFuthIEiAA */
  initial: 'kv_idle',
  context: {
    uid: new Date().toISOString(),
    cachedConnection: null,
    cachedKvm: null,
    kvmOpts: undefined,
    subscriptions: new Map<string, QueuedIterator<KvWatchEntry>>(),
    subscriptionConfigs: new Map<string, KvSubscriptionConfig>(),
    syncRequired: 0,
    trafficMetrics: createEmptyTrafficMetrics(),
  },
  on: {
    'METRICS.RECORD': {
      actions: assign({
        trafficMetrics: ({ context, event }) => recordTrafficMetric(context.trafficMetrics, event),
      }),
    },
    'METRICS.RESET_UPLOAD': {
      actions: assign({
        trafficMetrics: ({ context, event }) =>
          resetTrafficUpload(context.trafficMetrics, event.source),
      }),
    },
    'METRICS.RESET_DOWNLOAD': {
      actions: assign({
        trafficMetrics: ({ context, event }) =>
          resetTrafficDownload(context.trafficMetrics, event.source),
      }),
    },
    'METRICS.RESET_ALL': {
      actions: assign({
        trafficMetrics: ({ context, event }) =>
          resetTrafficAll(context.trafficMetrics, event.source),
      }),
    },
    'KV.SUBSCRIBE': {
      actions: [
        assign(({ context, event, self }) => {
          if (event.type !== 'KV.SUBSCRIBE') return {}

          const config: KvSubscriptionConfig = {
            ...event.config,
            onDownloadBytes: (bytes) =>
              self.send({
                type: 'METRICS.RECORD',
                source: 'nats.kv',
                downloadBytes: bytes,
              }),
          }
          return updateSubscriptionTarget(
            context,
            KvSubscriptionKey.key(config.bucket, config.key),
            config,
          )
        }),
      ],
    },
    'KV.UNSUBSCRIBE': {
      actions: assign(({ context, event }) => {
        return updateSubscriptionTarget(context, KvSubscriptionKey.key(event.bucket, event.key))
      }),
    },
    'KV.UNSUBSCRIBE_ALL': {
      actions: assign({
        subscriptionConfigs: new Map(),
        syncRequired: ({ context }) => context.syncRequired + 1,
      }),
    },
  },
  states: {
    kv_idle: {
      on: {
        'KV.DISCONNECTED': {},
        'KV.CONNECT': {
          target: 'kv_check_sync',
          actions: [
            assign({
              cachedConnection: ({ event }) => event.connection,
              cachedKvm: ({ event }) => new Kvm(event.connection),
            }),
          ],
        },
      },
    },
    kv_disconnecting: {
      entry: [
        // dont close the connection here, it will be closed by the nats connection machine
        assign(({ context }) => ({
          ...detachSubscriptions(context),
          cachedConnection: null,
          cachedKvm: null,
        })),
        sendParent({ type: 'KV.DISCONNECTED' }),
      ],
      always: {
        target: 'kv_idle',
      },
    },
    kv_connected: {
      entry: [sendParent({ type: 'KV.CONNECTED' })],
      always: {
        target: 'kv_syncing',
        guard: 'hasPendingSync',
      },
      on: {
        'KV.CONNECT': {
          target: 'kv_check_sync',
          actions: [
            assign({
              cachedConnection: ({ event }) => event.connection,
              cachedKvm: ({ event }) => new Kvm(event.connection),
            }),
          ],
        },
        'KV.DISCONNECTED': {
          target: 'kv_disconnecting',
        },
        'KV.BUCKET_LIST': {
          actions: async ({ context, event }) => {
            try {
              if (!context.cachedKvm) {
                event.onResult({ error: new Error('KVM not initialized') })
                return
              }

              const results = []
              if (event.bucket) {
                const bucket = await context.cachedKvm.open(event.bucket)
                for await (const key of await bucket.keys()) {
                  results.push(key)
                }
                event.onResult(results)
              } else {
                for await (const status of await context.cachedKvm.list()) {
                  results.push(status)
                }
                event.onResult(results)
              }
            } catch (error) {
              event.onResult({ error: error as Error })
            }
          },
        },
        'KV.BUCKET_CREATE': {
          actions: async ({ context, event }) => {
            try {
              if (!context.cachedKvm) throw new Error('KVM not initialized')

              for await (const status of context.cachedKvm.list()) {
                if (status.bucket === event.bucket) {
                  event.onResult?.({ ok: false })
                  return
                }
              }

              await context.cachedKvm.create(event.bucket)
              event.onResult?.({ ok: true })
            } catch (error) {
              event.onResult?.({ error: error as Error })
            }
          },
        },
        'KV.BUCKET_DELETE': {
          actions: async ({ event }) => {
            try {
              try {
                // workaround: theres no delete bucket method on the kvm
                const connection = (event as any).connection as NatsConnection
                const js = jetstream(connection!)
                const jsm = await js.jetstreamManager()
                const res = await jsm.streams.delete(`KV_${event.bucket}`)
                event.onResult({ ok: res })
              } catch (streamError) {
                // Stream deletion might fail, but that's okay
                event.onResult({ ok: false })
              }
            } catch (error) {
              event.onResult({ error: error as Error })
            }
          },
        },
        'KV.GET': {
          actions: async ({ context, event }) => {
            try {
              const kv = await context.cachedKvm?.open(event.bucket)
              if (!kv) {
                event.onResult({ error: new Error(`Bucket '${event.bucket}' not found`) })
                return
              }
              const entry = await kv.get(event.key)
              event.onResult(entry)
            } catch (error) {
              event.onResult({ error: error as Error })
            }
          },
        },
        'KV.PUT': {
          actions: async ({ context, event, self }) => {
            try {
              const kv = await context.cachedKvm?.open(event.bucket)
              if (!kv) {
                event.onResult({ error: new Error(`Bucket '${event.bucket}' not found`) })
                return
              }

              self.send({
                type: 'METRICS.RECORD',
                source: 'nats.kv',
                uploadBytes: byteLength(event.value),
              })
              await kv.put(event.key, event.value)
              event.onResult({ ok: true })
            } catch (error) {
              event.onResult({ error: error as Error })
            }
          },
        },
        'KV.DELETE': {
          actions: async ({ context, event }) => {
            try {
              const kv = await context.cachedKvm?.open(event.bucket)
              if (!kv) {
                event.onResult({ error: new Error(`Bucket '${event.bucket}' not found`) })
                return
              }
              await kv.delete(event.key)
              event.onResult({ ok: true })
            } catch (error) {
              event.onResult({ error: error as Error })
            }
          },
        },
      },
    },
    kv_check_sync: {
      on: {
        'KV.DISCONNECTED': {
          target: 'kv_disconnecting',
        },
      },
      always: [
        {
          target: 'kv_syncing',
          guard: 'hasPendingSync',
        },
        {
          target: 'kv_connected',
        },
      ],
    },
    kv_syncing: {
      on: {
        'KV.DISCONNECTED': {
          target: 'kv_disconnecting',
        },
      },
      entry: [assign(() => beginSubscriptionSync())],
      invoke: {
        id: 'single-instance-sync',
        src: 'kvConsolidateState',
        input: ({ context }: { context: Context }) => ({
          kvm: context.cachedKvm!,
          connection: context.cachedConnection!,
          currentState: context.subscriptions,
          targetState: context.subscriptionConfigs,
        }),
        onDone: [
          {
            target: 'kv_error',
            guard: ({ event }) => !!event.output.error,
            actions: assign(({ context, event }) => finishSubscriptionSync(context, event.output)),
          },
          {
            target: 'kv_connected',
            actions: assign(({ context, event }) => finishSubscriptionSync(context, event.output)),
          },
        ],
        onError: {
          target: 'kv_error',
          actions: assign(({ context, event }) =>
            finishSubscriptionSync(context, {
              subscriptions: context.subscriptions,
              error: event.error instanceof Error ? event.error : new Error(String(event.error)),
            }),
          ),
        },
      },
    },
    kv_error: {
      on: {
        'KV.DISCONNECTED': {
          target: 'kv_disconnecting',
        },
        'KV.CONNECT': {
          target: 'kv_check_sync',
          actions: [
            assign({
              cachedConnection: ({ event }) => event.connection,
              cachedKvm: ({ event }) => new Kvm(event.connection),
            }),
          ],
        },
      },
    },
  },
})