import { Channel, ClientFactory, Status } from 'nice-grpc';
import { getGrpcErrorCode } from '@util/grpc-error';
import { EventEmitter, on } from 'events';
import {
  DispatcherClient as PbDispatcherClient,
  DispatcherDefinition,
  ResourceEventType,
  ResourceType,
  DispatcherClient,
  WorkflowEvent,
} from '@hatchet/protoc/dispatcher';
import { ClientConfig } from '@clients/hatchet-client/client-config';
import HatchetError from '@util/errors/hatchet-error';
import { Logger } from '@hatchet/util/logger';
import sleep from '@hatchet/util/sleep';
import { Api } from '../../rest';
import { RunGrpcPooledListener } from './pooled-child-listener-client';

const DEFAULT_EVENT_LISTENER_RETRY_INTERVAL = 5; // seconds
const DEFAULT_EVENT_LISTENER_RETRY_COUNT = 5;

export enum RunEventType {
  STEP_RUN_EVENT_TYPE_STARTED = 'STEP_RUN_EVENT_TYPE_STARTED',
  STEP_RUN_EVENT_TYPE_COMPLETED = 'STEP_RUN_EVENT_TYPE_COMPLETED',
  STEP_RUN_EVENT_TYPE_FAILED = 'STEP_RUN_EVENT_TYPE_FAILED',
  STEP_RUN_EVENT_TYPE_CANCELLED = 'STEP_RUN_EVENT_TYPE_CANCELLED',
  STEP_RUN_EVENT_TYPE_TIMED_OUT = 'STEP_RUN_EVENT_TYPE_TIMED_OUT',
  STEP_RUN_EVENT_TYPE_STREAM = 'STEP_RUN_EVENT_TYPE_STREAM',
  WORKFLOW_RUN_EVENT_TYPE_STARTED = 'WORKFLOW_RUN_EVENT_TYPE_STARTED',
  WORKFLOW_RUN_EVENT_TYPE_COMPLETED = 'WORKFLOW_RUN_EVENT_TYPE_COMPLETED',
  WORKFLOW_RUN_EVENT_TYPE_FAILED = 'WORKFLOW_RUN_EVENT_TYPE_FAILED',
  WORKFLOW_RUN_EVENT_TYPE_CANCELLED = 'WORKFLOW_RUN_EVENT_TYPE_CANCELLED',
  WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT = 'WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT',
}

const stepEventTypeMap: Record<ResourceEventType, RunEventType | undefined> = {
  [ResourceEventType.RESOURCE_EVENT_TYPE_STARTED]: RunEventType.STEP_RUN_EVENT_TYPE_STARTED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_COMPLETED]: RunEventType.STEP_RUN_EVENT_TYPE_COMPLETED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_FAILED]: RunEventType.STEP_RUN_EVENT_TYPE_FAILED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_CANCELLED]: RunEventType.STEP_RUN_EVENT_TYPE_CANCELLED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_TIMED_OUT]: RunEventType.STEP_RUN_EVENT_TYPE_TIMED_OUT,
  [ResourceEventType.RESOURCE_EVENT_TYPE_STREAM]: RunEventType.STEP_RUN_EVENT_TYPE_STREAM,
  [ResourceEventType.RESOURCE_EVENT_TYPE_UNKNOWN]: undefined,
  [ResourceEventType.UNRECOGNIZED]: undefined,
};

const workflowEventTypeMap: Record<ResourceEventType, RunEventType | undefined> = {
  [ResourceEventType.RESOURCE_EVENT_TYPE_STARTED]: RunEventType.WORKFLOW_RUN_EVENT_TYPE_STARTED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_COMPLETED]: RunEventType.WORKFLOW_RUN_EVENT_TYPE_COMPLETED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_FAILED]: RunEventType.WORKFLOW_RUN_EVENT_TYPE_FAILED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_CANCELLED]: RunEventType.WORKFLOW_RUN_EVENT_TYPE_CANCELLED,
  [ResourceEventType.RESOURCE_EVENT_TYPE_TIMED_OUT]: RunEventType.WORKFLOW_RUN_EVENT_TYPE_TIMED_OUT,
  [ResourceEventType.RESOURCE_EVENT_TYPE_UNKNOWN]: undefined,
  [ResourceEventType.RESOURCE_EVENT_TYPE_STREAM]: undefined,
  [ResourceEventType.UNRECOGNIZED]: undefined,
};

const resourceTypeMap: Record<
  ResourceType,
  Record<ResourceEventType, RunEventType | undefined> | undefined
> = {
  [ResourceType.RESOURCE_TYPE_STEP_RUN]: stepEventTypeMap,
  [ResourceType.RESOURCE_TYPE_WORKFLOW_RUN]: workflowEventTypeMap,
  [ResourceType.RESOURCE_TYPE_UNKNOWN]: undefined,
  [ResourceType.UNRECOGNIZED]: undefined,
};

