#!/usr/bin/env node
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import axios, { AxiosError } from "axios";
import * as z from "zod";
import fs from "node:fs";

// --- Config from env (same as backend) ---
const RASPBERRY_PI_HOST = process.env.RASPBERRY_PI_HOST ?? "http://raspberrypi.local";
const RASPBERRY_PI_PORT = Number(process.env.RASPBERRY_PI_PORT) || 8000;
const QUIBOT_TOKEN = process.env.QUIBOT_TOKEN ?? "MY_SECRET_TOKEN";

const RPI_URL = `${RASPBERRY_PI_HOST}:${RASPBERRY_PI_PORT}`;

function rpiUrl(path: string, query?: Record<string, string>): string {
  const url = `${RPI_URL}${path}`;
  if (!query) return url;
  const q = new URLSearchParams({ token: QUIBOT_TOKEN, ...query });
  return `${url}?${q}`;
}

// --- Helpers ---
async function rpiPost(path: string, query?: Record<string, string>, body?: unknown): Promise<unknown> {
  try {
    const res = await axios.post(rpiUrl(path, query), body, { timeout: 10000 });
    return res.data;
  } catch (err) {
    if (err instanceof AxiosError && err.response) {
      throw new Error(`Pi error ${err.response.status}: ${JSON.stringify(err.response.data)}`);
    }
    throw err;
  }
}

async function rpiGet(path: string, query?: Record<string, string>): Promise<unknown> {
  try {
    const res = await axios.get(rpiUrl(path, query), { timeout: 10000 });
    return res.data;
  } catch (err) {
    if (err instanceof AxiosError && err.response) {
      throw new Error(`Pi error ${err.response.status}: ${JSON.stringify(err.response.data)}`);
    }
    throw err;
  }
}

async function rpiPostMultipart(path: string, formData: FormData): Promise<unknown> {
  try {
    const res = await axios.post(rpiUrl(path), formData, {
      timeout: 30000,
      headers: { "Content-Type": "multipart/form-data" },
    });
    return res.data;
  } catch (err) {
    if (err instanceof AxiosError && err.response) {
      throw new Error(`Pi error ${err.response.status}: ${JSON.stringify(err.response.data)}`);
    }
    throw err;
  }
}

// --- MCP Server ---
const server = new McpServer({
  name: "quibot",
  version: "1.0.0",
});

// === TOOLS ===

