#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  Tool,
} from '@modelcontextprotocol/sdk/types.js';
import { setupDatabase, initializeDatabase, getDatabase } from './database.js';
import {
  remixConversation,
  continueConversation,
  readConversationTree,
  searchParents,
  searchChildren,
} from './tools.js';

// Parse command-line arguments
function parseArgs(): { 
  host?: string;
  port?: number;
  user?: string;
  password?: string;
  database?: string;
  connectionString?: string;
} {
  const args = process.argv.slice(2);
  const result: { 
    host?: string;
    port?: number;
    user?: string;
    password?: string;
    database?: string;
    connectionString?: string;
  } = {};
  
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--host' && i + 1 < args.length) {
      result.host = args[i + 1];
      i++;
    } else if (args[i].startsWith('--host=')) {
      result.host = args[i].substring('--host='.length);
    } else if ((args[i] === '--port' || args[i] === '-p') && i + 1 < args.length) {
      const port = parseInt(args[i + 1], 10);
      if (isNaN(port) || port < 1 || port > 65535) {
        console.error(`Invalid port number: ${args[i + 1]}. Must be between 1 and 65535.`);
        process.exit(1);
      }
      result.port = port;
      i++;
    } else if (args[i].startsWith('--port=')) {
      const port = parseInt(args[i].substring('--port='.length), 10);
      if (isNaN(port) || port < 1 || port > 65535) {
        console.error(`Invalid port number: ${args[i].substring('--port='.length)}. Must be between 1 and 65535.`);
        process.exit(1);
      }
      result.port = port;
    } else if ((args[i] === '--user' || args[i] === '-U') && i + 1 < args.length) {
      result.user = args[i + 1];
      i++;
    } else if (args[i].startsWith('--user=')) {
      result.user = args[i].substring('--user='.length);
    } else if ((args[i] === '--password' || args[i] === '-W') && i + 1 < args.length) {
      console.error(
        'Warning: Passing passwords via command-line arguments can expose them in process listings and shell history.\n' +
        'For production, use environment variables (PGPASSWORD) or connection strings instead.'
      );
      result.password = args[i + 1];
      i++;
    } else if (args[i].startsWith('--password=')) {
      console.error(
        'Warning: Passing passwords via command-line arguments can expose them in process listings and shell history.\n' +
        'For production, use environment variables (PGPASSWORD) or connection strings instead.'
      );
      result.password = args[i].substring('--password='.length);
    } else if ((args[i] === '--database' || args[i] === '-d') && i + 1 < args.length) {
      result.database = args[i + 1];
      i++;
    } else if (args[i].startsWith('--database=')) {
      result.database = args[i].substring('--database='.length);
    } else if (args[i] === '--connection-string' && i + 1 < args.length) {
      result.connectionString = args[i + 1];
      i++;
    } else if (args[i].startsWith('--connection-string=')) {
      result.connectionString = args[i].substring('--connection-string='.length);
    }
  }
  
  return result;
}

