import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { getMcpClient, type McpClientManager } from "./mcpClient";

function jsonSchemaToZodShape(jsonSchema: Record<string, any>): Record<string, any> {
  const properties = jsonSchema.properties ?? {};
  const required = new Set(jsonSchema.required ?? []);
  const shape: Record<string, any> = {};

  for (const [key, prop] of Object.entries(properties)) {
    let schema: any;
    const p = prop as Record<string, any>;

    switch (p.type) {
      case "string":
        schema = z.string();
        if (p.enum) schema = z.enum(p.enum);
        break;
      case "integer":
      case "number":
        schema = z.number();
        if (p.minimum !== undefined) schema = schema.min(p.minimum);
        if (p.maximum !== undefined) schema = schema.max(p.maximum);
        break;
      case "boolean":
        schema = z.boolean();
        break;
      case "array":
        if (p.items?.type === "string") {
          schema = z.array(z.string());
        } else if (p.items?.type === "number" || p.items?.type === "integer") {
          schema = z.array(z.number());
        } else {
          schema = z.array(z.any());
        }
        break;
      case "object":
        schema = z.record(z.any());
        break;
      default:
        schema = z.any();
    }

    if (p.description) schema = schema.describe(p.description);
    if (!required.has(key)) schema = schema.optional();

    shape[key] = schema;
  }
  return shape;
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  console.error(`[toolsProvider] Starting tools provider...`);
  
  try {
    const mcp: McpClientManager = await getMcpClient();
    console.error(`[toolsProvider] Getting MCP client...`);
    
    const remoteTools = await mcp.listTools();
    console.error(`[toolsProvider] Got ${remoteTools.length} remote tools`);
    
    const tools: Tool[] = [];

    for (const remote of remoteTools) {
      const name = remote.name;
      const description = remote.description ?? name;
      const inputSchema = (remote.inputSchema ?? { type: "object", properties: {} }) as Record<string, any>;
      const shape = jsonSchemaToZodShape(inputSchema);

      console.error(`[toolsProvider] Registering tool: ${name}`);

      const t = tool({
        name,
        description,
        parameters: shape,
        implementation: async (params: Record<string, any>, { status, signal }: { status: (s: string) => void; signal: AbortSignal }) => {
          status(`Calling ${name}...`);
          try {
            const result = await mcp.callTool(name, params, signal);
            if (result?.content && Array.isArray(result.content)) {
              return result.content
                .map((c: any) => (c.type === "text" ? c.text : JSON.stringify(c)))
                .join("\n");
            }
            return JSON.stringify(result);
          } catch (err: any) {
            return `Error calling ${name}: ${err.message ?? err}`;
          }
        },
      });

      tools.push(t);
    }

    console.error(`[toolsProvider] Successfully registered ${tools.length} tools`);
    return tools;
  } catch (err: any) {
    console.error(`[toolsProvider] ERROR: ${err.message}`);
    console.error(`[toolsProvider] Stack: ${err.stack}`);
    // 回傳空陣列而不是拋出錯誤，這樣 plugin 至少可以載入
    return [];
  }
}

import { tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { getMcpClient, type McpClientManager } from "./mcpClient";

function jsonSchemaToZodShape(jsonSchema: Record<string, any>): Record<string, any> {
  const properties = jsonSchema.properties ?? {};
  const required = new Set(jsonSchema.required ?? []);
  const shape: Record<string, any> = {};

  for (const [key, prop] of Object.entries(properties)) {
    let schema: any;
    const p = prop as Record<string, any>;

    switch (p.type) {
      case "string":
        schema = z.string();
        if (p.enum) schema = z.enum(p.enum);
        break;
      case "integer":
      case "number":
        schema = z.number();
        if (p.minimum !== undefined) schema = schema.min(p.minimum);
        if (p.maximum !== undefined) schema = schema.max(p.maximum);
        break;
      case "boolean":
        schema = z.boolean();
        break;
      case "array":
        if (p.items?.type === "string") {
          schema = z.array(z.string());
        } else if (p.items?.type === "number" || p.items?.type === "integer") {
          schema = z.array(z.number());
        } else {
          schema = z.array(z.any());
        }
        break;
      case "object":
        schema = z.record(z.any());
        break;
      default:
        schema = z.any();
    }

    if (p.description) schema = schema.describe(p.description);
    if (!required.has(key)) schema = schema.optional();

    shape[key] = schema;
  }
  return shape;
}

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  console.error(`[toolsProvider] Starting tools provider...`);
  
  try {
    const mcp: McpClientManager = await getMcpClient();
    console.error(`[toolsProvider] Getting MCP client...`);
    
    const remoteTools = await mcp.listTools();
    console.error(`[toolsProvider] Got ${remoteTools.length} remote tools`);
    
    const tools: Tool[] = [];

    for (const remote of remoteTools) {
      const name = remote.name;
      const description = remote.description ?? name;
      const inputSchema = (remote.inputSchema ?? { type: "object", properties: {} }) as Record<string, any>;
      const shape = jsonSchemaToZodShape(inputSchema);

      console.error(`[toolsProvider] Registering tool: ${name}`);

      const t = tool({
        name,
        description,
        parameters: shape,
        implementation: async (params: Record<string, any>, { status, signal }: { status: (s: string) => void; signal: AbortSignal }) => {
          status(`Calling ${name}...`);
          try {
            const result = await mcp.callTool(name, params, signal);
            if (result?.content && Array.isArray(result.content)) {
              return result.content
                .map((c: any) => (c.type === "text" ? c.text : JSON.stringify(c)))
                .join("\n");
            }
            return JSON.stringify(result);
          } catch (err: any) {
            return `Error calling ${name}: ${err.message ?? err}`;
          }
        },
      });

      tools.push(t);
    }

    console.error(`[toolsProvider] Successfully registered ${tools.length} tools`);
    return tools;
  } catch (err: any) {
    console.error(`[toolsProvider] ERROR: ${err.message}`);
    console.error(`[toolsProvider] Stack: ${err.stack}`);
    // 回傳空陣列而不是拋出錯誤，這樣 plugin 至少可以載入
    return [];
  }
}