import { randomUUID } from "node:crypto";

import type {
  CuaDriverLike,
  MacOsPermissionStatus,
  ToolResult,
} from "@trycua/cua-driver";

import type { ComputerAction } from "./schema.ts";
import { validateComputerAction } from "./schema.ts";

export type CuaModule = typeof import("@trycua/cua-driver");
export type CuaModuleLoader = () => Promise<CuaModule>;

export interface ComputerToolDetails {
  action: ComputerAction["type"];
  session: string;
  initializationWarnings: string[];
  nativeResults: Array<{
    text: string;
    isError: boolean;
    errorCode?: string;
    degraded: boolean;
    structuredJson?: string;
    rawJson: string;
  }>;
  observation: {
    text: string;
    degraded: boolean;
    structuredJson?: string;
    rawJson: string;
  };
  outcomeUnknown: boolean;
}

export interface ComputerToolResult {
  content: Array<
    | { type: "text"; text: string }
    | { type: "image"; data: string; mimeType: string }
  >;
  details: ComputerToolDetails;
}

export interface PermissionSetupResult {
  supported: boolean;
  before?: MacOsPermissionStatus;
  after?: MacOsPermissionStatus;
  openedScreenRecordingSettings: boolean;
}

export interface CuaComputerRuntimeOptions {
  loadCua?: CuaModuleLoader;
  operationTimeoutMs?: number;
  platform?: NodeJS.Platform;
  sessionId?: string;
  waitMs?: number;
}

const defaultCuaLoader: CuaModuleLoader = async () => {
  try {
    return await import("@trycua/cua-driver");
  } catch (error) {
    throw new Error(
      `Could not load the bundled Cua Driver native SDK for ${process.platform}/${process.arch}. ` +
        "Reinstall with npm optional dependencies enabled and confirm the host is macOS, Windows, " +
        `or glibc Linux on arm64/x64. Original error: ${String(error)}`,
      { cause: error },
    );
  }
};

function abortableDelay(ms: number, signal?: AbortSignal): Promise<void> {
  if (signal?.aborted) return Promise.reject(signal.reason ?? new Error("Operation aborted"));

  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => {
      signal?.removeEventListener("abort", onAbort);
      resolve();
    }, ms);
    const onAbort = () => {
      clearTimeout(timer);
      reject(signal?.reason ?? new Error("Operation aborted"));
    };
    signal?.addEventListener("abort", onAbort, { once: true });
  });
}

function operationSignal(signal: AbortSignal | undefined, timeoutMs: number): AbortSignal {
  const timeout = AbortSignal.timeout(timeoutMs);
  return signal ? AbortSignal.any([signal, timeout]) : timeout;
}

function normalizeKey(key: string, platform: NodeJS.Platform): string {
  const normalized = key.trim().toUpperCase();
  const aliases: Record<string, string> = {
    CONTROL: "ctrl",
    CTRL: "ctrl",
    SHIFT: "shift",
    ALT: platform === "darwin" ? "option" : "alt",
    OPTION: platform === "darwin" ? "option" : "alt",
    META: platform === "darwin" ? "cmd" : platform === "win32" ? "win" : "super",
    CMD: platform === "darwin" ? "cmd" : platform === "win32" ? "win" : "super",
    COMMAND: platform === "darwin" ? "cmd" : platform === "win32" ? "win" : "super",
    WIN: "win",
    WINDOWS: "win",
    SUPER: "super",
    RETURN: "enter",
    ENTER: "enter",
    ESC: "escape",
    ESCAPE: "escape",
    SPACE: "space",
    ARROWUP: "up",
    ARROWDOWN: "down",
    ARROWLEFT: "left",
    ARROWRIGHT: "right",
  };
  return aliases[normalized] ?? normalized.toLowerCase();
}

function summarizeResult(result: ToolResult) {
  return {
    text: result.text,
    isError: result.isError,
    ...(result.errorCode ? { errorCode: result.errorCode } : {}),
    degraded: result.degraded,
    ...(result.structuredJson ? { structuredJson: result.structuredJson } : {}),
    rawJson: result.rawJson,
  };
}

export class CuaComputerRuntime {
  readonly sessionId: string;

  private readonly loadCua: CuaModuleLoader;
  private readonly operationTimeoutMs: number;
  private readonly platform: NodeJS.Platform;
  private readonly waitMs: number;
  private cua?: CuaModule;
  private driver?: CuaDriverLike;
  private started = false;
  private closing = false;
  private closePromise?: Promise<void>;
  private queue: Promise<unknown> = Promise.resolve();
  private readonly initializationWarnings: string[] = [];

  constructor(options: CuaComputerRuntimeOptions = {}) {
    this.loadCua = options.loadCua ?? defaultCuaLoader;
    this.operationTimeoutMs = options.operationTimeoutMs ?? 30_000;
    this.platform = options.platform ?? process.platform;
    this.sessionId = options.sessionId ?? `pi-computer-${randomUUID().slice(0, 12)}`;
    this.waitMs = options.waitMs ?? 2_000;
  }

