import type { APIPromise } from '../core/api-promise';
import type { RequestOptions } from '../internal/request-options';
import type { GenerateRunModelParams, GenerateRunModelResponse } from '../resources/generate';
import { Models, type ModelRetrieveParams, type ModelRetrieveResponse } from '../resources/models/models';
import { Train, type TrainTriggerParams, type TrainTriggerResponse } from '../resources/models/train';
import { type WithJob, enhanceJob } from './job';
import { type Scope, effectiveScope, withScope } from './scope';
// Type-only import to break the circular `scenario.ts` ↔ `models.ts` graph at runtime.
import type { Scenario } from './scenario';

/**
 * Enhanced Model entity with helper methods.
 * Has all original model fields plus `.run()`.
 *
 * @example
 * ```ts
 * const { model } = await client.models.retrieve(modelId);
 * const gen = await model.run({ body: { prompt: 'a red car' } });
 * const completed = await gen.job.wait();
 * ```
 */
export class ModelEntity {
  /** @internal */ declare readonly _client: Scenario;
  /** @internal scope captured from the retrieve call — replayed when `.run()` is invoked without an explicit query override. */
  declare readonly _scope?: Scope;

  /**
   * Generate with this model. Shortcut for
   * `client.generate.runModel(this.id, params)`.
   *
   * Returns an enhanced response where `response.job` has `.wait()`.
   * If the model was fetched with a `projectId` scope (per-call override or
   * client default), that scope is reused unless overridden by
   * `options.query`.
   */
  run(
    this: ModelEntity & ModelRetrieveResponse.Model,
    params: GenerateRunModelParams,
    options?: RequestOptions,
  ): APIPromise<WithJob<GenerateRunModelResponse>> {
    return this._client.generate.runModel(this.id, params, withScope(this._scope, options));
  }

  /** @internal Create a ModelEntity wrapping raw model data, optionally remembering the originating scope. */
  static from(
    client: Scenario,
    data: ModelRetrieveResponse.Model,
    scope?: Scope,
  ): ModelRetrieveResponse.Model & ModelEntity {
    const entity = Object.assign(Object.create(ModelEntity.prototype), data);
    Object.defineProperty(entity, '_client', { value: client, enumerable: false });
    if (scope) Object.defineProperty(entity, '_scope', { value: scope, enumerable: false });
    return entity;
  }
}

/** Enhanced retrieve response where `model` has `.run()` on top of its original fields. */
export type EnhancedModelRetrieveResponse = Omit<ModelRetrieveResponse, 'model'> & {
  model: ModelRetrieveResponse.Model & ModelEntity;
};

/**
 * Enhanced Train resource.
 * `trigger()` returns a response where `job` has `.wait()`.
 */
export class EnhancedTrain extends Train {
  override trigger(
    modelID: string,
    params: TrainTriggerParams,
    options?: RequestOptions,
  ): APIPromise<WithJob<TrainTriggerResponse>> {
    return enhanceJob(this._client, super.trigger(modelID, params, options), options);
  }
}

/**
 * Enhanced Models resource.
 * All original methods inherited.
 * - `train` is an {@link EnhancedTrain}
 * - `retrieve()` returns a response where `model` has `.run()`
 */
export class EnhancedModels extends Models {
  override train: EnhancedTrain = new EnhancedTrain(this._client);

  override retrieve(
    modelID: string,
    query: ModelRetrieveParams | null | undefined = {},
    options?: RequestOptions,
  ): APIPromise<EnhancedModelRetrieveResponse> {
    const client = this._client as Scenario;
    const scope = effectiveScope(options, client.projectId);
    return super.retrieve(modelID, query, options)._thenUnwrap((data) => ({
      ...data,
      model: ModelEntity.from(client, data.model, scope),
    }));
  }
}