/**
 * Google Cloud Spanner tools for MCP
 */
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';
import { getProjectId } from '../../utils/auth.js';
import { getSpannerClient, getSpannerConfig } from './types.js';
import { getSpannerSchema } from './schema.js';

/**
 * Get detailed schema information for a Spanner database in a format suitable for query generation
 * 
 * @param instanceId Spanner instance ID
 * @param databaseId Spanner database ID
 * @returns Detailed schema information with table relationships
 */
async function getDetailedSchemaForQueryGeneration(instanceId: string, databaseId: string): Promise<string> {
  const schema = await getSpannerSchema(instanceId, databaseId);
  
  // Format the schema in a way that's useful for SQL generation
  let schemaText = `Database: ${databaseId}\n\nTables:\n`;
  
  for (const table of schema.tables) {
    schemaText += `\nTable: ${table.name}\n`;
    schemaText += `Columns:\n`;
    
    for (const column of table.columns) {
      schemaText += `  - ${column.name}: ${column.type}${column.nullable ? ' (nullable)' : ''}\n`;
    }
    
    if (table.indexes && table.indexes.length > 0) {
      schemaText += `Indexes:\n`;
      for (const index of table.indexes) {
        schemaText += `  - ${index.name}: ${index.columns.join(', ')}${index.unique ? ' (unique)' : ''}\n`;
      }
    }
    
    if (table.foreignKeys && table.foreignKeys.length > 0) {
      schemaText += `Foreign Keys:\n`;
      for (const fk of table.foreignKeys) {
        schemaText += `  - ${fk.name}: ${fk.columns.join(', ')} → ${fk.referencedTable}(${fk.referencedColumns.join(', ')})\n`;
      }
    }
  }
  
  return schemaText;
}

