import type { Scenario } from '../client';
import type { APIPromise } from '../core/api-promise';
import type { RequestOptions } from '../internal/request-options';
import {
  Uploads,
  type UploadCreateParams,
  type UploadCreateResponse,
  type UploadRetrieveParams,
  type UploadRetrieveResponse,
  type UploadTriggerActionParams,
  type UploadTriggerActionResponse,
} from '../resources/uploads';
import type { AssetRetrieveResponse } from '../resources/assets/assets';
import type { ModelRetrieveResponse } from '../resources/models/models';
import { type Scope, clientProjectId, effectiveScope, withScope } from './scope';

type UploadKind = NonNullable<UploadCreateParams['kind']>;

/** The resolved entity returned after a successful upload, narrowed by `kind`. */
export type UploadResult<K extends UploadKind> =
  K extends 'model' ? ModelRetrieveResponse : AssetRetrieveResponse;

/**
 * Accepted upload inputs.
 * - `string` — a local file path, read from disk (Node only; throws in browsers).
 * - `Uint8Array` — raw bytes. `Buffer` is a `Uint8Array` and works too.
 * - `Blob` / `File` — browser-native, or Node ≥ 18 via `fs.openAsBlob()` /
 *   constructors.
 */
export type UploadFileInput = string | Uint8Array | Blob;

export interface UploadFileParams<K extends UploadKind> {
  /**
   * The file to upload. Accepts a local path (Node), a `Uint8Array` / `Buffer`,
   * or a `Blob` / `File`.
   */
  file: UploadFileInput;
  /** Original filename (e.g. `"my-model.safetensors"`). */
  fileName: string;
  /** MIME type (e.g. `"application/octet-stream"`, `"image/png"`). */
  contentType: string;
  /** The kind of upload. Determines the returned entity type. */
  kind: K;
  /**
   * Extra asset options. Ignored when `kind === 'model'` — models don't
   * produce an asset, so these options are not applicable server-side.
   */
  assetOptions?: UploadCreateParams['assetOptions'];
  /**
   * Maximum number of parts uploaded to S3 in parallel.
   * @default 4
   */
  partConcurrency?: number;
  /**
   * Polling interval for the final "upload imported" check, in ms.
   * @default 2000
   */
  pollIntervalMs?: number;
  /**
   * Maximum time to wait for the upload to reach a terminal state, in ms.
   * Throws if exceeded.
   * @default 300000
   */
  pollTimeoutMs?: number;
}

// AWS S3 multipart upload limits.
const MIN_PART_SIZE = 6_000_000; // 6 MB (spec minimum is 5 MiB; a bit over to be safe)
const MAX_PART_SIZE = 5_368_709_120; // 5 GiB
const MAX_PARTS = 10_000;
const MAX_FILE_SIZE = 5_497_558_138_880; // 5 TiB

type UploadStatus = UploadRetrieveResponse.Upload['status'];

// `validated` is intermediate — the server has validated the upload but
// hasn't created the asset/model yet, so `entityId` isn't populated.
// Real terminals are `imported` (asset/model created), `complete`, and
// `failed`. Everything else means "keep polling".
const TERMINAL_STATUSES: Set<UploadStatus> = new Set(['complete', 'failed', 'imported']);

export interface UploadWaitOptions {
  /** Polling interval in milliseconds. Default: 2000 */
  intervalMs?: number;
  /** Maximum wait time in milliseconds. Default: 300000 */
  timeoutMs?: number;
  /**
   * Override the project scope used for the polling requests. Defaults to
   * the scope captured on the upload (per-call override on the originating
   * call, else the client's default project).
   */
  projectId?: string;
}

/** Upload methods added on top of the original upload fields. */
export class UploadMethods {
  /** @internal */ declare readonly _client: Scenario;
  /** @internal scope captured from the originating call — replayed on every `.wait()` poll. */
  declare readonly _scope?: Scope;

  /**
   * Poll until the upload has been processed into an entity (asset or model),
   * or reaches a terminal state (`imported`, `complete`, `failed`). Resolves
   * with the latest upload data so you can read `entityId`, `status`, or
   * `errorMessage` without a follow-up retrieve.
   *
   * @example
   * ```ts
   * const res = await client.uploads.retrieve(uploadId);
   * const done = await res.upload.wait();
   * if (done.status === 'failed') throw new Error(done.errorMessage);
   * console.log(done.entityId);
   * ```
   */
  async wait(this: Upload, options?: UploadWaitOptions): Promise<Upload> {
    const scope = waitOverrideScope(options) ?? this._scope;
    const final = await pollUntilImported(
      () => this._client.uploads.retrieve(this.id, undefined, withScope(scope)),
      options?.intervalMs ?? 2_000,
      options?.timeoutMs ?? 300_000,
    );
    return UploadMethods.from(this._client, final.upload, scope);
  }

