#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  CallToolRequest,
  ListToolsRequest,
  Tool,
  TextContent,
  CallToolResult,
  ListToolsResult,
} from '@modelcontextprotocol/sdk/types.js';

interface FileSubmission {
  name: string;
  content: string;
  mimeType?: string;
  size?: number;
}

interface ProcessFileSubmissionArgs {
  files: FileSubmission[];
  userPrompt?: string;
  promptType?: PromptType;
}

interface ConfigurePromptsArgs {
  defaultPrompt?: string;
  filePrompt?: string;
  imagePrompt?: string;
  codePrompt?: string;
  dataPrompt?: string;
}

type PromptType = 'default' | 'file' | 'image' | 'code' | 'data';

interface PromptConfig {
  defaultPrompt: string;
  filePrompt: string;
  imagePrompt: string;
  codePrompt: string;
  dataPrompt: string;
}

class AutoPromptServer {
  private server: Server;
  private config: PromptConfig;

  private readonly IMAGE_EXTENSIONS = new Set([
    'jpg', 'jpeg', 'png', 'gif', 'bmp', 'svg', 'webp'
  ]);

  private readonly CODE_EXTENSIONS = new Set([
    'js', 'ts', 'jsx', 'tsx', 'py', 'java', 'cpp', 'c', 'cs', 'php', 'rb', 'go',
    'rs', 'swift', 'kt', 'scala', 'html', 'css', 'scss', 'sass', 'vue', 'svelte'
  ]);

  private readonly DATA_EXTENSIONS = new Set([
    'csv', 'json', 'xml', 'xlsx', 'xls', 'sql'
  ]);

  constructor() {
    this.server = new Server({
      name: 'auto-prompt-server',
      version: '1.0.0',
      capabilities: {
        tools: {},
      },
    });

    // Default configuration - can be overridden via environment variables
    this.config = {
      defaultPrompt: process.env.AUTO_PROMPT_DEFAULT || 'Please analyze this file and provide insights.',
      filePrompt: process.env.AUTO_PROMPT_FILE || 'What can you tell me about this file?',
      imagePrompt: process.env.AUTO_PROMPT_IMAGE || 'Describe this image and its contents.',
      codePrompt: process.env.AUTO_PROMPT_CODE || 'Review this code and explain what it does.',
      dataPrompt: process.env.AUTO_PROMPT_DATA || 'Analyze this data and summarize key findings.',
    };

    this.setupHandlers();
  }

  private validateProcessFileSubmissionArgs(args: unknown): ProcessFileSubmissionArgs {
    if (typeof args !== 'object' || args === null) {
      throw new Error('Invalid arguments: expected object');
    }
    
    const obj = args as Record<string, unknown>;
    
    if (!Array.isArray(obj.files)) {
      throw new Error('Invalid arguments: files must be an array');
    }

    // Validate each file in the array
    for (const file of obj.files) {
      if (typeof file !== 'object' || file === null) {
        throw new Error('Invalid file: each file must be an object');
      }
      const fileObj = file as Record<string, unknown>;
      if (typeof fileObj.name !== 'string' || typeof fileObj.content !== 'string') {
        throw new Error('Invalid file: name and content must be strings');
      }
    }
    
    return {
      files: obj.files as FileSubmission[],
      userPrompt: typeof obj.userPrompt === 'string' ? obj.userPrompt : undefined,
      promptType: typeof obj.promptType === 'string' ? obj.promptType as PromptType : undefined,
    };
  }

  private validateConfigurePromptsArgs(args: unknown): ConfigurePromptsArgs {
    if (typeof args !== 'object' || args === null) {
      throw new Error('Invalid arguments: expected object');
    }
    
    return args as ConfigurePromptsArgs;
  }

