import ky, { KyInstance } from "ky";
import { DEFAULT_KIE_BASE_URL, ENDPOINTS } from "./endpoints";
import {
  jobSchema,
  KieJob,
  KieModel,
  modelSchema,
  queuePayloadSchema,
  QueuePayload,
  taskSchema
} from "./types";
import { z } from "zod";

const modelsResponseSchema = z.object({
  data: z.array(modelSchema)
});

const jobsResponseSchema = z.object({
  data: z.array(jobSchema)
});

type ClientOptions = {
  apiKey: string;
  baseUrl?: string;
};

export class KieClient {
  private http: KyInstance;

  constructor({ apiKey, baseUrl = DEFAULT_KIE_BASE_URL }: ClientOptions) {
    this.http = ky.create({
      prefixUrl: baseUrl.replace(/\/$/, ""),
      headers: {
        Authorization: `Bearer ${apiKey}`,
        "Content-Type": "application/json"
      },
      timeout: 30_000
    });
  }

  async listModels(): Promise<KieModel[]> {
    const res = await this.http.get(ENDPOINTS.models).json();
    const parsed = modelsResponseSchema.parse(res);
    return parsed.data;
  }

  async getModel(id: string): Promise<KieModel> {
    const res = await this.http.get(`${ENDPOINTS.models}/${id}`).json();
    return modelSchema.parse(res);
  }

  async listTasks() {
    const res = await this.http.get(ENDPOINTS.tasks).json();
    return z.array(taskSchema).parse(res);
  }

  async queueJob(payload: QueuePayload): Promise<KieJob> {
    const body = queuePayloadSchema.parse(payload);
    const res = await this.http.post(ENDPOINTS.jobs, { json: body }).json();
    return jobSchema.parse(res);
  }

  async getJob(id: string): Promise<KieJob> {
    const res = await this.http.get(`${ENDPOINTS.jobs}/${id}`).json();
    return jobSchema.parse(res);
  }

  async listJobs(): Promise<KieJob[]> {
    const res = await this.http.get(ENDPOINTS.jobs).json();
    const parsed = jobsResponseSchema.parse(res);
    return parsed.data;
  }

  async cancelJob(id: string): Promise<void> {
    await this.http.delete(`${ENDPOINTS.jobs}/${id}`);
  }
}