import { McpServer, type CallToolResult } from '@modelcontextprotocol/server';
import * as z from 'zod/v4';
import type { ComputerBackend } from '../backend/computer-backend.js';
import { isComputerError } from '../errors.js';
import type { ComputerAction, ObserveRequest } from '../types.js';

const rectSchema = z.object({
  x: z.number().int(), y: z.number().int(), width: z.number().int().positive(), height: z.number().int().positive(),
});
const displaySchema = rectSchema.extend({ id: z.string(), name: z.string(), primary: z.boolean(), pixelWidth: z.number().int().positive(), pixelHeight: z.number().int().positive(), scale: z.number().positive() });
const pointSchema = z.discriminatedUnion('space', [
  z.object({ space: z.literal('desktop'), x: z.number().int(), y: z.number().int() }),
  z.object({ space: z.literal('display'), displayId: z.string().min(1), x: z.number().int(), y: z.number().int() }),
]);
const regionSchema = z.discriminatedUnion('space', [
  z.object({ space: z.literal('desktop'), x: z.number().int(), y: z.number().int(), width: z.number().int().positive(), height: z.number().int().positive() }),
  z.object({ space: z.literal('display'), displayId: z.string().min(1), x: z.number().int(), y: z.number().int(), width: z.number().int().positive(), height: z.number().int().positive() }),
]);
const targetSchema = z.discriminatedUnion('kind', [
  z.object({ kind: z.literal('desktop') }),
  z.object({ kind: z.literal('display'), displayId: z.string().min(1) }),
  z.object({ kind: z.literal('region'), region: regionSchema }),
]);
const buttonSchema = z.enum(['left', 'middle', 'right']);
const actionSchema = z.discriminatedUnion('type', [
  z.object({ type: z.literal('move'), point: pointSchema }),
  z.object({ type: z.literal('click'), button: buttonSchema.default('left'), point: pointSchema.optional() }),
  z.object({ type: z.literal('doubleClick'), button: buttonSchema.default('left'), point: pointSchema.optional(), intervalMs: z.number().int().min(0).max(2000).default(120) }),
  z.object({ type: z.literal('scroll'), deltaX: z.number().int().default(0), deltaY: z.number().int().default(0), point: pointSchema.optional() }),
  z.object({ type: z.literal('drag'), button: buttonSchema.default('left'), from: pointSchema, to: pointSchema }),
  z.object({ type: z.literal('mouseDown'), button: buttonSchema.default('left'), point: pointSchema.optional() }),
  z.object({ type: z.literal('mouseUp'), button: buttonSchema.default('left'), point: pointSchema.optional() }),
  z.object({ type: z.literal('key'), keys: z.array(z.string().min(1).max(128)).min(1).max(32) }),
  z.object({ type: z.literal('text'), text: z.string(), delayMs: z.number().int().min(0).max(10_000).default(0) }),
]);

function failure(error: unknown, operation: string): CallToolResult {
  const body = isComputerError(error)
    ? { code: error.code, message: error.message, operation: error.operation, ...(error.details === undefined ? {} : { details: error.details }) }
    : { code: 'OS_ERROR', message: 'Unexpected computer backend failure.', operation };
  return { content: [{ type: 'text', text: JSON.stringify({ error: body }) }], isError: true };
}

export function registerTools(server: McpServer, backend: ComputerBackend): void {
  server.registerTool(
    'computer.info',
    {
      title: 'Computer Info',
      description: 'Inspect the visual Linux desktop backend, displays, coordinate spaces, virtual bounds, and current observe/input capabilities. Use this when display topology or coordinates matter.',
      inputSchema: z.object({}),
      outputSchema: z.object({
        backend: z.enum(['x11', 'wayland']),
        sessionType: z.string(),
        coordinateUnit: z.enum(['physical-pixel', 'logical-pixel']),
        coordinateSpaces: z.array(z.enum(['desktop', 'display'])),
        virtualBounds: rectSchema,
        displays: z.array(displaySchema),
        capabilities: z.object({ observe: z.boolean(), pointer: z.boolean(), keyboard: z.boolean(), text: z.boolean() }),
      }),
      annotations: { readOnlyHint: true, idempotentHint: true, openWorldHint: false },
    },
    async () => {
      try {
        const info = await backend.info();
        return { content: [{ type: 'text', text: JSON.stringify(info) }], structuredContent: { ...info, displays: [...info.displays], coordinateSpaces: [...info.coordinateSpaces] } };
      } catch (error) {
        return failure(error, 'computer.info');
      }
    },
  );

  server.registerTool(
    'computer.observe',
    {
      title: 'Observe Computer',
      description: 'Capture exactly one PNG screenshot of the whole desktop, one display, or one region. This tool does not OCR, interpret, retry, or take actions.',
      inputSchema: z.object({
        target: targetSchema.default({ kind: 'desktop' }),
        includePointer: z.boolean().default(true),
      }),
      outputSchema: z.object({
        backend: z.enum(['x11', 'wayland']),
        mimeType: z.literal('image/png'),
        bytes: z.number().int().nonnegative(),
        capturedAt: z.string(),
        region: rectSchema,
        pixelWidth: z.number().int().positive(),
        pixelHeight: z.number().int().positive(),
      }),
      annotations: { readOnlyHint: true, idempotentHint: false, openWorldHint: false },
    },
    async ({ target, includePointer }): Promise<CallToolResult> => {
      try {
        const request: ObserveRequest = { target, includePointer };
        const observation = await backend.observe(request);
        const metadata = {
          backend: observation.backend,
          mimeType: observation.mimeType,
          bytes: observation.bytes,
          capturedAt: observation.capturedAt,
          region: observation.region,
          pixelWidth: observation.pixelWidth,
          pixelHeight: observation.pixelHeight,
        };
        return {
          content: [
            { type: 'text', text: JSON.stringify(metadata) },
            { type: 'image', data: observation.data.toString('base64'), mimeType: observation.mimeType },
          ],
          structuredContent: metadata,
        };
      } catch (error) {
        return failure(error, 'computer.observe');
      }
    },
  );

  server.registerTool(
    'computer.act',
    {
      title: 'Act on Computer',
      description: 'Execute exactly one explicit desktop input action. The MCP does not observe afterward, retry, recover, plan, or run a hidden interaction loop.',
      inputSchema: z.object({ action: actionSchema }),
      outputSchema: z.object({ type: z.enum(['move', 'click', 'doubleClick', 'scroll', 'drag', 'mouseDown', 'mouseUp', 'key', 'text']), executed: z.literal(true), executedAt: z.string() }),
      annotations: { readOnlyHint: false, destructiveHint: true, idempotentHint: false, openWorldHint: false },
    },
    async ({ action }) => {
      try {
        const result = await backend.act(action as ComputerAction);
        return { content: [{ type: 'text', text: JSON.stringify(result) }], structuredContent: result };
      } catch (error) {
        return failure(error, 'computer.act');
      }
    },
  );
}