src/lib/server/mcp-client.ts

import { DynamicStructuredTool } from "@langchain/core/tools";
import { z } from "zod";
import type { WebSocket } from "ws";
import { consoleBuffer } from "./console-buffer";
import { virtualFileSystem } from "../services/virtual-fs";

interface EditorWebSocketConnection {
  send: (message: {
    type: string;
    payload: Record<string, unknown>;
    timestamp: number;
  }) => void;
}

let wsConnection: EditorWebSocketConnection | null = null;

export function setMCPWebSocketConnection(ws: WebSocket) {
  wsConnection = {
    send: (message: {
      type: string;
      payload: Record<string, unknown>;
      timestamp: number;
    }) => {
      if (ws && ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify(message));
      }
    },
  };
}

/**
 * Standardized tool response builder
 */
function buildToolResponse({
  success,
  action,
  error,
  gameState,
  consoleOutput,
}: {
  success: boolean;
  action: string;
  error?: string;
  gameState?: ReturnType<typeof consoleBuffer.getGameStateFromMessages>;
  consoleOutput?: string[];
}): string {
  const sections: string[] = [];

  // Status section
  if (success) {
    sections.push(`✅ ${action} completed successfully.`);
  } else {
    sections.push(`❌ ${action} failed.`);
  }

  // Error section
  if (error) {
    sections.push(`\nError: ${error}`);
  }

  // Game state section
  if (gameState) {
    if (gameState.isReady) {
      sections.push("\n🎮 Game Status: Running");
    } else if (gameState.hasError) {
      sections.push(`\n🎮 Game Status: Error\n   ${gameState.lastError}`);
    } else if (gameState.isLoading) {
      sections.push("\n🎮 Game Status: Loading...");
    } else {
      sections.push("\n🎮 Game Status: Unknown");
    }
  }

  // Console output section
  if (consoleOutput && consoleOutput.length > 0) {
    sections.push("\n📝 Console Output:");
    sections.push(consoleOutput.join("\n"));
  }

  return sections.join("\n");
}

/**
 * Wait for game state with improved console capture
 */
async function waitForGameState(
  toolName: string,
  maxWaitTime: number = 5000,
): Promise<{
  gameState: ReturnType<typeof consoleBuffer.getGameStateFromMessages>;
  consoleOutput: string[];
}> {
  // Set execution context for console messages
  consoleBuffer.setExecutionContext(toolName);
  consoleBuffer.clearToolMessages();

  const startTime = Date.now();

  // Initial wait for game to process changes (increased to match GameCanvas reload delay)
  await new Promise((resolve) => setTimeout(resolve, 2000));

  while (Date.now() - startTime < maxWaitTime) {
    const gameState = consoleBuffer.getGameStateFromMessages();

    // Check if we have a definitive state (not reloading, and either ready or error)
    if (!gameState.isReloading && (gameState.isReady || gameState.hasError)) {
      const messages = consoleBuffer.getMessagesSinceLastTool();
      const consoleOutput = messages
        .slice(-30)
        .map((m) => `[${m.type}] ${m.message}`);

      // Clear context
      consoleBuffer.setExecutionContext(null);
      consoleBuffer.markAsRead();

      return { gameState, consoleOutput };
    }

    // If still reloading, wait longer
    if (gameState.isReloading) {
      await new Promise((resolve) => setTimeout(resolve, 200));
    } else {
      await new Promise((resolve) => setTimeout(resolve, 100));
    }
  }

  // Timeout - return current state
  const gameState = consoleBuffer.getGameStateFromMessages();
  const messages = consoleBuffer.getMessagesSinceLastTool();
  const consoleOutput = messages
    .slice(-30)
    .map((m) => `[${m.type}] ${m.message}`);

  // Clear context
  consoleBuffer.setExecutionContext(null);
  consoleBuffer.markAsRead();

  return { gameState, consoleOutput };
}

/**
 * MCPClientManager provides MCP-style tools for editor operations
 * Currently uses local implementation, can be extended to use actual MCP servers
 */
export class MCPClientManager {
  private tools: DynamicStructuredTool[] = [];
  private initialized = false;

  async initialize(): Promise<void> {
    if (this.initialized) {
      return;
    }

    this.tools = this.createEditorTools();
    this.initialized = true;
  }

