using System.ComponentModel;
using System.Globalization;
using System.Text;
using System.Text.RegularExpressions;
using ModelContextProtocol;
using ModelContextProtocol.Server;
using Npgsql;

/// <summary>The three tools. Every query runs as the role in SCHEMA_MCP_CONNECTION.</summary>
public sealed class SchemaTools(NpgsqlDataSource db)
{
    private const int MaxRows = 100;
    private const string StatementTimeout = "5s";

    [McpServerTool(Name = "list_tables", ReadOnly = true)]
    [Description("Lists the tables the database role can read, as schema.table.")]
    public async Task<string> ListTables(CancellationToken cancellationToken)
    {
        const string sql = """
            SELECT table_schema, table_name
            FROM information_schema.tables
            WHERE table_type = 'BASE TABLE'
              AND table_schema NOT IN ('pg_catalog', 'information_schema')
            ORDER BY table_schema, table_name
            """;
        var text = new StringBuilder();
        await using var command = db.CreateCommand(sql);
        await using var reader = await command.ExecuteReaderAsync(cancellationToken);
        while (await reader.ReadAsync(cancellationToken))
        {
            text.AppendLine($"{reader.GetString(0)}.{reader.GetString(1)}");
        }
        return text.Length == 0 ? "No tables visible to this role." : text.ToString();
    }

    [McpServerTool(Name = "describe_table", ReadOnly = true)]
    [Description("Describes one table: columns with types and nullability, primary and foreign keys, " +
                 "indexes, and the foreign keys in other tables that point at it. Names are case-sensitive.")]
    public async Task<string> DescribeTable(
        [Description("Table name exactly as stored, for example Appointments")] string table,
        [Description("Schema name; defaults to public")] string schema = "public",
        CancellationToken cancellationToken = default)
    {
        await using var connection = await db.OpenConnectionAsync(cancellationToken);

        await using var find = new NpgsqlCommand("SELECT to_regclass(format('%I.%I', $1, $2))::oid", connection);
        find.Parameters.AddWithValue(schema);
        find.Parameters.AddWithValue(table);
        if (await find.ExecuteScalarAsync(cancellationToken) is not uint relation)
        {
            throw new McpException($"Table {schema}.{table} not found. Call list_tables for the exact names.");
        }

        var text = new StringBuilder().AppendLine($"Table {schema}.{table}").AppendLine().AppendLine("Columns:");
        await AppendRows(connection, relation, text, """
            SELECT a.attname || ' ' || format_type(a.atttypid, a.atttypmod)
                   || CASE WHEN a.attnotnull THEN ' NOT NULL' ELSE ' NULL' END
                   || CASE a.attidentity WHEN 'a' THEN ' GENERATED ALWAYS AS IDENTITY'
                                         WHEN 'd' THEN ' GENERATED BY DEFAULT AS IDENTITY' ELSE '' END
                   || COALESCE(' DEFAULT ' || pg_get_expr(d.adbin, d.adrelid), '')
            FROM pg_attribute a
            LEFT JOIN pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
            WHERE a.attrelid = $1 AND a.attnum > 0 AND NOT a.attisdropped
            ORDER BY a.attnum
            """, cancellationToken);

        text.AppendLine().AppendLine("Constraints (primary key, foreign keys, unique, check):");
        await AppendRows(connection, relation, text, """
            SELECT conname || ': ' || pg_get_constraintdef(oid)
            FROM pg_constraint WHERE conrelid = $1 AND contype <> 'n'  -- PostgreSQL 18 lists NOT NULL here too
            ORDER BY CASE contype WHEN 'p' THEN 0 WHEN 'f' THEN 1 ELSE 2 END, conname
            """, cancellationToken);

        text.AppendLine().AppendLine("Indexes:");
        await AppendRows(connection, relation, text,
            "SELECT pg_get_indexdef(indexrelid) FROM pg_index WHERE indrelid = $1 ORDER BY indexrelid::regclass::text",
            cancellationToken);

        text.AppendLine().AppendLine("Referenced by:");
        await AppendRows(connection, relation, text, """
            SELECT conrelid::regclass::text || ' ' || conname || ': ' || pg_get_constraintdef(oid)
            FROM pg_constraint WHERE confrelid = $1 AND contype = 'f' ORDER BY 1
            """, cancellationToken);

        return text.ToString();
    }