server.registerTool(
  "motor_step",
  {
    description: "Move stepper motors in a direction (fire-and-forget — motor runs until stop)",
    inputSchema: z.object({
      direction: z.enum(["forward", "backward", "left", "right"]),
    }),
  },
  async ({ direction }) => {
    const piPath = direction === "backward" ? "/motor/step/backwards" : `/motor/step/${direction}`;
    try {
      const result = await rpiPost(piPath);
      return {
        content: [{ type: "text", text: JSON.stringify(result) }],
      };
    } catch (err) {
      const msg = err instanceof Error ? err.message : String(err);
      return {
        content: [{ type: "text", text: `Error: ${msg}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "motor_stop",
  {
    description: "Stop all stepper motors immediately (disables driver via GPIO EN)",
    inputSchema: z.object({}),
  },
  async () => {
    try {
      const result = await rpiPost("/motor/stop");
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "audio_upload_transcribe",
  {
    description:
      "Upload an audio file to the Raspberry Pi and get transcription + LLM response. The Pi runs Whisper for transcription and llamacpp with preamble for response.",
    inputSchema: z.object({
      audioBase64: z.string().describe("Base64-encoded audio file"),
      format: z.string().describe("Audio format (wav, m4a, mp3, etc.)"),
    }),
  },
  async ({ audioBase64, format }) => {
    try {
      const buffer = Buffer.from(audioBase64, "base64");
      const formData = new FormData();
      const fname = `audio-upload-${Date.now()}.${format}`;
      formData.append("file", new Blob([buffer]), fname);

      const result = await rpiPostMultipart("/transcribe", formData);
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "motor_upload",
  {
    description: "Upload an audio file to the Raspberry Pi for processing",
    inputSchema: z.object({
      filePath: z.string().describe("Path to audio file on local machine"),
      format: z.string().describe("Audio format (wav, m4a, mp3, etc.)"),
    }),
  },
  async ({ filePath, format }) => {
    try {
      if (!fs.existsSync(filePath)) {
        return { content: [{ type: "text", text: `File not found: ${filePath}` }], isError: true };
      }
      const fd = new (await import("form-data")).default();
      fd.append("file", fs.createReadStream(filePath));
      fd.append("format", format);
      await axios.post(`${RPI_URL}/audio/upload?format=${format}&token=${QUIBOT_TOKEN}`, fd, {
        headers: fd.getHeaders(),
        timeout: 30000,
      });
      return { content: [{ type: "text", text: `Uploaded: ${filePath}` }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "audio_list",
  {
    description: "List incoming audio files on the Raspberry Pi",
    inputSchema: z.object({}),
  },
  async () => {
    try {
      const files = await rpiGet("/audio/incoming");
      return { content: [{ type: "text", text: JSON.stringify(files, null, 2) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "audio_lifecycle",
  {
    description: "Manage audio file lifecycle: lock, unlock, cancel, or process a file",
    inputSchema: z.object({
      filename: z.string().describe("Audio filename"),
      action: z.enum(["lock", "unlock", "cancel", "process"]).describe("Lifecycle action to perform"),
    }),
  },
  async ({ filename, action }) => {
    try {
      const result = await rpiPost(`/audio/${action}/${filename}`);
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "eye_set_shape",
  {
    description: "Set the LED eye shape on the robot's WS2811 matrix",
    inputSchema: z.object({
      shape: z.enum(["EYES_OPEN", "EYES_FW", "EYES_DOWN", "EYES_GESTURE"]).describe("Eye shape pattern"),
    }),
  },
  async ({ shape }) => {
    try {
      const result = await rpiPost(`/eye/shape/${shape}`);
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "eye_set_color",
  {
    description: "Set the LED eye color on the robot's WS2811 matrix",
    inputSchema: z.object({
      color: z.enum(["RED", "GREEN", "BLUE", "YELLOW", "CYAN", "MAGENTA", "WHITE", "OFF", "ORANGE", "VIOLET", "DARK_RED"]).describe("Eye color name"),
    }),
  },
  async ({ color }) => {
    try {
      const result = await rpiPost(`/eye/color/${color}`);
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "eye_toggle",
  {
    description: "Turn eyes on or off (enable/disable LED matrix breathing thread)",
    inputSchema: z.object({
      state: z.enum(["on", "off"]).describe("Eye power state"),
    }),
  },
  async ({ state }) => {
    try {
      const result = await rpiPost(`/eye/${state}`);
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "gesture_toggle_mode",
  {
    description: "Toggle between block mode and gesture mode on the robot",
    inputSchema: z.object({}),
  },
  async () => {
    try {
      const result = await rpiPost("/gesture/toggle");
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "gesture_on",
  {
    description: "Enable gesture sensor polling (PAJ7620U2)",
    inputSchema: z.object({}),
  },
  async () => {
    try {
      const result = await rpiPost("/gesture/on");
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "gesture_off",
  {
    description: "Disable gesture sensor polling",
    inputSchema: z.object({}),
  },
  async () => {
    try {
      const result = await rpiPost("/gesture/off");
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "run_command",
  {
    description: "Run a whitelisted system command on the Raspberry Pi",
    inputSchema: z.object({
      task: z.enum(["restart_nginx", "uptime", "update"]).describe("Command to run"),
    }),
  },
  async ({ task }) => {
    try {
      const result = await rpiPost("/run", { task });
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

server.registerTool(
  "tts_speak",
  {
    description: "Synthesize speech via Piper TTS and play it on the Raspberry Pi",
    inputSchema: z.object({
      text: z.string().describe("Text to speak"),
      lang: z.string().default("ca").describe("Language code (ca, es, en)"),
    }),
  },
  async ({ text, lang }) => {
    try {
      const result = await rpiGet("/tts", { text, lang: lang || "ca" });
      return { content: [{ type: "text", text: JSON.stringify(result) }] };
    } catch (err) {
      return {
        content: [{ type: "text", text: `Error: ${err instanceof Error ? err.message : String(err)}` }],
        isError: true,
      };
    }
  },
);

// === RESOURCES ===

server.registerResource(
  "config",
  "quibot://config",
  { description: "Current Raspberry Pi connection config and token" },
  async () => ({
    contents: [
      {
        uri: "quibot://config",
        name: "QuiBot Configuration",
        mimeType: "application/json",
        text: JSON.stringify(
          { raspberryPiHost: RASPBERRY_PI_HOST, raspberryPiPort: RASPBERRY_PI_PORT, token: QUIBOT_TOKEN },
          null,
          2,
        ),
      },
    ],
  }),
);

server.registerResource(
  "available-directions",
  "quibot://directions",
  { description: "Available motor movement directions" },
  async () => ({
    contents: [
      {
        uri: "quibot://directions",
        name: "Available Motor Directions",
        mimeType: "text/plain",
        text: ["forward", "backward"].map((d) => `  POST /motor/step/${d}`).join("\n"),
      },
    ],
  }),
);

server.registerResource(
  "eye-shapes",
  "quibot://eyes/shapes",
  { description: "Available LED eye shapes and their meanings" },
  async () => ({
    contents: [
      {
        uri: "quibot://eyes/shapes",
        name: "Available Eye Shapes",
        mimeType: "text/plain",
        text: [
          "  EYES_OPEN  — Normal resting eyes (default)",
          "  EYES_FW    — Forward-looking eyes",
          "  EYES_DOWN  — Downward/downcast eyes",
          "  EYES_GESTURE — Gesture-acknowledge eyes",
        ].join("\n"),
      },
    ],
  }),
);

server.registerResource(
  "color-actions",
  "quibot://blocks/colors",
  { description: "Color-to-action mapping for block recognition" },
  async () => ({
    contents: [
      {
        uri: "quibot://blocks/colors",
        name: "Color Block Actions",
        mimeType: "text/plain",
        text: [
          "  RED    → Advance forward",
          "  GREEN  → Turn right",
          "  BLUE   → Turn left",
          "  YELLOW → Take / pick up block",
          "  ORANGE → Leave / eject block",
          "  VIOLET → Idle",
          "  BLACK  → Reference / no block",
        ].join("\n"),
      },
    ],
  }),
);

server.registerResource(
  "gestures",
  "quibot://gestures",
  { description: "Available PAJ7620U2 gestures" },
  async () => ({
    contents: [
      {
        uri: "quibot://gestures",
        name: "Available Gestures",
        mimeType: "text/plain",
        text: [
          "  GS_FORWARD      → Hand moving forward (toward sensor)",
          "  GS_BACKWARD     → Hand moving backward (away from sensor)",
          "  GS_LEFT         → Hand moving left",
          "  GS_RIGHT        → Hand moving right",
          "  GS_UP           → Hand moving up",
          "  GS_DOWN         → Hand moving down",
          "  GS_CLOCKWISE    → Clockwise wave motion",
          "  GS_ANTICLOCKWISE→ Counter-clockwise wave motion",
          "  GS_WAVE         → Wave hello gesture",
        ].join("\n"),
      },
    ],
  }),
);

server.registerResource(
  "pi-status",
  "quibot://status/pi",
  { description: "Check if Raspberry Pi HTTP server is reachable" },
  async () => {
    try {
      const res = await axios.get(`${RPI_URL}/health`, { timeout: 5000 });
      return {
        contents: [
          {
            uri: "quibot://status/pi",
            name: "Pi Status",
            mimeType: "application/json",
            text: JSON.stringify({ status: "connected", data: res.data }, null, 2),
          },
        ],
      };
    } catch {
      return {
        contents: [
          {
            uri: "quibot://status/pi",
            name: "Pi Status",
            mimeType: "application/json",
            text: JSON.stringify({ status: "disconnected", error: "Cannot reach Raspberry Pi" }, null, 2),
          },
        ],
      };
    }
  },
);

// === PROMPTS ===

server.registerPrompt(
  "quibot-setup",
  {
    description: "Get a complete reference for controlling the QuiBot robot via MCP",
  },
  () => ({
    messages: [
      {
        role: "user",
        content: {
          type: "text",
          text: `# QuiBot MCP Server

You can control the physical QuiBot robot with these tools:

## Motor Control
- motor_step(direction) — Move forward, backward, left, or right
- motor_stop() — Stop all motors
- motor_upload(filePath, format) — Upload audio file to Pi

## Audio
- audio_upload_transcribe(audioBase64, format) — Upload + Whisper transcribe + llamacpp response
- audio_list() — List incoming audio files
- audio_lifecycle(filename, action) — lock/unlock/cancel/process audio
- tts_speak(text, lang) — Synthesize and play speech via Piper TTS

## Eyes (WS2811 LED Matrix)
- eye_set_shape(shape) — Set face expression shape
- eye_set_color(color) — Change eye color
- eye_toggle(state) — Turn eyes on/off

## Gesture Sensor
- gesture_on() / gesture_off() — Enable/disable gesture polling
- gesture_toggle_mode() — Toggle between block/gesture mode

## System
- run_command(task) — Run system commands on Pi (uptime, restart_nginx)

## Resources
- quibot://status/pi — Check if Pi is reachable
- quibot://config — Current connection config
- quibot://blocks/colors — Color-to-action mapping
- quibot://gestures — Gesture reference
`,
        },
      },
    ],
  }),
);

// === START ===

async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  console.error("[quibot-mcp] Server connected, waiting for requests via stdio...");
}

main().catch((err) => {
  console.error("[quibot-mcp] Failed to start:", err);
  process.exit(1);
});
