// src/lib/code-session.ts
// Copyright (C) 2026 Robert Colbert <rob.colbert@openplatform.us>
// All Rights Reserved

import {
  GadgetSocket,
  SocketSession,
  SocketSessionType,
} from "./socket-session.js";
import {
  IChatSession,
  IDroneRegistration,
  IProject,
  IUser,
  ChatTurnStatus,
  GadgetId,
  ChatTurnDocument,
  WorkspaceMode,
  SubmitPromptCallback,
} from "@gadget/api";

import { ChatSessionService, SocketService } from "../services/index.ts";

export class CodeSession extends SocketSession {
  protected type: SocketSessionType = SocketSessionType.Code;

  protected project: IProject | undefined;
  protected chatSession: IChatSession | undefined;
  protected selectedDrone: IDroneRegistration | undefined;
  protected currentTurnId: GadgetId | undefined;
  protected workspaceMode: WorkspaceMode = WorkspaceMode.Idle;

  constructor(socket: GadgetSocket, user: IUser) {
    super(socket, user);
  }

  register() {
    super.register();

    this.socket.on("requestSessionLock", this.onRequestSessionLock.bind(this));
    this.socket.on(
      "requestWorkspaceMode",
      this.onRequestWorkspaceMode.bind(this),
    );
    this.socket.on("submitPrompt", this.onSubmitPrompt.bind(this));
    this.socket.on(
      "releaseSessionLock",
      this.onReleaseSessionLock.bind(this),
    );
    this.socket.on("sessionHeartbeat", this.onSessionHeartbeat.bind(this));
  }

  get hasLock(): boolean {
    return this.selectedDrone !== undefined && this.chatSession !== undefined;
  }

  get selectedDroneId(): GadgetId | undefined {
    return this.selectedDrone?._id;
  }

  get activeChatSession(): IChatSession | undefined {
    return this.chatSession;
  }

  get activeProject(): IProject | undefined {
    return this.project;
  }

  /**
   * Sets the selected drone for this code session.
   */
  setSelectedDrone(registration: IDroneRegistration): void {
    this.selectedDrone = registration;
  }

  /**
   * Sets the active chat session and project for this code session.
   */
  setChatSession(chatSession: IChatSession, project: IProject): void {
    this.chatSession = chatSession;
    this.project = project;
  }

  /**
   * Called when the IDE sends a requestSessionLock event to lock a gadget-drone
   * instance to this code session.
   * @param registration the gadget-drone registration to which the request will
   * be sent.
   * @param project the project we're locking the drone to
   * @param chatSession the chat session we're locking the drone to
   * @param cb response callback to call with the result of the request
   */
  onRequestSessionLock(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    cb: (success: boolean, chatSessionId: string) => void,
  ) {
    const droneSession = SocketService.getDroneSession(registration);
    droneSession.socket.emit(
      "requestSessionLock",
      registration,
      project,
      chatSession,
      (success: boolean, chatSessionId: string): void => {
        if (success) {
          this.selectedDrone = registration;
          this.chatSession = chatSession;
          this.project = project;
          SocketService.registerChatSession(chatSession._id, this);
          droneSession.setChatSessionId(chatSession._id);
        }
        cb(success, chatSessionId);
      },
    );
  }

  /**
   * Called when the IDE sends a requestWorkspaceMode event to change the drone's
   * workspace mode.
   * @param registration the gadget-drone registration
   * @param project the project
   * @param chatSession the chat session
   * @param mode the requested workspace mode
   * @param cb response callback with success and current mode
   */
  onRequestWorkspaceMode(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    mode: WorkspaceMode,
    cb: (success: boolean, currentMode: WorkspaceMode) => void,
  ) {
    if (!this.selectedDrone) {
      this.log.warn("workspace mode request rejected: no drone selected");
      return cb(false, this.workspaceMode);
    }

    const droneSession = SocketService.getDroneSession(this.selectedDrone);
    droneSession.socket.emit(
      "requestWorkspaceMode",
      registration,
      project,
      chatSession,
      mode,
      (success: boolean, currentMode: WorkspaceMode) => {
        if (success) {
          this.workspaceMode = currentMode;
        }
        cb(success, currentMode);
      },
    );
  }