  private setupHandlers(): void {
    this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
      tools: [
        {
          name: 'process_file_submission',
          description: 'Processes file submissions and adds appropriate prompts when none are provided',
          inputSchema: {
            type: 'object',
            properties: {
              files: {
                type: 'array',
                items: {
                  type: 'object',
                  properties: {
                    name: { type: 'string' },
                    content: { type: 'string' },
                    mimeType: { type: 'string' },
                    size: { type: 'number' }
                  },
                  required: ['name', 'content']
                }
              },
              userPrompt: {
                type: 'string',
                description: 'Optional user prompt - if provided, no auto-prompt will be added'
              },
              promptType: {
                type: 'string',
                enum: ['default', 'file', 'image', 'code', 'data'],
                description: 'Type of auto-prompt to use if no user prompt provided'
              }
            },
            required: ['files']
          }
        },
        {
          name: 'configure_prompts',
          description: 'Configure the default prompts used for different file types',
          inputSchema: {
            type: 'object',
            properties: {
              defaultPrompt: { type: 'string' },
              filePrompt: { type: 'string' },
              imagePrompt: { type: 'string' },
              codePrompt: { type: 'string' },
              dataPrompt: { type: 'string' }
            }
          }
        },
        {
          name: 'get_configuration',
          description: 'Get current prompt configuration',
          inputSchema: {
            type: 'object',
            properties: {}
          }
        }
      ],
    }));

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

      try {
        switch (name) {
          case 'process_file_submission':
            return await this.processFileSubmission(this.validateProcessFileSubmissionArgs(args));
          case 'configure_prompts':
            return await this.configurePrompts(this.validateConfigurePromptsArgs(args));
          case 'get_configuration':
            return await this.getConfiguration();
          default:
            throw new Error(`Unknown tool: ${name}`);
        }
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : 'Unknown error occurred';
        return {
          content: [
            {
              type: 'text',
              text: `Error: ${errorMessage}`,
            },
          ],
          isError: true,
        };
      }
    });
  }

  private async processFileSubmission(args: ProcessFileSubmissionArgs) {
    const { files, userPrompt, promptType = 'default' } = args;

    if (!files || !Array.isArray(files) || files.length === 0) {
      throw new Error('No files provided');
    }

    // If user provided a prompt, don't add auto-prompt
    if (userPrompt && userPrompt.trim()) {
      return {
        content: [
          {
            type: 'text',
            text: `User prompt detected: "${userPrompt}"\nNo auto-prompt added. Processing ${files.length} file(s) with user's prompt.`,
          },
        ],
      };
    }

    // No user prompt detected, add appropriate auto-prompt
    const autoPrompt = this.selectAutoPrompt(files, promptType);
    const filesList = files.map(file => `- ${file.name} (${file.mimeType || 'unknown type'})`).join('\n');

    return {
      content: [
        {
          type: 'text',
          text: `No user prompt detected. Auto-prompt added: "${autoPrompt}"\n\nFiles submitted:\n${filesList}\n\nProcessing with auto-prompt...`,
        },
      ],
    };
  }

  private selectAutoPrompt(files: FileSubmission[], promptType: PromptType): string {
    // If specific prompt type requested, use it
    if (promptType !== 'default') {
      const promptKey = `${promptType}Prompt` as keyof PromptConfig;
      return this.config[promptKey] || this.config.defaultPrompt;
    }

    // Auto-detect based on file types
    const fileExtensions = files.map(file => {
      const ext = file.name.toLowerCase().split('.').pop();
      return ext || '';
    });

    const mimeTypes = files.map(file => file.mimeType || '').filter(Boolean);

    // Check for images
    if (this.hasImageFiles(mimeTypes, fileExtensions)) {
      return this.config.imagePrompt;
    }

    // Check for code files
    if (this.hasCodeFiles(fileExtensions)) {
      return this.config.codePrompt;
    }

    // Check for data files
    if (this.hasDataFiles(fileExtensions, mimeTypes)) {
      return this.config.dataPrompt;
    }

    // Default for other file types
    return this.config.filePrompt;
  }

  private hasImageFiles(mimeTypes: string[], extensions: string[]): boolean {
    return mimeTypes.some(type => type.startsWith('image/')) ||
           extensions.some(ext => this.IMAGE_EXTENSIONS.has(ext));
  }

  private hasCodeFiles(extensions: string[]): boolean {
    return extensions.some(ext => this.CODE_EXTENSIONS.has(ext));
  }

  private hasDataFiles(extensions: string[], mimeTypes: string[]): boolean {
    return extensions.some(ext => this.DATA_EXTENSIONS.has(ext)) ||
           mimeTypes.some(type => 
             type.includes('spreadsheet') || 
             type.includes('json') || 
             type.includes('xml')
           );
  }

  private async configurePrompts(args: ConfigurePromptsArgs) {
    const validKeys: (keyof PromptConfig)[] = ['defaultPrompt', 'filePrompt', 'imagePrompt', 'codePrompt', 'dataPrompt'];
    const updates: string[] = [];

    for (const [key, value] of Object.entries(args)) {
      if (validKeys.includes(key as keyof PromptConfig) && typeof value === 'string' && value.trim()) {
        const configKey = key as keyof PromptConfig;
        this.config[configKey] = value.trim();
        updates.push(`${key}: "${value.trim()}"`);
      }
    }

    if (updates.length === 0) {
      throw new Error('No valid prompt configurations provided');
    }

    return {
      content: [
        {
          type: 'text',
          text: `Prompt configuration updated:\n${updates.join('\n')}`,
        },
      ],
    };
  }

  private async getConfiguration() {
    const configText = Object.entries(this.config)
      .map(([key, value]) => `${key}: "${value}"`)
      .join('\n');

    return {
      content: [
        {
          type: 'text',
          text: `Current prompt configuration:\n${configText}`,
        },
      ],
    };
  }

  public async run(): Promise<void> {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error('Auto-Prompt MCP server running on stdio');
  }
}

// Run the server
const isMainModule = process.argv[1] && process.argv[1].endsWith('server.js');
if (isMainModule) {
  const server = new AutoPromptServer();
  server.run().catch(console.error);
}

export { AutoPromptServer };
export default AutoPromptServer;