import { describe, expect, it, vi } from "vitest";

import {
  clearRecordedE2EUIResponses,
  getRecordedE2EUIResponses,
  recordSyntheticE2EUIRequest,
} from "../src/e2e-ui-harness-state.js";
import type { ClientMessage, ServerMessage, Session } from "../src/types.js";
import { WsMessageHandler, type WsMessageHandlerDeps } from "../src/ws-message-handler.js";

interface HandlerHarness {
  handler: WsMessageHandler;
  session: Session;
  sent: ServerMessage[];
  sessions: {
    sendPrompt: ReturnType<typeof vi.fn>;
    sendSteer: ReturnType<typeof vi.fn>;
    sendFollowUp: ReturnType<typeof vi.fn>;
    getMessageQueue: ReturnType<typeof vi.fn>;
    setMessageQueue: ReturnType<typeof vi.fn>;
    sendAbort: ReturnType<typeof vi.fn>;
    stopSession: ReturnType<typeof vi.fn>;
    getActiveSession: ReturnType<typeof vi.fn>;
    respondToUIRequest: ReturnType<typeof vi.fn>;
    forwardClientCommand: ReturnType<typeof vi.fn>;
  };
  ensureSessionContextWindow: ReturnType<typeof vi.fn>;
}

function makeSession(id = "s1"): Session {
  const now = Date.now();
  return {
    id,
    workspaceId: "w1",
    status: "ready",
    createdAt: now,
    lastActivity: now,
    messageCount: 0,
    tokens: { input: 0, output: 0 },
    cost: 0,
  };
}

function makeHarness(): HandlerHarness {
  const session = makeSession();
  const sent: ServerMessage[] = [];

  const sessions = {
    sendPrompt: vi.fn(async () => {}),
    sendSteer: vi.fn(async () => {}),
    sendFollowUp: vi.fn(async () => {}),
    getMessageQueue: vi.fn(() => ({ version: 0, steering: [], followUp: [] })),
    setMessageQueue: vi.fn(async () => ({ version: 0, steering: [], followUp: [] })),
    sendAbort: vi.fn(async () => {}),
    stopSession: vi.fn(async () => {}),
    getActiveSession: vi.fn(() => undefined as Session | undefined),
    respondToUIRequest: vi.fn(() => true),
    forwardClientCommand: vi.fn(async () => {}),
  };

  const ensureSessionContextWindow = vi.fn((value: Session) => value);

  const deps: WsMessageHandlerDeps = {
    sessions,
    ensureSessionContextWindow,
  };

  const handler = new WsMessageHandler(deps);

  return {
    handler,
    session,
    sent,
    sessions,
    ensureSessionContextWindow,
  };
}

function dispatch(harness: HandlerHarness, msg: ClientMessage): Promise<void> {
  return harness.handler.handleClientMessage(harness.session, msg, (outbound) => {
    harness.sent.push(outbound);
  });
}

