import type { SemanticState } from "./protocol.js";
import { emitHerdrBlocked } from "./herdrInterop.js";

export interface SemanticStateClient {
  updateSemanticState(state: SemanticState): boolean | void;
  shutdownSession?(reason?: SessionShutdownReason): boolean | void;
}

export interface SemanticStatePiApi {
  on(event: string, handler: (event: unknown, ctx: unknown) => unknown): void;
  events?: {
    emit?: (eventName: string, data: unknown) => void;
  };
}

export interface SemanticStateControllerOptions {
  idleDebounceMs?: number;
  askUserToolNames?: string[];
}

export type SessionShutdownReason = "quit" | "reload" | "new" | "resume" | "fork";

export interface SemanticStateController {
  readonly currentState: SemanticState;
  markWorking(): void;
  scheduleIdle(): void;
  shutdown(options?: { reason?: SessionShutdownReason; releaseSession?: boolean }): void;
  beginAskPostboxWait(label?: string): () => void;
  handleToolCall(event: ToolEventLike): void;
  handleToolResult(event: ToolEventLike): void;
}

export interface ToolEventLike {
  toolCallId?: string;
  toolName?: string;
}

const DEFAULT_ASK_USER_TOOL_NAMES = ["ask_user"];

export function createSemanticStateController(
  getClient: () => SemanticStateClient | undefined,
  pi?: Pick<SemanticStatePiApi, "events">,
  options: SemanticStateControllerOptions = {}
): SemanticStateController {
  const idleDebounceMs = options.idleDebounceMs ?? 750;
  const askUserToolNames = new Set(options.askUserToolNames ?? DEFAULT_ASK_USER_TOOL_NAMES);
  const localAskUserToolCalls = new Set<string>();
  let askPostboxWaits = 0;
  let agentActive = false;
  let currentState: SemanticState = "idle";
  let idleTimer: NodeJS.Timeout | undefined;

  const clearIdleTimer = () => {
    if (idleTimer) clearTimeout(idleTimer);
    idleTimer = undefined;
  };

  const publish = (nextState: SemanticState) => {
    currentState = nextState;
    getClient()?.updateSemanticState(nextState);
  };

  const recompute = () => {
    clearIdleTimer();
    if (askPostboxWaits > 0) {
      publish("waiting_for_postbox");
      return;
    }
    if (localAskUserToolCalls.size > 0) {
      publish("blocked");
      return;
    }
    publish(agentActive ? "working" : "idle");
  };

  return {
    get currentState() {
      return currentState;
    },

    markWorking() {
      agentActive = true;
      recompute();
    },

    scheduleIdle() {
      agentActive = false;
      clearIdleTimer();
      if (askPostboxWaits > 0) {
        publish("waiting_for_postbox");
        return;
      }
      if (localAskUserToolCalls.size > 0) {
        publish("blocked");
        return;
      }
      idleTimer = setTimeout(() => {
        idleTimer = undefined;
        publish("idle");
      }, idleDebounceMs);
      idleTimer.unref?.();
    },

    shutdown(options = {}) {
      clearIdleTimer();
      agentActive = false;
      askPostboxWaits = 0;
      localAskUserToolCalls.clear();
      emitHerdrBlocked(pi, false);
      getClient()?.updateSemanticState("idle");
      if (options.releaseSession ?? true) {
        getClient()?.shutdownSession?.(options.reason);
      }
    },

    beginAskPostboxWait(label = "Waiting for Postbox answer") {
      clearIdleTimer();
      askPostboxWaits += 1;
      emitHerdrBlocked(pi, true, label);
      publish("waiting_for_postbox");

      let released = false;
      return () => {
        if (released) return;
        released = true;
        askPostboxWaits = Math.max(0, askPostboxWaits - 1);
        emitHerdrBlocked(pi, false);
        recompute();
      };
    },

    handleToolCall(event: ToolEventLike) {
      if (!event.toolName || !askUserToolNames.has(event.toolName)) return;
      clearIdleTimer();
      localAskUserToolCalls.add(event.toolCallId ?? event.toolName);
      publish("blocked");
    },

    handleToolResult(event: ToolEventLike) {
      if (!event.toolName || !askUserToolNames.has(event.toolName)) return;
      localAskUserToolCalls.delete(event.toolCallId ?? event.toolName);
      recompute();
    }
  };
}

export function installSemanticStateHandlers(
  pi: SemanticStatePiApi,
  controller: SemanticStateController
): void {
  pi.on("agent_start", () => controller.markWorking());
  pi.on("agent_end", () => controller.scheduleIdle());
  pi.on("tool_call", (event) => controller.handleToolCall(event as ToolEventLike));
  pi.on("tool_result", (event) => controller.handleToolResult(event as ToolEventLike));
  pi.on("session_shutdown", (event) => {
    const reason = sessionShutdownReasonFromEvent(event);
    controller.shutdown({ reason, releaseSession: reason !== "reload" });
  });
}

function sessionShutdownReasonFromEvent(event: unknown): SessionShutdownReason | undefined {
  if (!event || typeof event !== "object" || !("reason" in event)) return undefined;
  const reason = (event as { reason?: unknown }).reason;
  return reason === "quit" || reason === "reload" || reason === "new" || reason === "resume" || reason === "fork" ? reason : undefined;
}