import {
  tool,
  Tool,
  ToolsProviderController,
} from "@lmstudio/sdk";
import { z } from "zod";
import path from "path";
import fs from "fs";

// We DO NOT top-level import from @modelcontextprotocol/sdk to avoid ESM/CJS conflicts.
// Instead, we use dynamic imports inside the async function.

async function getMcpClient() {
    const { Client } = await import("@modelcontextprotocol/sdk/client/index.js");
    const { StdioClientTransport } = await import("@modelcontextprotocol/sdk/client/stdio.js");
    
    // Use the direct path to the index file for maximum compatibility in CJS environments if needed
    // but @modelcontextprotocol/sdk's subpath exports should work with dynamic import().

    const memoryMeshDir = path.resolve(process.cwd(), "MemoryMesh");
    const indexJsPath = path.join(memoryMeshDir, "dist/index.js");

    if (!fs.existsSync(indexJsPath)) {
        throw new Error(`MemoryMesh not found at ${indexJsPath}`);
    }

    const transport = new StdioClientTransport({
        command: "node",
        args: [indexJsPath],
    });

    const client = new Client(
        { name: "lm-studio-wrapper-client", version: "1.0.0" },
        { capabilities: {} }
    );

    await client.connect(transport);
    return { client, transport };
}

/**
 * Helper to convert JSON Schema types to Zod schemas.
 */
function mapJsonSchemaToZod(schema: any): z.ZodTypeAny {
  if (!schema) return z.any();
  if (schema.enum && Array.isArray(schema.enum)) {
    return z.enum(schema.enum as [string, ...string[]]);
  }

  const type = schema.type;
  switch (type) {
    case "string": return z.string();
    case "number": case "integer": return z.number();
    case "boolean": return z.boolean();
    case "array": 
        const itemsType = schema.items ? mapJsonSchemaToZod(schema.items) : z.any();
        return z.array(itemsType);
    case "object":
      if (schema.properties) {
        const shape: Record<string, z.ZodTypeAny> = {};
        for (const key of Object.keys(schema.properties)) {
          let validator = mapJsonSchemaToZod(schema.properties[key]);
          if (!schema.required?.includes(key)) validator = validator.optional();
          shape[key] = validator;
        }
        return z.object(shape);
      }
      return z.record(z.any());
    default: return z.any();
  }
}

function createZodParamsShape(inputSchema: any): z.ZodRawShape {
  if (inputSchema && inputSchema.type === "object" && inputSchema.properties) {
    const shape: Record<string, z.ZodTypeAny> = {};
    for (const key of Object.keys(inputSchema.properties)) {
      let validator = mapJsonSchemaToZod(inputSchema.properties[key]);
      if (!inputSchema.required?.includes(key)) validator = validator.optional();
      shape[key] = validator;
    }
    return shape;
  }
  return {} as z.ZodRawShape;
}

let cachedClient: any = null;
let cachedTransport: any = null;

export async function toolsProvider(
  ctl: ToolsProviderController,
): Promise<Tool[]> {
  try {
    if (!cachedClient) {
        const mcpInstance = await getMcpClient();
        cachedClient = mcpInstance.client;
        cachedTransport = mcpInstance.transport;
    }

    const mcpToolsResponse = await cachedClient.listTools();
    const tools: Tool[] = [];

    for (const mcpTool of mcpToolsResponse.tools) {
      console.log(`[Wrapper] Registering tool: ${mcpTool.name}`);
      const zodParams = createZodParamsShape(mcpTool.inputSchema);

      tools.push(tool({
        name: "mm_" + mcpTool.name,
        description: mcpTool.description || `MemoryMesh tool: ${mcpTool.name}`,
        parameters: zodParams as any, 
        implementation: async (args) => {
          try {
            const result = await cachedClient.callTool({
              name: mcpTool.name,
              arguments: args as any,
            });

            if (Array.isArray(result.content)) {
              return result.content
                .filter((c: any) => c.type === 'text')
                .map((c: any) => c.text)
                .join("\n") || JSON.stringify(result.content);
            }
            return JSON.stringify(result.content);
          } catch (err: any) {
            return `Error during tool execution [${mcpTool.name}]: ${err.message}`;
          }
        },
      }));
    }
    return tools;
  } catch (err: any) {
    console.error("[Wrapper] Error in toolsProvider:", err);
    return [];
  }
}

