import { io, type Socket } from "socket.io-client";
import type {
  ChatEvent,
  ChatHistoryPayload,
  LiveRuntimeInfo,
  ModelListPayload,
  ModelRef,
  PersistedSessionInfo,
  RuntimeCreatedPayload,
} from "./types";

export type SocketCallbacks = {
  onConnect?: () => void;
  onDisconnect?: () => void;
  onError?: (message: string) => void;
  onSessionsList?: (payload: {
    persisted: PersistedSessionInfo[];
    live: LiveRuntimeInfo[];
  }) => void;
  onRuntimeCreated?: (payload: RuntimeCreatedPayload) => void;
  onRuntimeOpened?: (payload: RuntimeCreatedPayload) => void;
  onChatEvent?: (payload: ChatEvent) => void;
  onChatHistory?: (payload: ChatHistoryPayload) => void;
};

export function createSocket(baseUrl: string, callbacks: SocketCallbacks): Socket {
  const socket = io(baseUrl, {
    transports: ["websocket"],
    autoConnect: true,
  });

  socket.on("connect", () => callbacks.onConnect?.());
  socket.on("disconnect", () => callbacks.onDisconnect?.());
  socket.on("error", (payload: { message?: string }) => {
    callbacks.onError?.(payload.message ?? "Socket error");
  });
  socket.on("sessions:list:result", (payload) => callbacks.onSessionsList?.(payload));
  socket.on("runtime:created", (payload) => callbacks.onRuntimeCreated?.(payload));
  socket.on("runtime:opened", (payload) => callbacks.onRuntimeOpened?.(payload));
  socket.on("chat:event", (payload) => callbacks.onChatEvent?.(payload));
  socket.on("chat:history", (payload) => callbacks.onChatHistory?.(payload));

  return socket;
}

export function listModels(
  socket: Socket,
  payload: { runtimeId?: string; cwd?: string },
  timeoutMs = 10000,
): Promise<ModelListPayload> {
  return new Promise((resolve, reject) => {
    socket.timeout(timeoutMs).emit(
      "models:list",
      payload,
      (error: Error | null, response: ModelListPayload | { error: string }) => {
        if (error) {
          reject(new Error("Timed out loading models"));
        } else if ("error" in response) {
          reject(new Error(response.error));
        } else {
          resolve(response);
        }
      },
    );
  });
}

export function setRuntimeModel(
  socket: Socket,
  runtimeId: string,
  model: ModelRef,
  timeoutMs = 10000,
): Promise<LiveRuntimeInfo> {
  return new Promise((resolve, reject) => {
    socket.timeout(timeoutMs).emit(
      "runtime:model:set",
      { runtimeId, model },
      (
        error: Error | null,
        response: { runtime: LiveRuntimeInfo } | { error: string },
      ) => {
        if (error) {
          reject(new Error("Timed out setting model"));
        } else if ("error" in response) {
          reject(new Error(response.error));
        } else {
          resolve(response.runtime);
        }
      },
    );
  });
}