#!/usr/bin/env node

/**
 * Dexto MCP Server
 *
 * MCP server for the Dexto agent harness, exposing session management,
 * orchestration, and MCP client/server capabilities.
 */

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import * as fs from "fs";
import * as path from "path";
import { spawn, ChildProcess } from "child_process";

// Tool input schemas
const CreateAgentSchema = z.object({
  name: z.string(),
  config: z.string(), // YAML config
});

const RunAgentSchema = z.object({
  agent: z.string(),
  input: z.string(),
  session: z.string().optional(),
});

const ListSessionsSchema = z.object({});

const ResumeSessionSchema = z.object({
  sessionId: z.string(),
});

const OrchestrateSchema = z.object({
  workflow: z.string(), // YAML workflow definition
  input: z.string(),
});

class DextoMCPServer {
  private server: Server;
  private dextoPath: string;
  private activeSessions: Map<string, ChildProcess> = new Map();

  constructor(dextoPath: string) {
    this.dextoPath = dextoPath;
    this.server = new Server(
      {
        name: "dexto-mcp",
        version: "0.1.0",
      },
      {
        capabilities: {
          tools: {},
        },
      }
    );

    this.setupHandlers();
  }

  private setupHandlers() {
    // List available tools
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      return {
        tools: [
          {
            name: "dexto_create_agent",
            description: "Create a custom agent from YAML configuration",
            inputSchema: {
              type: "object",
              properties: {
                name: { type: "string", description: "Agent name" },
                config: { type: "string", description: "YAML agent configuration" },
              },
              required: ["name", "config"],
            },
          },
          {
            name: "dexto_run_agent",
            description: "Run a configured agent with input",
            inputSchema: {
              type: "object",
              properties: {
                agent: { type: "string", description: "Agent name or ID" },
                input: { type: "string", description: "Input for the agent" },
                session: { type: "string", description: "Optional session ID to resume" },
              },
              required: ["agent", "input"],
            },
          },
          {
            name: "dexto_list_sessions",
            description: "List all active and historical sessions",
            inputSchema: {
              type: "object",
              properties: {},
            },
          },
          {
            name: "dexto_resume_session",
            description: "Resume a previous session",
            inputSchema: {
              type: "object",
              properties: {
                sessionId: { type: "string", description: "Session ID to resume" },
              },
              required: ["sessionId"],
            },
          },
          {
            name: "dexto_orchestrate",
            description: "Orchestrate multi-agent workflow",
            inputSchema: {
              type: "object",
              properties: {
                workflow: { type: "string", description: "YAML workflow definition" },
                input: { type: "string", description: "Workflow input" },
              },
              required: ["workflow", "input"],
            },
          },
          {
            name: "dexto_mcp_connect",
            description: "Connect to an MCP server",
            inputSchema: {
              type: "object",
              properties: {
                serverName: { type: "string", description: "Name for the MCP server" },
                command: { type: "string", description: "Command to start MCP server" },
                args: { type: "array", items: { type: "string" }, description: "Arguments for MCP server" },
              },
              required: ["serverName", "command"],
            },
          },
          {
            name: "dexto_mcp_list_tools",
            description: "List available tools from connected MCP servers",
            inputSchema: {
              type: "object",
              properties: {
                serverName: { type: "string", description: "MCP server name" },
              },
            },
          },
          {
            name: "dexto_memory_store",
            description: "Store information in agent memory",
            inputSchema: {
              type: "object",
              properties: {
                key: { type: "string", description: "Memory key" },
                value: { type: "string", description: "Value to store" },
                session: { type: "string", description: "Optional session ID" },
              },
              required: ["key", "value"],
            },
          },
          {
            name: "dexto_memory_retrieve",
            description: "Retrieve from agent memory",
            inputSchema: {
              type: "object",
              properties: {
                key: { type: "string", description: "Memory key" },
                session: { type: "string", description: "Optional session ID" },
              },
              required: ["key"],
            },
          },
        ],
      };
    });

    // Handle tool calls
    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
      const { name, arguments: args } = request.params;

      try {
        switch (name) {
          case "dexto_create_agent":
            return await this.createAgent(args);
          case "dexto_run_agent":
            return await this.runAgent(args);
          case "dexto_list_sessions":
            return await this.listSessions();
          case "dexto_resume_session":
            return await this.resumeSession(args);
          case "dexto_orchestrate":
            return await this.orchestrate(args);
          case "dexto_mcp_connect":
            return await this.mcpConnect(args);
          case "dexto_mcp_list_tools":
            return await this.mcpListTools(args);
          case "dexto_memory_store":
            return await this.memoryStore(args);
          case "dexto_memory_retrieve":
            return await this.memoryRetrieve(args);
          default:
            throw new Error(`Unknown tool: ${name}`);
        }
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        return {
          content: [
            {
              type: "text",
              text: JSON.stringify({ error: errorMessage }),
            },
          ],
        };
      }
    });
  }

  private async createAgent(args: any) {
    const { name, config } = args;

    // Save agent config to Dexto's agents directory
    const agentsDir = path.join(this.dextoPath, "agents");
    const agentFile = path.join(agentsDir, `${name}.yaml`);

    await fs.promises.mkdir(agentsDir, { recursive: true });
    await fs.promises.writeFile(agentFile, config);

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({
            success: true,
            agent: name,
            path: agentFile,
            message: `Agent '${name}' created successfully`,
          }),
        },
      ],
    };
  }

  private async runAgent(args: any) {
    const { agent, input, session } = args;

    // Run Dexto agent via CLI
    const dextoArgs = ["--agent", agent, input];
    if (session) {
      dextoArgs.push("--session", session);
    }

    const result = await this.runDextoCommand(dextoArgs);

    return {
      content: [
        {
          type: "text",
          text: result,
        },
      ],
    };
  }

  private async listSessions() {
    const sessionsDir = path.join(this.dextoPath, "sessions");

    if (!fs.existsSync(sessionsDir)) {
      return {
        content: [
          {
            type: "text",
            text: JSON.stringify({ sessions: [] }),
          },
        ],
      };
    }

    const sessions = fs.readdirSync(sessionsDir)
      .filter(f => f.endsWith(".json"))
      .map(f => {
        const sessionData = JSON.parse(fs.readFileSync(path.join(sessionsDir, f), "utf-8"));
        return {
          id: f.replace(".json", ""),
          agent: sessionData.agent,
          created: sessionData.created,
          status: sessionData.status,
        };
      });

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({ sessions }),
        },
      ],
    };
  }

  private async resumeSession(args: any) {
    const { sessionId } = args;

    const result = await this.runDextoCommand(["--resume", sessionId]);

    return {
      content: [
        {
          type: "text",
          text: result,
        },
      ],
    };
  }

  private async orchestrate(args: any) {
    const { workflow, input } = args;

    // Save workflow temporarily
    const workflowFile = path.join(this.dextoPath, ".temp-workflow.yaml");
    await fs.promises.writeFile(workflowFile, workflow);

    const result = await this.runDextoCommand(["--workflow", workflowFile, input]);

    // Cleanup
    fs.unlinkSync(workflowFile);

    return {
      content: [
        {
          type: "text",
          text: result,
        },
      ],
    };
  }

  private async mcpConnect(args: any) {
    const { serverName, command, args: cmdArgs } = args;

    // This would integrate with Dexto's MCP client capabilities
    // For now, return a placeholder
    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({
            message: `MCP connection to '${serverName}' queued`,
            command,
            args: cmdArgs,
            note: "Full MCP client integration requires Dexto's MCP module",
          }),
        },
      ],
    };
  }

  private async mcpListTools(args: any) {
    const { serverName } = args;

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({
            server: serverName,
            tools: [],
            note: "Tool listing requires active MCP connection",
          }),
        },
      ],
    };
  }

  private async memoryStore(args: any) {
    const { key, value, session } = args;

    const memoryDir = path.join(this.dextoPath, "memory");
    await fs.promises.mkdir(memoryDir, { recursive: true });

    const memoryFile = session
      ? path.join(memoryDir, `${session}.json`)
      : path.join(memoryDir, "default.json");

    let memory: Record<string, any> = {};
    if (fs.existsSync(memoryFile)) {
      memory = JSON.parse(fs.readFileSync(memoryFile, "utf-8"));
    }

    memory[key] = { value, timestamp: Date.now() };
    await fs.promises.writeFile(memoryFile, JSON.stringify(memory, null, 2));

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({ success: true, key, stored: true }),
        },
      ],
    };
  }

  private async memoryRetrieve(args: any) {
    const { key, session } = args;

    const memoryDir = path.join(this.dextoPath, "memory");
    const memoryFile = session
      ? path.join(memoryDir, `${session}.json`)
      : path.join(memoryDir, "default.json");

    if (!fs.existsSync(memoryFile)) {
      return {
        content: [
          {
            type: "text",
            text: JSON.stringify({ error: "Memory not found" }),
          },
        ],
      };
    }

    const memory = JSON.parse(fs.readFileSync(memoryFile, "utf-8"));
    const value = memory[key];

    return {
      content: [
        {
          type: "text",
          text: JSON.stringify({ key, value }),
        },
      ],
    };
  }

  private async runDextoCommand(args: string[]): Promise<string> {
    return new Promise((resolve, reject) => {
      const dextoProcess = spawn("dexto", args, {
        cwd: this.dextoPath,
        stdio: ["pipe", "pipe", "pipe"],
      });

      let stdout = "";
      let stderr = "";

      dextoProcess.stdout?.on("data", (data) => {
        stdout += data.toString();
      });

      dextoProcess.stderr?.on("data", (data) => {
        stderr += data.toString();
      });

      dextoProcess.on("close", (code) => {
        if (code === 0) {
          resolve(stdout);
        } else {
          reject(new Error(`Dexto exited with code ${code}: ${stderr}`));
        }
      });

      dextoProcess.on("error", (error) => {
        reject(error);
      });
    });
  }

  async start() {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
  }
}

// CLI entry point
async function main() {
  const args = process.argv.slice(2);
  let dextoPath = process.env.DEXTO_PATH || process.cwd();

  for (let i = 0; i < args.length; i++) {
    if (args[i] === "--config" && i + 1 < args.length) {
      const configPath = args[i + 1];
      dextoPath = path.dirname(configPath);
      break;
    }
    if ((args[i] === "--dexto-path" || args[i] === "-d") && i + 1 < args.length) {
      dextoPath = args[i + 1];
      break;
    }
  }

  const server = new DextoMCPServer(dextoPath);
  await server.start();
}

main().catch(console.error);