// File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.

import { APIResource } from '../../resource';
import { isRequestOptions } from '../../core';
import * as Core from '../../core';
import * as WorkPoolsAPI from './work-pools';
import * as Shared from '../shared';
import * as QueuesAPI from './queues';
import * as WorkersAPI from './workers';

export class WorkPools extends APIResource {
  queues: QueuesAPI.Queues = new QueuesAPI.Queues(this._client);
  workers: WorkersAPI.Workers = new WorkersAPI.Workers(this._client);

  /**
   * Creates a new work pool. If a work pool with the same name already exists, an
   * error will be raised.
   */
  create(params: WorkPoolCreateParams, options?: Core.RequestOptions): Core.APIPromise<WorkPool> {
    const { 'x-prefect-api-version': xPrefectAPIVersion, ...body } = params;
    return this._client.post('/api/work_pools/', {
      body,
      ...options,
      headers: {
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Read a work pool by name
   */
  retrieve(
    name: string,
    params?: WorkPoolRetrieveParams,
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPool>;
  retrieve(name: string, options?: Core.RequestOptions): Core.APIPromise<WorkPool>;
  retrieve(
    name: string,
    params: WorkPoolRetrieveParams | Core.RequestOptions = {},
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPool> {
    if (isRequestOptions(params)) {
      return this.retrieve(name, {}, params);
    }
    const { 'x-prefect-api-version': xPrefectAPIVersion } = params;
    return this._client.get(`/api/work_pools/${name}`, {
      ...options,
      headers: {
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Update a work pool
   */
  update(name: string, params: WorkPoolUpdateParams, options?: Core.RequestOptions): Core.APIPromise<void> {
    const { 'x-prefect-api-version': xPrefectAPIVersion, ...body } = params;
    return this._client.patch(`/api/work_pools/${name}`, {
      body,
      ...options,
      headers: {
        Accept: '*/*',
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Delete a work pool
   */
  delete(name: string, params?: WorkPoolDeleteParams, options?: Core.RequestOptions): Core.APIPromise<void>;
  delete(name: string, options?: Core.RequestOptions): Core.APIPromise<void>;
  delete(
    name: string,
    params: WorkPoolDeleteParams | Core.RequestOptions = {},
    options?: Core.RequestOptions,
  ): Core.APIPromise<void> {
    if (isRequestOptions(params)) {
      return this.delete(name, {}, params);
    }
    const { 'x-prefect-api-version': xPrefectAPIVersion } = params;
    return this._client.delete(`/api/work_pools/${name}`, {
      ...options,
      headers: {
        Accept: '*/*',
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Count work pools
   */
  count(params?: WorkPoolCountParams, options?: Core.RequestOptions): Core.APIPromise<WorkPoolCountResponse>;
  count(options?: Core.RequestOptions): Core.APIPromise<WorkPoolCountResponse>;
  count(
    params: WorkPoolCountParams | Core.RequestOptions = {},
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolCountResponse> {
    if (isRequestOptions(params)) {
      return this.count({}, params);
    }
    const { 'x-prefect-api-version': xPrefectAPIVersion, ...body } = params;
    return this._client.post('/api/work_pools/count', {
      body,
      ...options,
      headers: {
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Load scheduled runs for a worker
   */
  getScheduledFlowRuns(
    name: string,
    params?: WorkPoolGetScheduledFlowRunsParams,
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolGetScheduledFlowRunsResponse>;
  getScheduledFlowRuns(
    name: string,
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolGetScheduledFlowRunsResponse>;
  getScheduledFlowRuns(
    name: string,
    params: WorkPoolGetScheduledFlowRunsParams | Core.RequestOptions = {},
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolGetScheduledFlowRunsResponse> {
    if (isRequestOptions(params)) {
      return this.getScheduledFlowRuns(name, {}, params);
    }
    const { 'x-prefect-api-version': xPrefectAPIVersion, ...body } = params;
    return this._client.post(`/api/work_pools/${name}/get_scheduled_flow_runs`, {
      body,
      ...options,
      headers: {
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }

  /**
   * Read multiple work pools
   */
  listFiltered(
    params?: WorkPoolListFilteredParams,
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolListFilteredResponse>;
  listFiltered(options?: Core.RequestOptions): Core.APIPromise<WorkPoolListFilteredResponse>;
  listFiltered(
    params: WorkPoolListFilteredParams | Core.RequestOptions = {},
    options?: Core.RequestOptions,
  ): Core.APIPromise<WorkPoolListFilteredResponse> {
    if (isRequestOptions(params)) {
      return this.listFiltered({}, params);
    }
    const { 'x-prefect-api-version': xPrefectAPIVersion, ...body } = params;
    return this._client.post('/api/work_pools/filter', {
      body,
      ...options,
      headers: {
        ...(xPrefectAPIVersion != null ? { 'x-prefect-api-version': xPrefectAPIVersion } : undefined),
        ...options?.headers,
      },
    });
  }
}

/**
 * An ORM representation of a work pool
 */
export interface WorkPool {
  /**
   * The name of the work pool.
   */
  name: string;

  /**
   * The work pool type.
   */
  type: string;

  id?: string;

  /**
   * The work pool's base job template.
   */
  base_job_template?: unknown;

  /**
   * A concurrency limit for the work pool.
   */
  concurrency_limit?: number | null;

  created?: string | null;

  /**
   * The id of the pool's default queue.
   */
  default_queue_id?: string | null;

  /**
   * A description of the work pool.
   */
  description?: string | null;

  /**
   * Pausing the work pool stops the delivery of all work.
   */
  is_paused?: boolean;

  /**
   * Enumeration of work pool statuses.
   */
  status?: 'READY' | 'NOT_READY' | 'PAUSED' | null;

  updated?: string | null;
}

export interface WorkerFlowRunResponse {
  /**
   * An ORM representation of flow run data.
   */
  flow_run: WorkerFlowRunResponse.FlowRun;

  work_pool_id: string;

  work_queue_id: string;
}

export namespace WorkerFlowRunResponse {
  /**
   * An ORM representation of flow run data.
   */
  export interface FlowRun {
    /**
     * The id of the flow being run.
     */
    flow_id: string;

    id?: string;

    /**
     * Whether or not the flow run was automatically scheduled.
     */
    auto_scheduled?: boolean;

    /**
     * Additional context for the flow run.
     */
    context?: unknown;

    created?: string | null;

    /**
     * Optional information about the creator of this flow run.
     */
    created_by?: FlowRun.CreatedBy | null;

    /**
     * The id of the deployment associated with this flow run, if available.
     */
    deployment_id?: string | null;

    /**
     * The version of the deployment associated with this flow run.
     */
    deployment_version?: string | null;

    /**
     * Defines of how a flow run should retry.
     */
    empirical_policy?: FlowRun.EmpiricalPolicy;

    /**
     * The actual end time.
     */
    end_time?: string | null;

    /**
     * A real-time estimate of the total run time.
     */
    estimated_run_time?: number;

    /**
     * The difference between actual and expected start time.
     */
    estimated_start_time_delta?: number;

    /**
     * The flow run's expected start time.
     */
    expected_start_time?: string | null;

    /**
     * The version of the flow executed in this flow run.
     */
    flow_version?: string | null;

    /**
     * An optional idempotency key for the flow run. Used to ensure the same flow run
     * is not created multiple times.
     */
    idempotency_key?: string | null;

    /**
     * The block document defining infrastructure to use this flow run.
     */
    infrastructure_document_id?: string | null;

    /**
     * The id of the flow run as returned by an infrastructure block.
     */
    infrastructure_pid?: string | null;

    /**
     * Variables used as overrides in the base job template
     */
    job_variables?: unknown | null;

    /**
     * The name of the flow run. Defaults to a random slug if not specified.
     */
    name?: string;

    /**
     * The next time the flow run is scheduled to start.
     */
    next_scheduled_start_time?: string | null;

    /**
     * Parameters for the flow run.
     */
    parameters?: unknown;

    /**
     * If the flow run is a subflow, the id of the 'dummy' task in the parent flow used
     * to track subflow state.
     */
    parent_task_run_id?: string | null;

    /**
     * The number of times the flow run was executed.
     */
    run_count?: number;

    /**
     * The actual start time.
     */
    start_time?: string | null;

    /**
     * Represents the state of a run.
     */
    state?: Shared.State | null;

    /**
     * The id of the flow run's current state.
     */
    state_id?: string | null;

    /**
     * The name of the current flow run state.
     */
    state_name?: string | null;

    /**
     * Enumeration of state types.
     */
    state_type?:
      | 'SCHEDULED'
      | 'PENDING'
      | 'RUNNING'
      | 'COMPLETED'
      | 'FAILED'
      | 'CANCELLED'
      | 'CRASHED'
      | 'PAUSED'
      | 'CANCELLING'
      | null;

    /**
     * A list of tags on the flow run
     */
    tags?: Array<string>;

    /**
     * Total run time. If the flow run was executed multiple times, the time of each
     * run will be summed.
     */
    total_run_time?: number;

    updated?: string | null;

    /**
     * The id of the run's work pool queue.
     */
    work_queue_id?: string | null;

    /**
     * The work queue that handled this flow run.
     */
    work_queue_name?: string | null;
  }

  export namespace FlowRun {
    /**
     * Optional information about the creator of this flow run.
     */
    export interface CreatedBy {
      /**
       * The id of the creator of the object.
       */
      id?: string | null;

      /**
       * The display value for the creator.
       */
      display_value?: string | null;

      /**
       * The type of the creator of the object.
       */
      type?: string | null;
    }

    /**
     * Defines of how a flow run should retry.
     */
    export interface EmpiricalPolicy {
      /**
       * @deprecated: The maximum number of retries. Field is not used. Please use
       * `retries` instead.
       */
      max_retries?: number;

      /**
       * Tracks pauses this run has observed.
       */
      pause_keys?: Array<unknown> | null;

      /**
       * Indicates if this run is resuming from a pause.
       */
      resuming?: boolean | null;

      /**
       * The number of retries.
       */
      retries?: number | null;

      /**
       * The delay time between retries, in seconds.
       */
      retry_delay?: number | null;

      /**
       * @deprecated: The delay between retries. Field is not used. Please use
       * `retry_delay` instead.
       */
      retry_delay_seconds?: number;
    }
  }
}

export type WorkPoolCountResponse = number;

export type WorkPoolGetScheduledFlowRunsResponse = Array<WorkerFlowRunResponse>;

export type WorkPoolListFilteredResponse = Array<WorkPool>;

export interface WorkPoolCreateParams {
  /**
   * Body param: The name of the work pool.
   */
  name: string;

  /**
   * Body param: The work pool's base job template.
   */
  base_job_template?: unknown;

  /**
   * Body param: A concurrency limit for the work pool.
   */
  concurrency_limit?: number | null;

  /**
   * Body param: The work pool description.
   */
  description?: string | null;

  /**
   * Body param: Pausing the work pool stops the delivery of all work.
   */
  is_paused?: boolean;

  /**
   * Body param: The work pool type.
   */
  type?: string;

  /**
   * Header param:
   */
  'x-prefect-api-version'?: string;
}

export interface WorkPoolRetrieveParams {
  'x-prefect-api-version'?: string;
}

export interface WorkPoolUpdateParams {
  /**
   * Body param:
   */
  base_job_template?: unknown | null;

  /**
   * Body param:
   */
  concurrency_limit?: number | null;

  /**
   * Body param:
   */
  description?: string | null;

  /**
   * Body param:
   */
  is_paused?: boolean | null;

  /**
   * Header param:
   */
  'x-prefect-api-version'?: string;
}

export interface WorkPoolDeleteParams {
  'x-prefect-api-version'?: string;
}

export interface WorkPoolCountParams {
  /**
   * Body param: Filter work pools. Only work pools matching all criteria will be
   * returned
   */
  work_pools?: WorkPoolCountParams.WorkPools | null;

  /**
   * Header param:
   */
  'x-prefect-api-version'?: string;
}

export namespace WorkPoolCountParams {
  /**
   * Filter work pools. Only work pools matching all criteria will be returned
   */
  export interface WorkPools {
    /**
     * Filter by `WorkPool.id`.
     */
    id?: WorkPools.ID | null;

    /**
     * Filter by `WorkPool.name`.
     */
    name?: WorkPools.Name | null;

    /**
     * Operator for combining filter criteria. Defaults to 'and\_'.
     */
    operator?: 'and_' | 'or_';

    /**
     * Filter by `WorkPool.type`.
     */
    type?: WorkPools.Type | null;
  }

  export namespace WorkPools {
    /**
     * Filter by `WorkPool.id`.
     */
    export interface ID {
      /**
       * A list of work pool ids to include
       */
      any_?: Array<string> | null;
    }

    /**
     * Filter by `WorkPool.name`.
     */
    export interface Name {
      /**
       * A list of work pool names to include
       */
      any_?: Array<string> | null;
    }

    /**
     * Filter by `WorkPool.type`.
     */
    export interface Type {
      /**
       * A list of work pool types to include
       */
      any_?: Array<string> | null;
    }
  }
}

export interface WorkPoolGetScheduledFlowRunsParams {
  /**
   * Body param: Defaults to PREFECT_API_DEFAULT_LIMIT if not provided.
   */
  limit?: number;

  /**
   * Body param: The minimum time to look for scheduled flow runs
   */
  scheduled_after?: string;

  /**
   * Body param: The maximum time to look for scheduled flow runs
   */
  scheduled_before?: string;

  /**
   * Body param: The names of work pool queues
   */
  work_queue_names?: Array<string>;

  /**
   * Header param:
   */
  'x-prefect-api-version'?: string;
}

export interface WorkPoolListFilteredParams {
  /**
   * Body param: Defaults to PREFECT_API_DEFAULT_LIMIT if not provided.
   */
  limit?: number;

  /**
   * Body param:
   */
  offset?: number;

  /**
   * Body param: Filter work pools. Only work pools matching all criteria will be
   * returned
   */
  work_pools?: WorkPoolListFilteredParams.WorkPools | null;

  /**
   * Header param:
   */
  'x-prefect-api-version'?: string;
}

export namespace WorkPoolListFilteredParams {
  /**
   * Filter work pools. Only work pools matching all criteria will be returned
   */
  export interface WorkPools {
    /**
     * Filter by `WorkPool.id`.
     */
    id?: WorkPools.ID | null;

    /**
     * Filter by `WorkPool.name`.
     */
    name?: WorkPools.Name | null;

    /**
     * Operator for combining filter criteria. Defaults to 'and\_'.
     */
    operator?: 'and_' | 'or_';

    /**
     * Filter by `WorkPool.type`.
     */
    type?: WorkPools.Type | null;
  }

  export namespace WorkPools {
    /**
     * Filter by `WorkPool.id`.
     */
    export interface ID {
      /**
       * A list of work pool ids to include
       */
      any_?: Array<string> | null;
    }

    /**
     * Filter by `WorkPool.name`.
     */
    export interface Name {
      /**
       * A list of work pool names to include
       */
      any_?: Array<string> | null;
    }

    /**
     * Filter by `WorkPool.type`.
     */
    export interface Type {
      /**
       * A list of work pool types to include
       */
      any_?: Array<string> | null;
    }
  }
}

export namespace WorkPools {
  export import WorkPool = WorkPoolsAPI.WorkPool;
  export import WorkerFlowRunResponse = WorkPoolsAPI.WorkerFlowRunResponse;
  export import WorkPoolCountResponse = WorkPoolsAPI.WorkPoolCountResponse;
  export import WorkPoolGetScheduledFlowRunsResponse = WorkPoolsAPI.WorkPoolGetScheduledFlowRunsResponse;
  export import WorkPoolListFilteredResponse = WorkPoolsAPI.WorkPoolListFilteredResponse;
  export import WorkPoolCreateParams = WorkPoolsAPI.WorkPoolCreateParams;
  export import WorkPoolRetrieveParams = WorkPoolsAPI.WorkPoolRetrieveParams;
  export import WorkPoolUpdateParams = WorkPoolsAPI.WorkPoolUpdateParams;
  export import WorkPoolDeleteParams = WorkPoolsAPI.WorkPoolDeleteParams;
  export import WorkPoolCountParams = WorkPoolsAPI.WorkPoolCountParams;
  export import WorkPoolGetScheduledFlowRunsParams = WorkPoolsAPI.WorkPoolGetScheduledFlowRunsParams;
  export import WorkPoolListFilteredParams = WorkPoolsAPI.WorkPoolListFilteredParams;
  export import Queues = QueuesAPI.Queues;
  export import QueueFilterResponse = QueuesAPI.QueueFilterResponse;
  export import QueueCreateParams = QueuesAPI.QueueCreateParams;
  export import QueueRetrieveParams = QueuesAPI.QueueRetrieveParams;
  export import QueueUpdateParams = QueuesAPI.QueueUpdateParams;
  export import QueueDeleteParams = QueuesAPI.QueueDeleteParams;
  export import QueueFilterParams = QueuesAPI.QueueFilterParams;
  export import Workers = WorkersAPI.Workers;
  export import WorkerResponse = WorkersAPI.WorkerResponse;
  export import WorkerFilterResponse = WorkersAPI.WorkerFilterResponse;
  export import WorkerDeleteParams = WorkersAPI.WorkerDeleteParams;
  export import WorkerFilterParams = WorkersAPI.WorkerFilterParams;
  export import WorkerHeartbeatParams = WorkersAPI.WorkerHeartbeatParams;
}