  /**
   * Called when the IDE submits a prompt to be processed by the agent.
   * Creates a ChatTurn document and sends a work order to the selected drone.
   */
  async onSubmitPrompt(
    content: string,
    cb: SubmitPromptCallback,
  ): Promise<void> {
    if (!this.selectedDrone) {
      this.log.warn("prompt rejected: no drone selected");
      return cb(false, { message: "No drone selected" });
    }
    if (!this.chatSession) {
      this.log.warn("prompt rejected: no chat session active");
      return cb(false, { message: "No chat session active" });
    }
    if (!this.project) {
      this.log.warn("prompt rejected: no project selected");
      return cb(false, { message: "No project selected" });
    }

    const droneSession = SocketService.getDroneSession(this.selectedDrone);

    let turn: ChatTurnDocument = await ChatSessionService.createTurn(
      this.chatSession,
      content,
    );
    this.currentTurnId = turn._id;

    this.log.info("ChatTurn created", {
      turnId: turn._id,
      chatSessionId: this.chatSession._id,
    });

    droneSession.setCurrentTurnId(turn._id);
    cb(true, { turnId: turn._id, message: "turn created successfully" });

    droneSession.socket.emit(
      "processWorkOrder",
      this.selectedDrone,
      turn,
      async (success: boolean, message?: string) => {
        if (success) {
          this.log.info("work order accepted by drone", {
            turnId: turn._id,
            message,
          });
        } else {
          this.log.error("work order rejected by drone", {
            turnId: turn._id,
            message,
          });
          turn.status = ChatTurnStatus.Error;
          turn.errorMessage = message || "Drone rejected work order";
          await turn.save();
        }
      },
    );
  }

  /**
   * Called when the IDE sends a releaseSessionLock event to release a
   * previously-acquired session lock on a gadget-drone instance.
   */
  onReleaseSessionLock(
    registration: IDroneRegistration,
    project: IProject,
    chatSession: IChatSession,
    cb: (success: boolean) => void,
  ) {
    const droneSession = SocketService.getDroneSession(registration);
    droneSession.socket.emit(
      "releaseSessionLock",
      registration,
      project,
      chatSession,
      (success: boolean) => {
        if (success) {
          SocketService.unregisterChatSession(chatSession._id);
          droneSession.chatSessionId = undefined;
          this.selectedDrone = undefined;
          this.chatSession = undefined;
          this.project = undefined;
        }
        cb(success);
      },
    );
  }

  /**
   * Called when the IDE sends a sessionHeartbeat event to keep the session
   * lock alive. Forwards to the drone which maintains a timeout.
   */
  onSessionHeartbeat(cb: (ack: boolean) => void) {
    if (!this.selectedDrone) {
      return cb(false);
    }
    const droneSession = SocketService.getDroneSession(this.selectedDrone);
    droneSession.socket.emit("sessionHeartbeat", (ack: boolean) => {
      cb(ack);
    });
  }

  /**
   * Called by DroneSession when the drone emits a workspace mode change.
   * Updates local state and forwards to the IDE socket.
   */
  onWorkspaceModeChanged(mode: WorkspaceMode): void {
    this.workspaceMode = mode;
    this.socket.emit("workspaceModeChanged", mode);
  }

  onStatus(content: string): void {
    this.socket.emit("status", content);
  }

  onThinking(content: string): void {
    this.socket.emit("thinking", content);
  }

  onResponse(content: string): void {
    this.socket.emit("response", content);
  }

  onToolCall(
    callId: string,
    name: string,
    params: string,
    response: string,
  ): void {
    this.socket.emit("toolCall", callId, name, params, response);
  }

  onWorkOrderComplete(
    turnId: string,
    success: boolean,
    message?: string,
  ): void {
    this.socket.emit("workOrderComplete", turnId, success, message);
  }
}