describe("WsMessageHandler", () => {
  it("rejects legacy raw image prompt transport", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "prompt",
      message: "hello",
      images: [{ data: "base64data", mimeType: "image/png" }],
      clientTurnId: "turn-1",
      requestId: "req-1",
      streamingBehavior: "steer",
    } as unknown as ClientMessage);

    expect(harness.sessions.sendPrompt).not.toHaveBeenCalled();

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "prompt",
        requestId: "req-1",
        success: false,
        error:
          "Raw base64 image transport is not supported; upload images as chat attachments first",
      },
    ]);
  });

  it("returns command_result failure when prompt handler throws and requestId exists", async () => {
    const harness = makeHarness();
    harness.sessions.sendPrompt.mockRejectedValueOnce(new Error("prompt failed"));

    await dispatch(harness, {
      type: "prompt",
      message: "hello",
      requestId: "req-2",
    });

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "prompt",
        requestId: "req-2",
        success: false,
        error: "prompt failed",
      },
    ]);
  });

  it("rethrows prompt handler errors when requestId is absent", async () => {
    const harness = makeHarness();
    harness.sessions.sendPrompt.mockRejectedValueOnce(new Error("prompt failed"));

    await expect(
      dispatch(harness, {
        type: "prompt",
        message: "hello",
      }),
    ).rejects.toThrow("prompt failed");

    expect(harness.sent).toEqual([]);
  });

  it("hydrates get_state through ensureSessionContextWindow", async () => {
    const harness = makeHarness();
    const activeSession: Session = {
      ...harness.session,
      model: "openai-codex/gpt-5.3-codex",
      contextWindow: 200000,
    };
    const normalizedSession: Session = {
      ...activeSession,
      contextWindow: 272000,
    };

    harness.sessions.getActiveSession.mockReturnValue(activeSession);
    harness.ensureSessionContextWindow.mockReturnValue(normalizedSession);

    await dispatch(harness, { type: "get_state", requestId: "req-3" });

    expect(harness.sessions.getActiveSession).toHaveBeenCalledWith("s1");
    expect(harness.ensureSessionContextWindow).toHaveBeenCalledWith(activeSession);
    expect(harness.sent).toEqual([
      {
        type: "state",
        session: normalizedSession,
      },
    ]);
  });

  it("returns queue_state and command_result for get_queue", async () => {
    const harness = makeHarness();
    const queue = {
      version: 3,
      steering: [{ id: "q1", message: "steer", createdAt: 1 }],
      followUp: [{ id: "q2", message: "follow", createdAt: 2 }],
    };
    harness.sessions.getMessageQueue.mockReturnValue(queue);

    await dispatch(harness, {
      type: "get_queue",
      requestId: "req-queue-1",
    });

    expect(harness.sessions.getMessageQueue).toHaveBeenCalledWith("s1");
    expect(harness.sent).toEqual([
      { type: "queue_state", queue },
      {
        type: "command_result",
        command: "get_queue",
        requestId: "req-queue-1",
        success: true,
        data: queue,
      },
    ]);
  });

  it("forwards set_queue and emits command_result", async () => {
    const harness = makeHarness();
    const queue = {
      version: 4,
      steering: [{ id: "q1", message: "steer", createdAt: 1 }],
      followUp: [],
    };
    harness.sessions.setMessageQueue.mockResolvedValue(queue);

    await dispatch(harness, {
      type: "set_queue",
      baseVersion: 3,
      steering: [{ id: "q1", message: "steer" }],
      followUp: [],
      requestId: "req-queue-2",
    });

    expect(harness.sessions.setMessageQueue).toHaveBeenCalledWith("s1", {
      baseVersion: 3,
      steering: [{ id: "q1", message: "steer" }],
      followUp: [],
    });

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "set_queue",
        requestId: "req-queue-2",
        success: true,
        data: queue,
      },
    ]);
  });

  it("reports missing extension UI requests", async () => {
    const harness = makeHarness();
    harness.sessions.respondToUIRequest.mockReturnValueOnce(false);

    await dispatch(harness, {
      type: "extension_ui_response",
      id: "ui-1",
      value: "approved",
      confirmed: true,
      requestId: "req-4",
    });

    expect(harness.sessions.respondToUIRequest).toHaveBeenCalledWith("s1", {
      type: "extension_ui_response",
      id: "ui-1",
      value: "approved",
      confirmed: true,
      cancelled: undefined,
    });

    expect(harness.sent).toEqual([
      {
        type: "error",
        error: "UI request not found: ui-1",
      },
    ]);
  });

  it("records harness extension UI responses without reporting synthetic missing requests", async () => {
    const originalFlag = process.env.OPPI_E2E_UI_HARNESS;
    process.env.OPPI_E2E_UI_HARNESS = "1";
    clearRecordedE2EUIResponses("s1");

    try {
      const harness = makeHarness();
      harness.sessions.respondToUIRequest.mockReturnValueOnce(false);
      recordSyntheticE2EUIRequest("s1", "ui-1");

      await dispatch(harness, {
        type: "extension_ui_response",
        id: "ui-1",
        value: "approved",
        requestId: "req-4",
      });

      expect(getRecordedE2EUIResponses("s1")).toMatchObject([
        {
          type: "extension_ui_response",
          sessionId: "s1",
          id: "ui-1",
          value: "approved",
          requestId: "req-4",
        },
      ]);
      expect(harness.sent).toEqual([]);
    } finally {
      clearRecordedE2EUIResponses("s1");
      if (originalFlag === undefined) {
        delete process.env.OPPI_E2E_UI_HARNESS;
      } else {
        process.env.OPPI_E2E_UI_HARNESS = originalFlag;
      }
    }
  });

  it("forwards RPC commands with request IDs", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "set_model",
      provider: "anthropic",
      modelId: "claude-sonnet-4-0",
      requestId: "req-5",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "set_model",
        provider: "anthropic",
        modelId: "claude-sonnet-4-0",
        requestId: "req-5",
      },
      "req-5",
    );
  });

  it("forwards get_commands requests", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "get_commands",
      requestId: "req-commands-1",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "get_commands",
        requestId: "req-commands-1",
      },
      "req-commands-1",
    );
  });

  it("forwards share_session requests", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "share_session",
      requestId: "req-share-1",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "share_session",
        requestId: "req-share-1",
      },
      "req-share-1",
    );
  });

  it("forwards get_fork_messages requests", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "get_fork_messages",
      requestId: "req-forks-1",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "get_fork_messages",
        requestId: "req-forks-1",
      },
      "req-forks-1",
    );
  });

  it("forwards get_session_tree requests", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "get_session_tree",
      filterMode: "no-tools",
      requestId: "req-tree-1",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "get_session_tree",
        filterMode: "no-tools",
        requestId: "req-tree-1",
      },
      "req-tree-1",
    );
  });

  it("returns command_result when passthrough forwarding throws", async () => {
    const harness = makeHarness();
    harness.sessions.forwardClientCommand.mockRejectedValueOnce(
      new Error("Session not active: s1"),
    );

    await dispatch(harness, {
      type: "get_session_tree",
      filterMode: "default",
      requestId: "req-tree-err",
    });

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "get_session_tree",
        requestId: "req-tree-err",
        success: false,
        error: "Session not active: s1",
      },
    ]);
  });

  it("forwards navigate_tree requests", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "navigate_tree",
      targetId: "entry-42",
      summarize: true,
      customInstructions: "Focus on TODOs",
      replaceInstructions: false,
      label: "Branch summary",
      requestId: "req-navigate-1",
    });

    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledTimes(1);
    expect(harness.sessions.forwardClientCommand).toHaveBeenCalledWith(
      "s1",
      {
        type: "navigate_tree",
        targetId: "entry-42",
        summarize: true,
        customInstructions: "Focus on TODOs",
        replaceInstructions: false,
        label: "Branch summary",
        requestId: "req-navigate-1",
      },
      "req-navigate-1",
    );
  });

  it("returns command_result error for unsupported command types", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "future_command_v99",
      requestId: "req-unknown",
    } as unknown as ClientMessage);

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "future_command_v99",
        requestId: "req-unknown",
        success: false,
        error: "Unsupported command type: future_command_v99",
      },
    ]);
  });

  it("handles stop aliases through sendAbort and emits command_result", async () => {
    const harness = makeHarness();

    await dispatch(harness, {
      type: "stop",
      requestId: "req-6",
    });

    expect(harness.sessions.sendAbort).toHaveBeenCalledWith("s1");
    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "stop",
        requestId: "req-6",
        success: true,
      },
    ]);
  });

  it("emits stop_session command_result failure when stopSession throws", async () => {
    const harness = makeHarness();
    harness.sessions.stopSession.mockRejectedValueOnce(new Error("stop failed"));

    await dispatch(harness, {
      type: "stop_session",
      requestId: "req-7",
    });

    expect(harness.sent).toEqual([
      {
        type: "command_result",
        command: "stop_session",
        requestId: "req-7",
        success: false,
        error: "stop failed",
      },
    ]);
  });
});