export interface StepRunEvent {
  type: RunEventType;
  payload: string;
  resourceId: string;
  workflowRunId: string;
}

export class RunEventListener {
  client: DispatcherClient;

  q: Array<StepRunEvent> = [];
  eventEmitter = new EventEmitter();

  pollInterval: ReturnType<typeof setInterval> | undefined;

  constructor(client: DispatcherClient) {
    this.client = client;
  }

  static forRunId(workflowRunId: string, client: DispatcherClient): RunEventListener {
    const listener = new RunEventListener(client);
    listener.listenForRunId(workflowRunId);
    return listener;
  }

  static forAdditionalMeta(key: string, value: string, client: DispatcherClient): RunEventListener {
    const listener = new RunEventListener(client);
    listener.listenForAdditionalMeta(key, value);
    return listener;
  }

  emit(event: StepRunEvent) {
    this.q.push(event);
    this.eventEmitter.emit('event');
  }

  async listenForRunId(workflowRunId: string) {
    const listenerFactory = () =>
      this.client.subscribeToWorkflowEvents({
        workflowRunId,
      });

    return this.listenLoop(listenerFactory);
  }

  async listenForAdditionalMeta(key: string, value: string) {
    const listenerFactory = () =>
      this.client.subscribeToWorkflowEvents({
        additionalMetaKey: key,
        additionalMetaValue: value,
      });

    return this.listenLoop(listenerFactory);
  }

  async listenLoop(listenerFactory: () => AsyncIterable<WorkflowEvent>) {
    let listener = listenerFactory();

    while (true) {
      try {
        for await (const workflowEvent of listener) {
          const eventType = resourceTypeMap[workflowEvent.resourceType]?.[workflowEvent.eventType];
          if (eventType) {
            this.emit({
              type: eventType,
              payload: workflowEvent.eventPayload,
              resourceId: workflowEvent.resourceId,
              workflowRunId: workflowEvent.workflowRunId,
            });
          }
        }

        this.eventEmitter.emit('complete');
        return;
      } catch (e: unknown) {
        if (getGrpcErrorCode(e) === Status.CANCELLED) {
          this.eventEmitter.emit('complete');
          return;
        }
        if (getGrpcErrorCode(e) === Status.UNAVAILABLE) {
          listener = await this.retrySubscribe(listenerFactory);
        } else {
          throw e;
        }
      }
    }
  }

  async retrySubscribe(listenerFactory: () => AsyncIterable<WorkflowEvent>) {
    let retries = 0;

    while (retries < DEFAULT_EVENT_LISTENER_RETRY_COUNT) {
      try {
        await sleep(DEFAULT_EVENT_LISTENER_RETRY_INTERVAL);
        return listenerFactory();
      } catch {
        retries += 1;
      }
    }

    throw new HatchetError(
      `Could not subscribe to the worker after ${DEFAULT_EVENT_LISTENER_RETRY_COUNT} retries`
    );
  }

  async *stream(): AsyncGenerator<StepRunEvent, void, unknown> {
    let completed = false;

    this.eventEmitter.once('complete', () => {
      completed = true;
      this.eventEmitter.emit('event');
    });

    for await (const _ of on(this.eventEmitter, 'event')) {
      while (this.q.length > 0) {
        const r = this.q.shift();
        if (r) {
          yield r;
        }
      }

      if (completed && this.q.length === 0) {
        break;
      }
    }
  }
}

export class RunListenerClient {
  config: ClientConfig;
  client: PbDispatcherClient;
  logger: Logger;
  api: Api;

  pooledListener: RunGrpcPooledListener | undefined;

  constructor(config: ClientConfig, channel: Channel, factory: ClientFactory, api: Api) {
    this.config = config;
    this.client = factory.create(DispatcherDefinition, channel);
    this.logger = config.logger(`Listener`, config.log_level);
    this.api = api;
  }

  get(workflowRunId: string) {
    if (!this.pooledListener) {
      this.pooledListener = new RunGrpcPooledListener(this, () => {
        this.pooledListener = undefined;
      });
    }

    return this.pooledListener.subscribe({
      workflowRunId,
    });
  }

  async stream(workflowRunId: string): Promise<AsyncGenerator<StepRunEvent, void, unknown>> {
    const listener = RunEventListener.forRunId(workflowRunId, this.client);
    return listener.stream();
  }

  async streamByRunId(workflowRunId: string): Promise<AsyncGenerator<StepRunEvent, void, unknown>> {
    return this.stream(workflowRunId);
  }

  async streamByAdditionalMeta(
    key: string,
    value: string
  ): Promise<AsyncGenerator<StepRunEvent, void, unknown>> {
    const listener = RunEventListener.forAdditionalMeta(key, value, this.client);
    return listener.stream();
  }
}