  private createEditorTools(): DynamicStructuredTool[] {
    const tools: DynamicStructuredTool[] = [];

    tools.push(
      new DynamicStructuredTool({
        name: "read_editor",
        description:
          "Read the complete editor content - use for initial exploration or when search returns no results",
        schema: z.object({}),
        func: async () => {
          const file = virtualFileSystem.getGameFile();
          return `Current editor content (html):\n${file.content}`;
        },
      }),
    );

    tools.push(
      new DynamicStructuredTool({
        name: "read_editor_lines",
        description:
          "Read specific lines from the editor - use AFTER search_editor to examine found code sections in detail",
        schema: z.object({
          startLine: z
            .number()
            .min(1)
            .describe("The starting line number (1-indexed)"),
          endLine: z
            .number()
            .min(1)
            .optional()
            .describe(
              "The ending line number (inclusive). If not provided, only the start line is returned",
            ),
        }),
        func: async (input: { startLine: number; endLine?: number }) => {
          const result = virtualFileSystem.getLines(
            input.startLine,
            input.endLine,
          );

          if (result.error) {
            return `Error: ${result.error}`;
          }

          return result.content;
        },
      }),
    );

    tools.push(
      new DynamicStructuredTool({
        name: "search_editor",
        description:
          "Search for code elements and get line numbers - use FIRST to locate specific functions, classes, or components before reading or editing",
        schema: z.object({
          query: z.string().describe("Text or regex pattern to search for"),
          mode: z
            .enum(["text", "regex"])
            .optional()
            .describe(
              "Search mode: 'text' for literal text search, 'regex' for pattern matching (default: text)",
            ),
          contextLines: z
            .number()
            .min(0)
            .max(5)
            .optional()
            .describe(
              "Number of context lines before/after match (default: 2, max: 5)",
            ),
        }),
        func: async (input: {
          query: string;
          mode?: "text" | "regex";
          contextLines?: number;
        }) => {
          const mode = input.mode || "text";
          const results = virtualFileSystem.searchContent(input.query, mode);

          if (results.length === 0) {
            return `No matches found for "${input.query}" in editor content`;
          }

          const totalMatches = results.length;
          const displayMatches = results.slice(0, 10);

          let output = `Found ${totalMatches} match${totalMatches > 1 ? "es" : ""} for "${input.query}":\n\n`;

          displayMatches.forEach((match, index) => {
            if (index > 0) output += "\n---\n\n";
            output += match.context.join("\n");
          });

          if (totalMatches > 10) {
            output += `\n\n(Showing first 10 of ${totalMatches} matches. Use more specific search terms to narrow results)`;
          }

          return output;
        },
      }),
    );

    tools.push(
      new DynamicStructuredTool({
        name: "edit_editor",
        description:
          "Replace specific text in the editor - use for SMALL, targeted changes (max ~20 lines). For large changes, use multiple edit_editor calls with plan_tasks",
        schema: z.object({
          oldText: z
            .string()
            .describe(
              "The exact text to find and replace (keep small - max ~20 lines)",
            ),
          newText: z.string().describe("The text to replace it with"),
        }),
        func: async (input: { oldText: string; newText: string }) => {
          const currentContent = virtualFileSystem.getGameFile().content;
          if (!currentContent.includes(input.oldText)) {
            const shortPreview =
              input.oldText.substring(0, 50) +
              (input.oldText.length > 50 ? "..." : "");

            return buildToolResponse({
              success: false,
              action: "Text replacement",
              error: `Text not found: "${shortPreview}". This might be due to a previous edit already modifying this text. Please verify the current content with read_editor or search_editor.`,
            });
          }

          const result = virtualFileSystem.editContent(
            input.oldText,
            input.newText,
          );

          if (!result.success) {
            return buildToolResponse({
              success: false,
              action: "Text replacement",
              error: result.error,
            });
          }

          const file = virtualFileSystem.getGameFile();
          this.syncEditorContent(file.content);

          // Wait for game to reload and capture console
          const { gameState, consoleOutput } =
            await waitForGameState("edit_editor");

          return buildToolResponse({
            success: true,
            action: "Text replacement",
            gameState,
            consoleOutput,
          });
        },
      }),
    );

    tools.push(
      new DynamicStructuredTool({
        name: "write_editor",
        description:
          "Replace entire editor content - use ONLY for creating new files or complete rewrites. For modifications, use edit_editor with plan_tasks instead",
        schema: z.object({
          content: z
            .string()
            .describe("The complete code content to write to the editor"),
        }),
        func: async (input: { content: string }) => {
          virtualFileSystem.updateGameContent(input.content);
          this.syncEditorContent(input.content);

          // Wait for game to reload and capture console
          const { gameState, consoleOutput } =
            await waitForGameState("write_editor");

          return buildToolResponse({
            success: true,
            action: "Editor content update",
            gameState,
            consoleOutput,
          });
        },
      }),
    );

    tools.push(...this.createContext7Tools());

    return tools;
  }