  execute(action: ComputerAction, signal?: AbortSignal): Promise<ComputerToolResult> {
    if (this.closing) return Promise.reject(new Error("Computer runtime is shutting down"));
    const run = this.queue.then(() => this.executeExclusive(action, signal));
    this.queue = run.catch(() => undefined);
    return run;
  }

  async setupPermissions(): Promise<PermissionSetupResult> {
    if (this.platform !== "darwin") {
      return { supported: false, openedScreenRecordingSettings: false };
    }

    const cua = await this.getCua();
    const before = cua.currentMacOsPermissionStatus();
    const after = cua.requestMacOsPermissions();
    let openedScreenRecordingSettings = false;
    if (!after.screenRecording) {
      cua.openMacOsScreenRecordingSettings();
      openedScreenRecordingSettings = true;
    }
    return { supported: true, before, after, openedScreenRecordingSettings };
  }

  close(): Promise<void> {
    if (this.closePromise) return this.closePromise;
    this.closing = true;
    this.closePromise = this.queue.then(() => this.closeExclusive());
    return this.closePromise;
  }

  private async executeExclusive(
    rawAction: ComputerAction,
    signal?: AbortSignal,
  ): Promise<ComputerToolResult> {
    signal?.throwIfAborted();
    const action = validateComputerAction(rawAction);
    await this.ensureStarted(signal);

    if (action.type === "screenshot") {
      return this.observeResult(action.type, [], false, undefined, signal);
    }
    if (action.type === "wait") {
      await abortableDelay(this.waitMs, signal);
      return this.observeResult(action.type, [], false, undefined, signal);
    }

    let nativeResults: ToolResult[] = [];
    let outcomeUnknown = false;
    let warning: string | undefined;
    try {
      nativeResults = await this.performAction(action, signal);
      const error = nativeResults.find((result) => result.isError);
      if (error) {
        outcomeUnknown = true;
        warning =
          `Computer action reported an error and its outcome may be unknown (${error.text || error.errorCode || "unknown error"}). ` +
          "A fresh screenshot follows. Do not retry until it proves the action did not land.";
      }
    } catch (error) {
      outcomeUnknown = true;
      warning =
        `Computer action outcome is unknown (${String(error)}). ` +
        "A fresh screenshot follows. Do not retry until it proves the action did not land.";
    }

    // Use a fresh timeout-only signal for recovery observation. If the caller's
    // signal fired after native dispatch, the screenshot is still needed to
    // determine whether the action landed.
    return this.observeResult(action.type, nativeResults, outcomeUnknown, warning, undefined);
  }

  private async ensureStarted(signal?: AbortSignal): Promise<void> {
    if (this.started) return;
    const cua = await this.getCua();

    if (this.platform === "darwin") {
      const status = cua.currentMacOsPermissionStatus();
      if (!status.accessibility || !status.screenRecording) {
        throw new Error(
          "Computer Use requires macOS Accessibility and Screen Recording permission. " +
            "Run /computer-permissions, grant both permissions, then restart Pi.",
        );
      }
    }

    const driver = cua.CuaDriver.create(undefined);
    this.driver = driver;
    let sessionStarted = false;
    try {
      const boundedSignal = operationSignal(signal, this.operationTimeoutMs);
      await driver.startSession(
        cua.StartSessionInput.new({
          session: this.sessionId,
          captureScope: cua.CaptureScope.Desktop,
        }),
        { signal: boundedSignal },
      );
      sessionStarted = true;
      const cursor = await driver.setAgentCursorEnabled(
        cua.SetAgentCursorEnabledInput.new({
          session: this.sessionId,
          enabled: true,
        }),
        { signal: boundedSignal },
      );
      if (cursor.isError) {
        this.initializationWarnings.push(
          `Cua agent cursor could not be enabled: ${cursor.text || cursor.errorCode || "unknown error"}`,
        );
      }
      this.started = true;
    } catch (error) {
      await this.disposeDriver(driver, sessionStarted);
      this.driver = undefined;
      throw error;
    }
  }

  private async getCua(): Promise<CuaModule> {
    if (!this.cua) this.cua = await this.loadCua();
    return this.cua;
  }