  /** @internal Create an Upload from raw upload data, optionally remembering the originating scope. */
  static from(client: Scenario, data: UploadRetrieveResponse['upload'], scope?: Scope): Upload {
    const upload = Object.assign(Object.create(UploadMethods.prototype), data) as Upload;
    Object.defineProperty(upload, '_client', { value: client, enumerable: false });
    if (scope) Object.defineProperty(upload, '_scope', { value: scope, enumerable: false });
    return upload;
  }
}

/** An upload with all original fields plus `.wait()`. */
export type Upload = UploadRetrieveResponse.Upload & UploadMethods;
export const Upload = UploadMethods;

/**
 * @internal Helper type: adds `.wait()` to the `upload` field of a response.
 * Uses intersection so the original response type is preserved — a `WithUpload<T>` is still assignable to `T`.
 */
export type WithUpload<T extends { upload: unknown }> = T & { upload: UploadMethods };

/**
 * @internal Wrap `_thenUnwrap` to replace `response.upload` with an enhanced Upload.
 * Captures the effective scope (per-call override or client default) so
 * `.wait()` can replay it on every poll — keeping follow-up calls in sync
 * with the project the upload was originally created in.
 */
export function enhanceUpload<T extends { upload: { id: string; status: string } }>(
  client: Scenario,
  promise: APIPromise<T>,
  options?: RequestOptions,
): APIPromise<WithUpload<T>> {
  const scope = effectiveScope(options, clientProjectId(client));
  return promise._thenUnwrap((data) => ({
    ...data,
    upload: Upload.from(client, data.upload as UploadRetrieveResponse['upload'], scope),
  }));
}

/**
 * Enhanced Uploads resource.
 * All original methods (`create`, `retrieve`, `triggerAction`) are inherited,
 * with each response's `upload` field enriched with a `.wait()` helper.
 * Adds `uploadFile()` — a one-call wrapper around the 4-step upload flow:
 * init → PUT parts to S3 → trigger complete → poll until imported.
 */
export class EnhancedUploads extends Uploads {
  override create(
    body: UploadCreateParams,
    options?: RequestOptions,
  ): APIPromise<WithUpload<UploadCreateResponse>> {
    return enhanceUpload(this._client, super.create(body, options), options);
  }

  override retrieve(
    uploadID: string,
    query: UploadRetrieveParams | null | undefined = {},
    options?: RequestOptions,
  ): APIPromise<WithUpload<UploadRetrieveResponse>> {
    return enhanceUpload(this._client, super.retrieve(uploadID, query, options), options);
  }

  override triggerAction(
    uploadID: string,
    body: UploadTriggerActionParams,
    options?: RequestOptions,
  ): APIPromise<WithUpload<UploadTriggerActionResponse>> {
    return enhanceUpload(this._client, super.triggerAction(uploadID, body, options), options);
  }