// Define MCP tools
const TOOLS: Tool[] = [
  {
    name: 'remix_conversation',
    description: 
      'Allows to save the last X chat messages to create a new conversation or branch from an existing one. ' +
      'This tool accepts a messages[] array where each message contains "author" (either "User" or "Copilot") and "payload" (the message content). ' +
      'The MCP server automatically calculates message IDs, previous/next message IDs, and timestamps for the linked list structure. ' +
      'If old_conversation_id is provided, the new conversation will be linked as a child of that conversation in the tree structure.',
    inputSchema: {
      type: 'object',
      properties: {
        old_conversation_id: {
          type: 'string',
          description: 'Optional ID of the parent conversation to branch from',
        },
        new_conversation_id: {
          type: 'string',
          description: 'ID for the new conversation (must be unique)',
        },
        title: {
          type: 'string',
          description: 'Title for the conversation',
        },
        task_description: {
          type: 'string',
          description: 'Description of the task or context for this conversation',
        },
        messages: {
          type: 'array',
          description: 'Array of messages to save',
          items: {
            type: 'object',
            properties: {
              author: {
                type: 'string',
                enum: ['User', 'Copilot'],
                description: 'Who wrote the message',
              },
              payload: {
                type: 'string',
                description: 'The message content',
              },
            },
            required: ['author', 'payload'],
          },
        },
      },
      required: ['new_conversation_id', 'title', 'task_description', 'messages'],
    },
  },
  {
    name: 'continue_conversation',
    description:
      'Retrieves all messages from a conversation as an array. ' +
      'Each message includes its ID, previous message ID, next message ID, author, payload, and timestamp. ' +
      'Use this to load a saved conversation and continue it in a new AI chat session.',
    inputSchema: {
      type: 'object',
      properties: {
        conversation_id: {
          type: 'string',
          description: 'ID of the conversation to retrieve',
        },
      },
      required: ['conversation_id'],
    },
  },
  {
    name: 'read_conversation_tree',
    description:
      'Reads and returns the conversation tree structure as JSON. ' +
      'Each node contains the conversation ID, title, task description, and children (nested conversations). ' +
      'If root_conversation is specified, returns the tree starting from that conversation. ' +
      'If root_conversation is not specified, returns all root-level conversations (those without parents) and their descendants.',
    inputSchema: {
      type: 'object',
      properties: {
        root_conversation: {
          type: 'string',
          description: 'Optional root conversation ID to start the tree from',
        },
      },
    },
  },
  {
    name: 'search_parents',
    description:
      'Performs full-text search on messages from a child conversation up to all parent conversations. ' +
      'Walks up the conversation tree from the specified child and searches all messages in that path. ' +
      'Returns matching messages with their conversation ID, message ID, prev/next message IDs, author, payload, and timestamp.',
    inputSchema: {
      type: 'object',
      properties: {
        keywords: {
          type: 'array',
          items: {
            type: 'string',
          },
          description: 'Array of keywords to search for',
        },
        child_conversation_id: {
          type: 'string',
          description: 'ID of the child conversation to start searching from',
        },
      },
      required: ['keywords', 'child_conversation_id'],
    },
  },
  {
    name: 'search_children',
    description:
      'Performs full-text search on messages from a parent conversation down to all child conversations recursively. ' +
      'Walks down the conversation tree from the specified parent and searches all messages in the descendants. ' +
      'Returns matching messages with their conversation ID, message ID, prev/next message IDs, author, payload, and timestamp.',
    inputSchema: {
      type: 'object',
      properties: {
        keywords: {
          type: 'array',
          items: {
            type: 'string',
          },
          description: 'Array of keywords to search for',
        },
        parent_conversation_id: {
          type: 'string',
          description: 'ID of the parent conversation to start searching from',
        },
      },
      required: ['keywords', 'parent_conversation_id'],
    },
  },
];

// Create MCP server
const server = new Server(
  {
    name: 'remix-mcp',
    version: '0.3.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

// Handle tool listing
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: TOOLS,
  };
});

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

  try {
    if (!args) {
      throw new Error('Missing arguments');
    }

    switch (name) {
      case 'remix_conversation': {
        const result = await remixConversation(
          args.new_conversation_id as string,
          args.title as string,
          args.task_description as string,
          args.messages as Array<{ author: 'User' | 'Copilot'; payload: string }>,
          args.old_conversation_id as string | undefined
        );
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(result, null, 2),
            },
          ],
        };
      }

      case 'continue_conversation': {
        const messages = await continueConversation(args.conversation_id as string);
        if (messages === null) {
          return {
            content: [
              {
                type: 'text',
                text: JSON.stringify({
                  error: 'Conversation not found',
                  conversation_id: args.conversation_id,
                }, null, 2),
              },
            ],
            isError: true,
          };
        }
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ messages }, null, 2),
            },
          ],
        };
      }

      case 'read_conversation_tree': {
        const tree = await readConversationTree(args.root_conversation as string | undefined);
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ tree }, null, 2),
            },
          ],
        };
      }

      case 'search_parents': {
        const results = await searchParents(
          args.keywords as string[],
          args.child_conversation_id as string
        );
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ results, count: results.length }, null, 2),
            },
          ],
        };
      }

      case 'search_children': {
        const results = await searchChildren(
          args.keywords as string[],
          args.parent_conversation_id as string
        );
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ results, count: results.length }, null, 2),
            },
          ],
        };
      }

      default:
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ error: `Unknown tool: ${name}` }),
            },
          ],
          isError: true,
        };
    }
  } catch (error) {
    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify({
            error: error instanceof Error ? error.message : String(error),
          }),
        },
      ],
      isError: true,
    };
  }
});

// Start server
async function main() {
  // Initialize database with optional custom connection parameters
  try {
    const config = parseArgs();
    await setupDatabase(config);
    await initializeDatabase();
  } catch (error) {
    console.error(
      'Failed to initialize database:',
      error instanceof Error ? error.message : String(error)
    );
    process.exit(1);
  }

  const transport = new StdioServerTransport();
  await server.connect(transport);
  
  // Setup graceful shutdown
  const gracefulShutdown = async () => {
    console.error('Shutting down gracefully...');
    try {
      // Close server transport first to stop accepting new requests
      await server.close();
      console.error('Server transport closed');
      
      // Then close database pool
      const db = getDatabase();
      await db.end();
      console.error('Database pool closed');
    } catch (error) {
      console.error('Error during shutdown:', error);
    }
    process.exit(0);
  };

  process.on('SIGINT', gracefulShutdown);
  process.on('SIGTERM', gracefulShutdown);
  
  console.error('Remix MCP Server running on stdio');
}

main().catch((error) => {
  console.error('Server error:', error);
  process.exit(1);
});