export function registerSpannerTools(server: McpServer): void {
  // Tool to execute SQL queries
  server.tool(
    'execute-spanner-query',
    {
      sql: z.string().describe('The SQL query to execute'),
      instanceId: z.string().optional().describe('Spanner instance ID (defaults to SPANNER_INSTANCE env var)'),
      databaseId: z.string().optional().describe('Spanner database ID (defaults to SPANNER_DATABASE env var)'),
      params: z.record(z.any()).optional().describe('Query parameters')
    },
    async ({ sql, instanceId, databaseId, params }, _extra) => {
      try {
        const projectId = await getProjectId();
        const config = await getSpannerConfig(
          Array.isArray(instanceId) ? instanceId[0] : instanceId,
          Array.isArray(databaseId) ? databaseId[0] : databaseId
        );
        
        const spanner = getSpannerClient();
        const instance = spanner.instance(config.instanceId);
        const database = instance.database(config.databaseId);
        
        // Execute the query
        const [result] = await database.run({
          sql,
          params: params || {}
        });
        
        if (!result || result.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\nQuery executed successfully. No results returned.`
            }]
          };
        }
        
        // Convert to markdown table
        const columns = Object.keys(result[0]);
        
        let markdown = `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n`;
        markdown += `SQL: \`${sql}\`\n\n`;
        markdown += `Rows: ${result.length}\n\n`;
        
        // Table header
        markdown += '| ' + columns.join(' | ') + ' |\n';
        markdown += '| ' + columns.map(() => '---').join(' | ') + ' |\n';
        
        // Table rows (limit to 100 rows for display)
        const displayRows = result.slice(0, 100);
        for (const row of displayRows) {
          const rowValues = columns.map(col => {
            const value = (row as any)[col];
            if (value === null || value === undefined) return 'NULL';
            if (typeof value === 'object') return JSON.stringify(value);
            return String(value);
          });
          
          markdown += '| ' + rowValues.join(' | ') + ' |\n';
        }
        
        if (result.length > 100) {
          markdown += '\n*Results truncated. Showing 100 of ' + result.length + ' rows.*';
        }
        
        return {
          content: [{
            type: 'text',
            text: markdown
          }]
        };
      } catch (error: any) {
        console.error('Error executing Spanner query:', error);
        throw error;
      }
    }
  );

  // Tool to list tables
  server.tool(
    'list-spanner-tables',
    {
      instanceId: z.string().optional().describe('Spanner instance ID (defaults to SPANNER_INSTANCE env var)'),
      databaseId: z.string().optional().describe('Spanner database ID (defaults to SPANNER_DATABASE env var)')
    },
    async ({ instanceId, databaseId }, _extra) => {
      try {
        const projectId = await getProjectId();
        const config = await getSpannerConfig(
          Array.isArray(instanceId) ? instanceId[0] : instanceId,
          Array.isArray(databaseId) ? databaseId[0] : databaseId
        );
        
        const spanner = getSpannerClient();
        const instance = spanner.instance(config.instanceId);
        const database = instance.database(config.databaseId);
        
        // Query for tables
        // Execute query to list tables
        const [tablesResult] = await database.run({
          sql: `SELECT t.table_name, 
                    (SELECT COUNT(1) FROM information_schema.columns 
                     WHERE table_name = t.table_name) as column_count
              FROM information_schema.tables t
              WHERE t.table_catalog = '' AND t.table_schema = ''
              ORDER BY t.table_name`
        });
        
        if (!tablesResult || tablesResult.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Spanner Tables\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\nNo tables found in the database.`
            }]
          };
        }
        
        let markdown = `# Spanner Tables\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n`;
        
        // Table header
        markdown += '| Table Name | Column Count |\n';
        markdown += '|------------|-------------|\n';
        
        // Table rows
        for (const row of tablesResult) {
          // Access the row properties directly
          
          // Extract table name and column count
          const tableName = (row as any).table_name as string || 'unknown';
          const columnCount = (row as any).column_count as number || 0;
          
          markdown += `| ${tableName} | ${columnCount} |\n`;
        }
        
        // Add resource links for further exploration
        markdown += '\n## Available Resources\n\n';
        markdown += `- Schema: \`gcp-spanner://${projectId}/${config.instanceId}/${config.databaseId}/schema\`\n`;
        
        for (const row of tablesResult) {
          const tableName = (row as any).table_name as string || 'unknown';
          markdown += `- Table Preview: \`gcp-spanner://${projectId}/${config.instanceId}/${config.databaseId}/tables/${tableName}/preview\`\n`;
        }
        
        return {
          content: [{
            type: 'text',
            text: markdown
          }]
        };
      } catch (error: any) {
        console.error('Error listing Spanner tables:', error);
        throw error;
      }
    }
  );

  // Tool to list instances
  server.tool(
    'list-spanner-instances',
    // Define an empty schema with a dummy parameter that's optional
    // This ensures compatibility with clients that expect an object parameter
    {
      _dummy: z.string().optional().describe('Not used, just to ensure parameter compatibility')
    },
    async (_params, _extra) => {
      try {
        const projectId = await getProjectId();
        const spanner = getSpannerClient();
        
        const [instances] = await spanner.getInstances();
        
        if (!instances || instances.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Spanner Instances\n\nProject: ${projectId}\n\nNo instances found in the project.`
            }]
          };
        }
        
        let markdown = `# Spanner Instances\n\nProject: ${projectId}\n\n`;
        
        // Table header
        markdown += '| Instance ID | State | Config | Nodes |\n';
        markdown += '|-------------|-------|--------|-------|\n';
        
        // Table rows
        for (const instance of instances) {
          const metadata = instance.metadata || {};
          markdown += `| ${instance.id || 'unknown'} | ${metadata.state || 'unknown'} | ${metadata.config?.split('/').pop() || 'unknown'} | ${metadata.nodeCount || 'unknown'} |\n`;
        }
        
        // Add resource links for further exploration
        markdown += '\n## Available Resources\n\n';
        markdown += `- All Instances: \`gcp-spanner://${projectId}/instances\`\n`;
        
        for (const instance of instances) {
          markdown += `- Databases in ${instance.id}: \`gcp-spanner://${projectId}/${instance.id}/databases\`\n`;
        }
        
        return {
          content: [{
            type: 'text',
            text: markdown
          }]
        };
      } catch (error: any) {
        console.error('Error listing Spanner instances:', error);
        throw error;
      }
    }
  );

  // Tool to list databases
  server.tool(
    'list-spanner-databases',
    {
      instanceId: z.string().describe('Spanner instance ID')
    },
    async ({ instanceId }, _extra) => {
      try {
        const projectId = await getProjectId();
        const spanner = getSpannerClient();
        const instance = spanner.instance(Array.isArray(instanceId) ? instanceId[0] : instanceId);
        
        const [databases] = await instance.getDatabases();
        
        if (!databases || databases.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Spanner Databases\n\nProject: ${projectId}\nInstance: ${Array.isArray(instanceId) ? instanceId[0] : instanceId}\n\nNo databases found in the instance.`
            }]
          };
        }
        
        let markdown = `# Spanner Databases\n\nProject: ${projectId}\nInstance: ${Array.isArray(instanceId) ? instanceId[0] : instanceId}\n\n`;
        
        // Table header
        markdown += '| Database ID | State |\n';
        markdown += '|-------------|-------|\n';
        
        // Table rows
        for (const database of databases) {
          const metadata = database.metadata || {};
          markdown += `| ${database.id || 'unknown'} | ${metadata.state || 'unknown'} |\n`;
        }
        
        // Add resource links for further exploration
        markdown += '\n## Available Resources\n\n';
        
        for (const database of databases) {
          markdown += `- Tables in ${database.id}: \`gcp-spanner://${projectId}/${Array.isArray(instanceId) ? instanceId[0] : instanceId}/${database.id}/tables\`\n`;
          markdown += `- Schema for ${database.id}: \`gcp-spanner://${projectId}/${Array.isArray(instanceId) ? instanceId[0] : instanceId}/${database.id}/schema\`\n`;
        }
        
        return {
          content: [{
            type: 'text',
            text: markdown
          }]
        };
      } catch (error: any) {
        console.error('Error listing Spanner databases:', error);
        throw error;
      }
    }
  );

  // Tool to execute natural language queries against Spanner
  server.tool(
    'natural-language-spanner-query',
    {
      query: z.string().describe('Natural language description of the query you want to execute'),
      instanceId: z.string().optional().describe('Spanner instance ID (defaults to SPANNER_INSTANCE env var)'),
      databaseId: z.string().optional().describe('Spanner database ID (defaults to SPANNER_DATABASE env var)')
    },
    async ({ query, instanceId, databaseId }, _extra) => {
      try {
        const projectId = await getProjectId();
        const config = await getSpannerConfig(
          Array.isArray(instanceId) ? instanceId[0] : instanceId,
          Array.isArray(databaseId) ? databaseId[0] : databaseId
        );
        
        // Get the schema for the database
        const schemaInfo = await getDetailedSchemaForQueryGeneration(config.instanceId, config.databaseId);
        
        // Get a list of tables
        const spanner = getSpannerClient();
        const instance = spanner.instance(config.instanceId);
        const database = instance.database(config.databaseId);
        
        const [tablesResult] = await database.run({
          sql: `SELECT table_name FROM information_schema.tables 
                WHERE table_catalog = '' AND table_schema = ''
                ORDER BY table_name`
        });
        
        const tableNames = tablesResult.map((row: any) => row.table_name as string);
        
        if (!tableNames || tableNames.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Query Error\n\nNo tables found in database ${config.databaseId}.`
            }]
          };
        }
        
        // Generate a SQL query based on the natural language query and schema
        let generatedSql = '';
        
        // For simple queries about table structure, generate SQL directly
        if (query.toLowerCase().includes('list tables') || 
            query.toLowerCase().includes('show tables') || 
            query.toLowerCase().includes('what tables')) {
          generatedSql = `SELECT table_name FROM information_schema.tables 
                          WHERE table_catalog = '' AND table_schema = ''
                          ORDER BY table_name`;
        } 
        else if (query.toLowerCase().includes('schema') || 
                 query.toLowerCase().includes('structure') || 
                 query.toLowerCase().includes('columns')) {
          // Extract table name if specified
          let tableName = '';
          for (const name of tableNames) {
            if (query.toLowerCase().includes(name.toLowerCase())) {
              tableName = name;
              break;
            }
          }
          
          if (tableName) {
            generatedSql = `SELECT column_name, spanner_type, is_nullable 
                            FROM information_schema.columns 
                            WHERE table_catalog = '' AND table_schema = '' AND table_name = '${tableName}'
                            ORDER BY ordinal_position`;
          } else {
            generatedSql = `SELECT table_name, column_name, spanner_type 
                            FROM information_schema.columns 
                            WHERE table_catalog = '' AND table_schema = ''
                            ORDER BY table_name, ordinal_position`;
          }
        }
        // For simple 'show all data' queries
        else if (query.toLowerCase().includes('all data') || 
                 query.toLowerCase().includes('all rows')) {
          // Extract table name if specified
          let tableName = '';
          for (const name of tableNames) {
            if (query.toLowerCase().includes(name.toLowerCase())) {
              tableName = name;
              break;
            }
          }
          
          if (tableName) {
            generatedSql = `SELECT * FROM ${tableName} LIMIT 100`;
          } else {
            // If no specific table mentioned, return an error
            return {
              content: [{
                type: 'text',
                text: `# Query Error\n\nPlease specify which table you want to see data from. Available tables: ${tableNames.join(', ')}`
              }]
            };
          }
        }
        // For count queries
        else if (query.toLowerCase().includes('count') || 
                 query.toLowerCase().includes('how many')) {
          // Extract table name if specified
          let tableName = '';
          for (const name of tableNames) {
            if (query.toLowerCase().includes(name.toLowerCase())) {
              tableName = name;
              break;
            }
          }
          
          if (tableName) {
            generatedSql = `SELECT COUNT(*) as count FROM ${tableName}`;
          } else {
            // If no specific table mentioned, count rows in all tables
            const countQueries = tableNames.map(name => `SELECT '${name}' as table_name, COUNT(*) as row_count FROM ${name}`);
            generatedSql = countQueries.join(' UNION ALL ');
          }
        }
        // For more complex queries, provide schema information and ask the user to use execute-spanner-query
        else {
          return {
            content: [{
              type: 'text',
              text: `# Complex Query Detected\n\nYour query requires a custom SQL statement. Here's the database schema to help you formulate your query:\n\n\`\`\`\n${schemaInfo}\n\`\`\`\n\nPlease use the \`execute-spanner-query\` tool with a specific SQL statement to query this data.\n\nExample:\n\`\`\`sql\nSELECT * FROM [table_name] WHERE [condition] LIMIT 100\n\`\`\``
            }]
          };
        }
        
        // Execute the generated SQL query
        // Execute the generated SQL query
        const [result] = await database.run({
          sql: generatedSql
        });
        
        if (!result || result.length === 0) {
          return {
            content: [{
              type: 'text',
              text: `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\nNatural Language Query: ${query}\n\nGenerated SQL: \`${generatedSql}\`\n\nQuery executed successfully. No results returned.`
            }]
          };
        }
        
        // Convert to markdown table
        const columns = Object.keys(result[0]);
        
        let markdown = `# Query Results\n\nProject: ${projectId}\nInstance: ${config.instanceId}\nDatabase: ${config.databaseId}\n\n`;
        markdown += `Natural Language Query: ${query}\n\n`;
        markdown += `Generated SQL: \`${generatedSql}\`\n\n`;
        markdown += `Rows: ${result.length}\n\n`;
        
        // Table header
        markdown += '| ' + columns.join(' | ') + ' |\n';
        markdown += '| ' + columns.map(() => '---').join(' | ') + ' |\n';
        
        // Table rows (limit to 100 rows for display)
        const displayRows = result.slice(0, 100);
        for (const row of displayRows) {
          const rowValues = columns.map(col => {
            const value = (row as any)[col];
            if (value === null || value === undefined) return 'NULL';
            if (typeof value === 'object') return JSON.stringify(value);
            return String(value);
          });
          
          markdown += '| ' + rowValues.join(' | ') + ' |\n';
        }
        
        if (result.length > 100) {
          markdown += '\n*Results truncated. Showing 100 of ' + result.length + ' rows.*';
        }
        
        // Add a note about using execute-spanner-query for more complex queries
        markdown += '\n\n## Need a more complex query?\n\n';
        markdown += 'For more complex queries, use the `execute-spanner-query` tool with a specific SQL statement.';
        
        return {
          content: [{
            type: 'text',
            text: markdown
          }]
        };
      } catch (error: any) {
        console.error('Error executing natural language Spanner query:', error);
        return {
          content: [{
            type: 'text',
            text: `# Query Error\n\nFailed to execute query: ${error.message}\n\nIf this is a complex query, please use the \`execute-spanner-query\` tool with a specific SQL statement.`
          }]
        };
      }
    }
  );
}