import { IntegrationTaskKey } from "@trigger.dev/sdk";
import { Model, ModelVersion } from "{{ sdkPackage }}";

import { {{ identifier | capitalize }}RunTask } from "./index";
import { modelProperties } from "./utils";
import { {{ identifier | capitalize }}ReturnType } from "./types";

export class Models {
  constructor(private runTask: {{ identifier | capitalize }}RunTask) {}

  get(
    key: IntegrationTaskKey,
    params: {
      model_owner: string;
      model_name: string;
    }
  ): {{ identifier | capitalize }}ReturnType<Model> {
    return this.runTask(
      key,
      (client) => {
        return client.models.get(params.model_owner, params.model_name);
      },
      {
        name: "Get Model",
        params,
        properties: modelProperties(params),
      }
    );
  }

  get versions() {
    return new Versions(this.runTask);
  }
}

class Versions {
  constructor(private runTask: {{ identifier | capitalize }}RunTask) {}

  get(
    key: IntegrationTaskKey,
    params: {
      model_owner: string;
      model_name: string;
      version_id: string;
    }
  ): {{ identifier | capitalize }}ReturnType<ModelVersion> {
    return this.runTask(
      key,
      (client) => {
        return client.models.versions.get(params.model_owner, params.model_name, params.version_id);
      },
      {
        name: "Get Model Version",
        params,
        properties: modelProperties(params),
      }
    );
  }

  list(
    key: IntegrationTaskKey,
    params: {
      model_owner: string;
      model_name: string;
    }
  ): {{ identifier | capitalize }}ReturnType<ModelVersion[]> {
    return this.runTask(
      key,
      (client) => {
        return client.models.versions.list(params.model_owner, params.model_name);
      },
      {
        name: "List Models",
        params,
        properties: modelProperties(params),
      }
    );
  }
}