import { z } from 'zod';
import type { ToolContext, ToolPrivilegeLevel } from './types.js';
import { handleSqlResponse, executeSqlWithFallback } from './utils.js';

// Output schema for constraints
const ListConstraintsOutputSchema = z.array(z.object({
    schema_name: z.string(),
    table_name: z.string(),
    constraint_name: z.string(),
    constraint_type: z.string(), // PRIMARY KEY, FOREIGN KEY, UNIQUE, CHECK, EXCLUDE
    columns: z.array(z.string()),
    definition: z.string(),
    is_deferrable: z.boolean(),
    initially_deferred: z.boolean(),
}));

// Input schema with optional filters
const ListConstraintsInputSchema = z.object({
    schema: z.string().optional().describe('Filter by schema name.'),
    table: z.string().optional().describe('Filter by table name.'),
    constraint_type: z.enum(['PRIMARY KEY', 'FOREIGN KEY', 'UNIQUE', 'CHECK', 'EXCLUDE']).optional().describe('Filter by constraint type.'),
    include_system: z.boolean().optional().default(false).describe('Include constraints in system schemas.'),
});
type ListConstraintsInput = z.infer<typeof ListConstraintsInputSchema>;

// Static JSON Schema for MCP capabilities
const mcpInputSchema = {
    type: 'object',
    properties: {
        schema: {
            type: 'string',
            description: 'Filter by schema name.',
        },
        table: {
            type: 'string',
            description: 'Filter by table name.',
        },
        constraint_type: {
            type: 'string',
            enum: ['PRIMARY KEY', 'FOREIGN KEY', 'UNIQUE', 'CHECK', 'EXCLUDE'],
            description: 'Filter by constraint type.',
        },
        include_system: {
            type: 'boolean',
            default: false,
            description: 'Include constraints in system schemas.',
        },
    },
    required: [],
};

// SQL identifier validation pattern
const identifierPattern = /^[a-zA-Z_][a-zA-Z0-9_$]*$/;

export const listConstraintsTool = {
    name: 'list_constraints',
    description: 'Lists all constraints (PRIMARY KEY, FOREIGN KEY, UNIQUE, CHECK, EXCLUDE) in the database. Can filter by schema, table, and type.',
    privilegeLevel: 'regular' as ToolPrivilegeLevel,
    inputSchema: ListConstraintsInputSchema,
    mcpInputSchema: mcpInputSchema,
    outputSchema: ListConstraintsOutputSchema,

    execute: async (input: ListConstraintsInput, context: ToolContext) => {
        const client = context.selfhostedClient;
        const { schema, table, constraint_type, include_system } = input;

        // Validate identifiers if provided
        if (schema && !identifierPattern.test(schema)) {
            throw new Error(`Invalid schema name: ${schema}`);
        }
        if (table && !identifierPattern.test(table)) {
            throw new Error(`Invalid table name: ${table}`);
        }

        // Build WHERE conditions
        const conditions: string[] = [];

        if (!include_system) {
            conditions.push("n.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'auth', 'storage', 'extensions', 'graphql', 'graphql_public', 'pgbouncer', 'realtime', 'supabase_functions', 'supabase_migrations', '_realtime')");
        }

        if (schema) {
            conditions.push(`n.nspname = '${schema}'`);
        }
        if (table) {
            conditions.push(`rel.relname = '${table}'`);
        }
        if (constraint_type) {
            // Use Map to prevent object injection attacks
            const typeMap = new Map<string, string>([
                ['PRIMARY KEY', 'p'],
                ['FOREIGN KEY', 'f'],
                ['UNIQUE', 'u'],
                ['CHECK', 'c'],
                ['EXCLUDE', 'x'],
            ]);
            const typeCode = typeMap.get(constraint_type);
            if (typeCode) {
                conditions.push(`c.contype = '${typeCode}'`);
            }
        }

        const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';

        const sql = `
            SELECT
                n.nspname AS schema_name,
                rel.relname AS table_name,
                c.conname AS constraint_name,
                CASE c.contype
                    WHEN 'p' THEN 'PRIMARY KEY'
                    WHEN 'f' THEN 'FOREIGN KEY'
                    WHEN 'u' THEN 'UNIQUE'
                    WHEN 'c' THEN 'CHECK'
                    WHEN 'x' THEN 'EXCLUDE'
                    ELSE c.contype::text
                END AS constraint_type,
                ARRAY(
                    SELECT a.attname
                    FROM unnest(c.conkey) WITH ORDINALITY AS k(attnum, ord)
                    JOIN pg_catalog.pg_attribute a ON a.attrelid = c.conrelid AND a.attnum = k.attnum
                    ORDER BY k.ord
                ) AS columns,
                pg_get_constraintdef(c.oid) AS definition,
                c.condeferrable AS is_deferrable,
                c.condeferred AS initially_deferred
            FROM pg_catalog.pg_constraint c
            JOIN pg_catalog.pg_class rel ON rel.oid = c.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = rel.relnamespace
            ${whereClause}
            ORDER BY n.nspname, rel.relname, c.conname
        `;

        const result = await executeSqlWithFallback(client, sql, true);
        return handleSqlResponse(result, ListConstraintsOutputSchema);
    },
};