import { StringEnum, Type } from "@earendil-works/pi-ai";
import type { Static } from "typebox";

export const COMPUTER_ACTION_TYPES = [
  "screenshot",
  "click",
  "double_click",
  "scroll",
  "type",
  "wait",
  "move",
  "keypress",
  "drag",
] as const;

export const COMPUTER_MOUSE_BUTTONS = ["left", "right", "wheel", "back", "forward"] as const;

const coordinate = Type.Integer({ minimum: 0 });

const computerActionSchema = Type.Object(
  {
    type: StringEnum(COMPUTER_ACTION_TYPES, {
      description: "The computer action to perform.",
    }),
    x: Type.Optional(coordinate),
    y: Type.Optional(coordinate),
    button: Type.Optional(
      StringEnum(COMPUTER_MOUSE_BUTTONS, {
        description: "Mouse button for click. wheel is mapped to a middle click.",
      }),
    ),
    scroll_x: Type.Optional(
      Type.Integer({ description: "Horizontal scroll delta. Positive scrolls right." }),
    ),
    scroll_y: Type.Optional(
      Type.Integer({ description: "Vertical scroll delta. Positive scrolls down." }),
    ),
    text: Type.Optional(Type.String({ description: "Text to type verbatim." })),
    keys: Type.Optional(
      Type.Array(Type.String({ minLength: 1 }), {
        minItems: 1,
        maxItems: 16,
        description: "A key or key chord, such as [\"CTRL\", \"L\"].",
      }),
    ),
    path: Type.Optional(
      Type.Array(
        Type.Object(
          {
            x: coordinate,
            y: coordinate,
          },
          { additionalProperties: false },
        ),
        {
          minItems: 2,
          maxItems: 200,
          description: "Drag path in coordinates from the latest screenshot.",
        },
      ),
    ),
  },
  { additionalProperties: false },
);

export const computerParameters = Type.Object(
  {
    action: computerActionSchema,
  },
  { additionalProperties: false },
);

export type ComputerParameters = Static<typeof computerParameters>;
export type ComputerAction = ComputerParameters["action"];
export type ComputerActionType = (typeof COMPUTER_ACTION_TYPES)[number];

export class ComputerActionValidationError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "ComputerActionValidationError";
  }
}

const ACTION_FIELDS: Record<ComputerActionType, ReadonlySet<keyof ComputerAction>> = {
  screenshot: new Set(["type"]),
  click: new Set(["type", "x", "y", "button"]),
  double_click: new Set(["type", "x", "y"]),
  scroll: new Set(["type", "x", "y", "scroll_x", "scroll_y"]),
  type: new Set(["type", "text"]),
  wait: new Set(["type"]),
  move: new Set(["type", "x", "y"]),
  keypress: new Set(["type", "keys"]),
  drag: new Set(["type", "path"]),
};

function requireField<K extends keyof ComputerAction>(
  action: ComputerAction,
  field: K,
): asserts action is ComputerAction & Required<Pick<ComputerAction, K>> {
  if (action[field] === undefined) {
    throw new ComputerActionValidationError(`${action.type} requires action.${String(field)}`);
  }
}

export function validateComputerAction(action: ComputerAction): ComputerAction {
  const allowed = ACTION_FIELDS[action.type];
  if (!allowed) {
    throw new ComputerActionValidationError(
      `Unsupported computer action type: ${String(action.type)}`,
    );
  }
  for (const key of Object.keys(action) as Array<keyof ComputerAction>) {
    if (!allowed.has(key)) {
      throw new ComputerActionValidationError(`${action.type} does not accept action.${String(key)}`);
    }
  }

  switch (action.type) {
    case "click":
      requireField(action, "x");
      requireField(action, "y");
      requireField(action, "button");
      if (action.button === "back" || action.button === "forward") {
        throw new ComputerActionValidationError(
          `Cua Driver does not support the ${action.button} mouse button`,
        );
      }
      break;
    case "double_click":
    case "move":
      requireField(action, "x");
      requireField(action, "y");
      break;
    case "scroll":
      requireField(action, "x");
      requireField(action, "y");
      requireField(action, "scroll_x");
      requireField(action, "scroll_y");
      if (action.scroll_x === 0 && action.scroll_y === 0) {
        throw new ComputerActionValidationError("scroll requires a non-zero scroll_x or scroll_y");
      }
      break;
    case "type":
      requireField(action, "text");
      break;
    case "keypress":
      requireField(action, "keys");
      if (action.keys.length === 0) {
        throw new ComputerActionValidationError("keypress requires at least one key");
      }
      break;
    case "drag":
      requireField(action, "path");
      if (action.path.length < 2) {
        throw new ComputerActionValidationError("drag requires at least two path points");
      }
      break;
    case "screenshot":
    case "wait":
      break;
  }

  return action;
}