import {
  tool,
  Tool,
  ToolsProviderController,
} from "@lmstudio/sdk";
import { z } from "zod";
import path from "path";
import fs from "fs";

// We DO NOT top-level import from @modelcontextprotocol/sdk to avoid ESM/CJS conflicts.
// Instead, we use dynamic imports inside the async function.

async function getMcpClient() {
    const { Client } = await import("@modelcontextprotocol/sdk/client/index.js");
    const { StdioClientTransport } = await import("@modelcontextprotocol/sdk/client/stdio.js");
    
    // Use the direct path to the index file for maximum compatibility in CJS environments if needed
    // but @modelcontextprotocol/sdk's subpath exports should work with dynamic import().

    const memoryMeshDir = path.resolve(process.cwd(), "MemoryMesh");
    const indexJsPath = path.join(memoryMeshDir, "dist/index.js");

    if (!fs.existsSync(indexJsPath)) {
        throw new Error(`MemoryMesh not found at ${indexJsPath}`);
    }

    const transport = new StdioClientTransport({
        command: "node",
        args: [indexJsPath],
    });

    const client = new Client(
        { name: "lm-studio-wrapper-client", version: "1.0.0" },
        { capabilities: {} }
    );

    await client.connect(transport);
    return { client, transport };
}

/**
 * Helper to convert JSON Schema types to Zod schemas.
 */
function mapJsonSchemaToZod(schema: any): z.ZodTypeAny {
  if (!schema) return z.any();
  if (schema.enum && Array.isArray(schema.enum)) {
    return z.enum(schema.enum as [string, ...string[]]);
  }

  const type = schema.type;
  switch (type) {
    case "string": return z.string();
    case "number": case "integer": return z.number();
    case "boolean": return z.boolean();
    case "array": 
        const itemsType = schema.items ? mapJsonSchemaToZod(schema.items) : z.any();
        return z.array(itemsType);
    case "object":
      if (schema.properties) {
        const shape: Record<string, z.ZodTypeAny> = {};
        for (const key of Object.keys(schema.properties)) {
          let validator = mapJsonSchemaToZod(schema.properties[key]);
          if (!schema.required?.includes(key)) validator = validator.optional();
          shape[key] = validator;
        }
        return z.object(shape);
      }
      return z.record(z.any());
    default: return z.any();
  }
}

function createZodParamsShape(inputSchema: any): z.ZodRawShape {
  if (inputSchema && inputSchema.type === "object" && inputSchema.properties) {
    const shape: Record<string, z.ZodTypeAny> = {};
    for (const key of Object.keys(inputSchema.properties)) {
      let validator = mapJsonSchemaToZod(inputSchema.properties[key]);
      if (!inputSchema.required?.includes(key)) validator = validator.optional();
      shape[key] = validator;
    }
    return shape;
  }
  return {} as z.ZodRawShape;
}

let cachedClient: any = null;
let cachedTransport: any = null;

export async function toolsProvider(
  ctl: ToolsProviderController,
): Promise<Tool[]> {
  try {
    if (!cachedClient) {
        const mcpInstance = await getMcpClient();
        cachedClient = mcpInstance.client;
        cachedTransport = mcpInstance.transport;
    }

    const mcpToolsResponse = await cachedClient.listTools();
    const tools: Tool[] = [];

    for (const mcpTool of mcpToolsResponse.tools) {
      console.log(`[Wrapper] Registering tool: ${mcpTool.name}`);
      const zodParams = createZodParamsShape(mcpTool.inputSchema);

      tools.push(tool({
        name: "mm_" + mcpTool.name,
        description: mcpTool.description || `MemoryMesh tool: ${mcpTool.name}`,
        parameters: zodParams as any, 
        implementation: async (args) => {
          try {
            const result = await cachedClient.callTool({
              name: mcpTool.name,
              arguments: args as any,
            });

            if (Array.isArray(result.content)) {
              return result.content
                .filter((c: any) => c.type === 'text')
                .map((c: any) => c.text)
                .join("\n") || JSON.stringify(result.content);
            }
            return JSON.stringify(result.content);
          } catch (err: any) {
            return `Error during tool execution [${mcpTool.name}]: ${err.message}`;
          }
        },
      }));
    }
    return tools;
  } catch (err: any) {
    console.error("[Wrapper] Error in toolsProvider:", err);
    return [];
  }
}