  /**
   * Upload a file end-to-end. Computes an optimal part size (within AWS's
   * 5 MiB / 5 GiB / 10,000-part bounds), uploads all parts in parallel,
   * triggers completion, polls until the server has validated and imported
   * the file, and returns the resolved entity.
   *
   * The returned type narrows on `kind`: `'model'` returns a Model,
   * everything else returns an Asset.
   *
   * @example
   * ```ts
   * // From a file path (Node)
   * const { asset } = await client.uploads.uploadFile({
   *   file: './photo.jpg',
   *   fileName: 'photo.jpg',
   *   contentType: 'image/jpeg',
   *   kind: 'image',
   * });
   * console.log(asset.id, asset.url);
   *
   * // Or from bytes / Blob — anything you already have in memory
   * await client.uploads.uploadFile({
   *   file: new Blob([pngBytes], { type: 'image/png' }),
   *   fileName: 'generated.png',
   *   contentType: 'image/png',
   *   kind: 'image',
   * });
   * ```
   */
  async uploadFile<K extends UploadKind>(
    params: UploadFileParams<K>,
    options?: RequestOptions,
  ): Promise<UploadResult<K>> {
    const bytes = await normalizeToBytes(params.file);
    const fileSize = bytes.byteLength;

    if (fileSize === 0) {
      throw new Error('Cannot upload an empty file');
    }
    if (fileSize > MAX_FILE_SIZE) {
      throw new Error(
        `File size (${fileSize} bytes) exceeds the maximum allowed size of ${MAX_FILE_SIZE} bytes (5 TiB)`,
      );
    }

    // Pick the smallest legal part size that fits within MAX_PARTS.
    const requiredPartSize = Math.ceil(fileSize / MAX_PARTS);
    const partSize = Math.max(MIN_PART_SIZE, requiredPartSize);
    if (partSize > MAX_PART_SIZE) {
      throw new Error(
        `File size (${fileSize} bytes) requires a part size of ${partSize} bytes, ` +
          `which exceeds the maximum part size of ${MAX_PART_SIZE} bytes (5 GiB)`,
      );
    }
    const partsCount = Math.max(1, Math.ceil(fileSize / partSize));

    // Step 1 — init upload session and receive presigned S3 URLs.
    const createBody: UploadCreateParams = {
      fileName: params.fileName,
      fileSize,
      contentType: params.contentType,
      kind: params.kind,
      parts: partsCount,
      ...(params.kind !== 'model' && params.assetOptions ? { assetOptions: params.assetOptions } : {}),
    };
    const createRes: UploadCreateResponse = await this.create(createBody, options);
    const session = createRes.upload;
    const uploadId = session.id;
    const parts = session.parts;
    if (parts?.length !== partsCount) {
      throw new Error(`Expected ${partsCount} presigned part URL(s), got ${parts?.length ?? 0}`);
    }

    // Step 2 — PUT each chunk to its presigned S3 URL, with a concurrency cap.
    const concurrency = Math.max(1, params.partConcurrency ?? 4);
    await uploadPartsWithConcurrency(bytes, parts, partSize, concurrency);

    // Step 3 — tell the server all parts are uploaded.
    await this.triggerAction(uploadId, { action: 'complete' }, options);

    // Step 4 — poll until the upload has been processed into an entity.
    const final = await pollUntilImported(
      () => this.retrieve(uploadId, undefined, options),
      params.pollIntervalMs ?? 2_000,
      params.pollTimeoutMs ?? 300_000,
    );

    if (final.upload.status === 'failed') {
      throw new Error(`Upload ${uploadId} failed: ${final.upload.errorMessage ?? '<no error message>'}`);
    }
    const entityId = final.upload.entityId;
    if (!entityId) {
      throw new Error(`Upload ${uploadId} reached status "${final.upload.status}" but has no entityId`);
    }

    // Fetch and return the resolved entity (narrowed by kind). Forward the
    // original options so any per-call `projectId` query is passed through
    // to the final retrieve (necessary under bearer auth with per-call
    // scope overrides).
    if (params.kind === 'model') {
      const model = await this._client.models.retrieve(entityId, undefined, options);
      return model as UploadResult<K>;
    }
    const asset = await this._client.assets.retrieve(entityId, undefined, options);
    return asset as UploadResult<K>;
  }
}

/** Resolve the scope override carried by an `UploadWaitOptions`, if any. */
function waitOverrideScope(options?: UploadWaitOptions): Scope | undefined {
  if (!options?.projectId) return undefined;
  return { projectId: options.projectId };
}

// ─────────────────────── helpers ───────────────────────

async function normalizeToBytes(input: UploadFileInput): Promise<Uint8Array> {
  if (typeof input === 'string') {
    // Dynamic import so browser bundlers don't pull in node:fs.
    const { readFile } = await import('node:fs/promises');
    return new Uint8Array(await readFile(input));
  }
  if (input instanceof Uint8Array) {
    return input;
  }
  if (typeof Blob !== 'undefined' && input instanceof Blob) {
    return new Uint8Array(await input.arrayBuffer());
  }
  throw new Error('uploadFile: unsupported input type — expected a path string, Uint8Array, or Blob');
}

async function uploadPartsWithConcurrency(
  bytes: Uint8Array,
  parts: Array<{ number: number; url: string; expires: string }>,
  partSize: number,
  concurrency: number,
): Promise<void> {
  // Sort by part number so slicing lines up with the presigned URL order.
  const sorted = [...parts].sort((a, b) => a.number - b.number);

  let nextIndex = 0;
  const workers = Array.from({ length: concurrency }, () => worker());
  await Promise.all(workers);

  async function worker() {
    while (true) {
      const i = nextIndex++;
      if (i >= sorted.length) return;
      const part = sorted[i];
      if (!part) continue;
      const start = i * partSize;
      const end = Math.min(bytes.byteLength, start + partSize);
      const chunk = bytes.subarray(start, end);
      const res = await fetch(part.url, { method: 'PUT', body: chunk });
      if (!res.ok) {
        throw new Error(
          `Failed to upload part ${part.number}/${sorted.length}: ${res.status} ${res.statusText}`,
        );
      }
    }
  }
}

async function pollUntilImported(
  retrieve: () => Promise<UploadRetrieveResponse>,
  intervalMs: number,
  timeoutMs: number,
): Promise<UploadRetrieveResponse> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    const res = await retrieve();
    const { status } = res.upload;
    // Return as soon as we have an entityId OR the upload reached a terminal state.
    if (res.upload.entityId || TERMINAL_STATUSES.has(status)) {
      return res;
    }
    await new Promise((r) => setTimeout(r, intervalMs));
  }
  throw new Error(`Upload did not reach a terminal status within ${timeoutMs / 1000}s`);
}