  private createContext7Tools(): DynamicStructuredTool[] {
    const tools: DynamicStructuredTool[] = [];
    const apiKey = process.env.CONTEXT7_API_KEY;

    if (!apiKey) {
      console.warn(
        "CONTEXT7_API_KEY not set, Context7 tools will not be available",
      );
      return tools;
    }

    tools.push(
      new DynamicStructuredTool({
        name: "resolve_library_id",
        description: "Resolve a library name to Context7-compatible library ID",
        schema: z.object({
          libraryName: z
            .string()
            .describe("The name of the library to resolve"),
        }),
        func: async (input: { libraryName: string }) => {
          try {
            const response = await globalThis.fetch(
              "https://mcp.context7.com/mcp",
              {
                method: "POST",
                headers: {
                  "Content-Type": "application/json",
                  Accept: "application/json, text/event-stream",
                  CONTEXT7_API_KEY: apiKey,
                },
                body: JSON.stringify({
                  jsonrpc: "2.0",
                  id: Math.floor(Math.random() * 10000),
                  method: "tools/call",
                  params: {
                    name: "resolve-library-id",
                    arguments: input,
                  },
                }),
              },
            );

            if (!response.ok) {
              throw new Error(
                `HTTP ${response.status}: ${response.statusText}`,
              );
            }

            const text = await response.text();

            const lines = text.split("\n");
            let jsonData = null;

            for (const line of lines) {
              if (line.startsWith("data: ")) {
                try {
                  jsonData = JSON.parse(line.substring(6));
                  break;
                } catch {
                  // Continue looking for valid JSON
                }
              }
            }

            if (!jsonData) {
              throw new Error("No valid JSON data found in response");
            }

            if (jsonData.error) {
              throw new Error(
                jsonData.error.message || JSON.stringify(jsonData.error),
              );
            }

            return JSON.stringify(jsonData.result, null, 2);
          } catch (error) {
            return `Error resolving library ID for "${input.libraryName}": ${error instanceof Error ? error.message : String(error)}`;
          }
        },
      }),
    );

    tools.push(
      new DynamicStructuredTool({
        name: "get_library_docs",
        description:
          "Fetch up-to-date documentation for a Context7-compatible library",
        schema: z.object({
          context7CompatibleLibraryID: z
            .string()
            .describe("The Context7 library ID (e.g., '/greensock/gsap')"),
          tokens: z
            .number()
            .optional()
            .describe("Maximum tokens to retrieve (default: 5000)"),
          topic: z
            .string()
            .optional()
            .describe("Specific topic to focus on (e.g., 'animations')"),
        }),
        func: async (input: {
          context7CompatibleLibraryID: string;
          tokens?: number;
          topic?: string;
        }) => {
          try {
            const response = await globalThis.fetch(
              "https://mcp.context7.com/mcp",
              {
                method: "POST",
                headers: {
                  "Content-Type": "application/json",
                  Accept: "application/json, text/event-stream",
                  CONTEXT7_API_KEY: apiKey,
                },
                body: JSON.stringify({
                  jsonrpc: "2.0",
                  id: Math.floor(Math.random() * 10000),
                  method: "tools/call",
                  params: {
                    name: "get-library-docs",
                    arguments: {
                      context7CompatibleLibraryID:
                        input.context7CompatibleLibraryID,
                      tokens: input.tokens || 5000,
                      topic: input.topic,
                    },
                  },
                }),
              },
            );

            if (!response.ok) {
              throw new Error(
                `HTTP ${response.status}: ${response.statusText}`,
              );
            }

            const text = await response.text();

            const lines = text.split("\n");
            let jsonData = null;

            for (const line of lines) {
              if (line.startsWith("data: ")) {
                try {
                  jsonData = JSON.parse(line.substring(6));
                  break;
                } catch {
                  // Continue looking for valid JSON
                }
              }
            }

            if (!jsonData) {
              throw new Error("No valid JSON data found in response");
            }

            if (jsonData.error) {
              throw new Error(
                jsonData.error.message || JSON.stringify(jsonData.error),
              );
            }

            return JSON.stringify(jsonData.result, null, 2);
          } catch (error) {
            return `Error fetching docs for "${input.context7CompatibleLibraryID}": ${error instanceof Error ? error.message : String(error)}`;
          }
        },
      }),
    );

    return tools;
  }

  private syncEditorContent(content: string): void {
    if (wsConnection) {
      wsConnection.send({
        type: "editor_update",
        payload: { content },
        timestamp: Date.now(),
      });
    }
  }

  getTools(): DynamicStructuredTool[] {
    return this.tools;
  }

  async cleanup(): Promise<void> {
    this.initialized = false;
  }
}

export const mcpClientManager = new MCPClientManager();