    [McpServerTool(Name = "run_select", ReadOnly = true)]
    [Description("Runs ONE SELECT statement in a read-only transaction and returns at most 100 rows. " +
                 "Anything else (INSERT, UPDATE, DELETE, DDL, several statements) is refused. " +
                 "Quote mixed-case names: SELECT \"FullName\" FROM \"Doctors\".")]
    public async Task<string> RunSelect(
        [Description("A single SELECT statement")] string sql,
        CancellationToken cancellationToken)
    {
        // A convenience check with a clear message. It is not the safety net: the read-only
        // transaction and the role's privileges are (a reader who connects as an owner loses both).
        if (RefusalReason(sql) is { } reason)
        {
            throw new McpException($"Refused by SchemaMcp: {reason}");
        }

        try
        {
            await using var connection = await db.OpenConnectionAsync(cancellationToken);
            await using var transaction = await connection.BeginTransactionAsync(cancellationToken);
            await using (var setup = new NpgsqlCommand(
                $"SET TRANSACTION READ ONLY; SET LOCAL statement_timeout = '{StatementTimeout}'", connection, transaction))
            {
                await setup.ExecuteNonQueryAsync(cancellationToken);
            }

            await using var command = new NpgsqlCommand(sql, connection, transaction);
            await using var reader = await command.ExecuteReaderAsync(cancellationToken);

            var text = new StringBuilder();
            text.AppendLine(string.Join(" | ", Enumerable.Range(0, reader.FieldCount).Select(reader.GetName)));
            var rows = 0;
            while (rows < MaxRows && await reader.ReadAsync(cancellationToken))
            {
                text.AppendLine(string.Join(" | ", Enumerable.Range(0, reader.FieldCount)
                    .Select(i => Format(reader.GetValue(i)))));
                rows++;
            }
            var more = rows == MaxRows && await reader.ReadAsync(cancellationToken);
            text.Append(more ? $"(first {MaxRows} rows; more exist)" : $"({rows} rows)");
            return text.ToString(); // no commit: disposing the transaction rolls it back
        }
        catch (PostgresException e)
        {
            throw new McpException($"PostgreSQL refused the query: {e.SqlState}: {e.MessageText}");
        }
    }

    private static string? RefusalReason(string sql)
    {
        var text = sql.Trim().TrimEnd(';').TrimEnd();
        if (text.Contains(';'))
        {
            return "only one statement is allowed.";
        }
        return Regex.IsMatch(text, @"^SELECT\b", RegexOptions.IgnoreCase)
            ? null
            : "only a single SELECT statement is allowed.";
    }

    private static async Task AppendRows(NpgsqlConnection connection, uint relation, StringBuilder text,
        string sql, CancellationToken cancellationToken)
    {
        await using var command = new NpgsqlCommand(sql, connection);
        command.Parameters.Add(new NpgsqlParameter { Value = relation, DataTypeName = "oid" });
        await using var reader = await command.ExecuteReaderAsync(cancellationToken);
        var any = false;
        while (await reader.ReadAsync(cancellationToken))
        {
            text.AppendLine("  " + reader.GetString(0));
            any = true;
        }
        if (!any)
        {
            text.AppendLine("  (none)");
        }
    }

    private static string Format(object value) => value switch
    {
        DBNull => "NULL",
        DateTime d => d.ToString("O", CultureInfo.InvariantCulture),
        IFormattable f => f.ToString(null, CultureInfo.InvariantCulture),
        _ => value.ToString() ?? "",
    };
}