  private async performAction(action: ComputerAction, signal?: AbortSignal): Promise<ToolResult[]> {
    const cua = this.cua;
    const driver = this.driver;
    if (!cua || !driver) throw new Error("Computer runtime is not initialized");
    const boundedSignal = operationSignal(signal, this.operationTimeoutMs);
    const asyncOptions = { signal: boundedSignal };
    const scope = cua.DesktopScope.Desktop;
    const session = this.sessionId;

    switch (action.type) {
      case "click": {
        const button =
          action.button === "right"
            ? cua.ClickButton.Right
            : action.button === "wheel"
              ? cua.ClickButton.Middle
              : cua.ClickButton.Left;
        return [
          await driver.click(
            cua.ClickInput.new({ x: action.x!, y: action.y!, scope, session, button, count: 1 }),
            asyncOptions,
          ),
        ];
      }
      case "double_click":
        return [
          await driver.click(
            cua.ClickInput.new({
              x: action.x!,
              y: action.y!,
              scope,
              session,
              button: cua.ClickButton.Left,
              count: 2,
            }),
            asyncOptions,
          ),
        ];
      case "move":
        return [
          await driver.moveCursor(
            cua.MoveCursorInput.new({ x: action.x!, y: action.y!, scope, session }),
            asyncOptions,
          ),
        ];
      case "type":
        return [
          await driver.typeText(
            cua.TypeTextInput.new({ text: action.text!, scope, session }),
            asyncOptions,
          ),
        ];
      case "keypress": {
        const keys = action.keys!.map((key) => normalizeKey(key, this.platform));
        if (keys.length === 1) {
          return [
            await driver.pressKey(
              cua.PressKeyInput.new({ key: keys[0]!, scope, session }),
              asyncOptions,
            ),
          ];
        }
        return [
          await driver.hotkey(cua.HotkeyInput.new({ keys, scope, session }), asyncOptions),
        ];
      }
      case "drag": {
        const path = action.path!;
        const from = path[0]!;
        const to = path[path.length - 1]!;
        return [
          await driver.drag(
            cua.DragInput.new({
              fromX: from.x,
              fromY: from.y,
              toX: to.x,
              toY: to.y,
              scope,
              session,
              durationMs: 500n,
              steps: BigInt(Math.min(path.length, 200)),
              button: cua.ClickButton.Left,
            }),
            asyncOptions,
          ),
        ];
      }
      case "scroll": {
        const results: ToolResult[] = [];
        const axes = [
          {
            delta: action.scroll_x!,
            positive: cua.ScrollDirection.Right,
            negative: cua.ScrollDirection.Left,
          },
          {
            delta: action.scroll_y!,
            positive: cua.ScrollDirection.Down,
            negative: cua.ScrollDirection.Up,
          },
        ];
        for (const axis of axes) {
          if (axis.delta === 0) continue;
          const amount = BigInt(Math.min(50, Math.max(1, Math.ceil(Math.abs(axis.delta) / 120))));
          const result = await driver.scroll(
            cua.ScrollInput.new({
              x: action.x!,
              y: action.y!,
              direction: axis.delta > 0 ? axis.positive : axis.negative,
              scope,
              session,
              by: cua.ScrollBy.Line,
              amount,
            }),
            asyncOptions,
          );
          results.push(result);
          if (result.isError) break;
        }
        return results;
      }
      case "screenshot":
      case "wait":
        return [];
    }
  }

  private async observeResult(
    action: ComputerAction["type"],
    nativeResults: ToolResult[],
    outcomeUnknown: boolean,
    warning: string | undefined,
    signal?: AbortSignal,
  ): Promise<ComputerToolResult> {
    const cua = this.cua;
    const driver = this.driver;
    if (!cua || !driver) throw new Error("Computer runtime is not initialized");
    const observation = await driver.getDesktopState(
      cua.GetDesktopStateInput.new({ session: this.sessionId }),
      { signal: operationSignal(signal, this.operationTimeoutMs) },
    );
    if (observation.isError) {
      throw new Error(
        `${warning ? `${warning} ` : ""}Desktop capture failed: ${observation.text || observation.errorCode || "unknown error"}`,
      );
    }
    const image = observation.images.find(
      (candidate) => candidate.mimeType.toLowerCase() === "image/png",
    );
    if (!image) {
      throw new Error(`${warning ? `${warning} ` : ""}Desktop capture returned no PNG image`);
    }

    const text = [
      ...this.initializationWarnings,
      warning,
      ...nativeResults.map((result) => result.text).filter(Boolean),
      observation.text,
    ]
      .filter((part): part is string => Boolean(part?.trim()))
      .join("\n");

    return {
      content: [
        { type: "text", text: text || "Desktop screenshot captured." },
        { type: "image", data: image.dataBase64, mimeType: image.mimeType },
      ],
      details: {
        action,
        session: this.sessionId,
        initializationWarnings: [...this.initializationWarnings],
        nativeResults: nativeResults.map(summarizeResult),
        observation: {
          text: observation.text,
          degraded: observation.degraded,
          ...(observation.structuredJson
            ? { structuredJson: observation.structuredJson }
            : {}),
          rawJson: observation.rawJson,
        },
        outcomeUnknown,
      },
    };
  }

  private async closeExclusive(): Promise<void> {
    const driver = this.driver;
    this.driver = undefined;
    if (!driver) return;
    await this.disposeDriver(driver, this.started);
    this.started = false;
  }

  private async disposeDriver(driver: CuaDriverLike, endSession: boolean): Promise<void> {
    const cua = this.cua;
    try {
      if (endSession && cua) {
        await driver.endSession(
          cua.EndSessionInput.new({ session: this.sessionId }),
          { signal: AbortSignal.timeout(this.operationTimeoutMs) },
        );
      }
    } finally {
      try {
        await driver.shutdown({ signal: AbortSignal.timeout(this.operationTimeoutMs) });
      } finally {
        (driver as CuaDriverLike & { uniffiDestroy(): void }).uniffiDestroy();
      }
    }
  }
}

export { normalizeKey };