using System.Diagnostics;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Npgsql;
using TenantRls;

const string Host = "Host=127.0.0.1;Port=5433;Database=tenants;";
const string Admin = Host + "Username=postgres;Password=postgres";
const string Owner = Host + "Username=app_owner;Password=owner";
const string App = Host + "Username=app_user;Password=app";

var tenant1 = Guid.Parse("00000000-0000-0000-0000-000000000001");
var tenant2 = Guid.Parse("00000000-0000-0000-0000-000000000002");

if (args is ["setup"])
{
    var sql = await File.ReadAllTextAsync(Path.Combine(AppContext.BaseDirectory, "setup.sql"));
    await using var admin = new NpgsqlConnection(Admin);
    await admin.OpenAsync();
    var sw = Stopwatch.StartNew();
    await new NpgsqlCommand(sql, admin) { CommandTimeout = 600 }.ExecuteNonQueryAsync();
    await new NpgsqlCommand("VACUUM ANALYZE invoices", admin) { CommandTimeout = 600 }.ExecuteNonQueryAsync();
    var rows = await new NpgsqlCommand("SELECT count(*) FROM invoices", admin).ExecuteScalarAsync();
    Console.WriteLine($"setup done: {rows:N0} invoices in {sw.Elapsed.TotalSeconds:F1}s");
    return;
}

Section("1. Two layers: the EF query filter and the policy");
await using (var scope = Request(App, tenant1, interceptor: true))
{
    var db = scope.Db;
    var newest = await db.Invoices.OrderByDescending(i => i.CreatedAt).Take(3).ToListAsync();
    Console.WriteLine($"filtered read as tenant 1: {newest.Count} rows, tenants seen: {Distinct(newest)}");

    var leaked = await db.Invoices.IgnoreQueryFilters().CountAsync();
    Console.WriteLine($"IgnoreQueryFilters() as app_user: {leaked:N0} rows (table has 1,000,000)");
}

Section("2. Who bypasses the policy");
await Exec(Admin, "ALTER TABLE invoices NO FORCE ROW LEVEL SECURITY");
await using (var scope = Request(Owner, tenant1, interceptor: true))
    Console.WriteLine($"app_owner, ENABLE only: {await scope.Db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows");
await Exec(Admin, "ALTER TABLE invoices FORCE ROW LEVEL SECURITY");
await using (var scope = Request(Owner, tenant1, interceptor: true))
    Console.WriteLine($"app_owner, FORCE:       {await scope.Db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows");
await using (var scope = Request(Admin, tenant1, interceptor: true))
    Console.WriteLine($"postgres (superuser):   {await scope.Db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows");

Section("3. Writes: attach a stub for another tenant's invoice");
var foreignId = (Guid)(await Scalar(Admin, $"SELECT id FROM invoices WHERE tenant_id = '{tenant2}' LIMIT 1"))!;
await using (var scope = Request(App, tenant1, interceptor: true))
{
    var db = scope.Db;
    var stub = new Invoice { Id = foreignId, TenantId = tenant1 };
    db.Attach(stub);
    stub.Status = "Paid";
    try
    {
        await db.SaveChangesAsync();
        Console.WriteLine("tracked UPDATE succeeded (it should not)");
    }
    catch (DbUpdateConcurrencyException ex)
    {
        Console.WriteLine($"tracked UPDATE: DbUpdateConcurrencyException: {FirstSentence(ex.Message)}");
    }

    var rows = await db.Invoices.IgnoreQueryFilters()
        .Where(i => i.Id == foreignId)
        .ExecuteUpdateAsync(s => s.SetProperty(i => i.Status, "Paid"));
    Console.WriteLine($"ExecuteUpdate by id only: {rows} rows affected");
}
Console.WriteLine($"stored status of that invoice: {await Scalar(Admin, $"SELECT status FROM invoices WHERE id = '{foreignId}'")}");

Section("4. Writes: insert a row for another tenant");
await using (var scope = Request(App, tenant1, interceptor: true))
{
    var db = scope.Db;
    db.Invoices.Add(new Invoice
    {
        Id = Guid.NewGuid(), TenantId = tenant2, Number = "INV-999999",
        Amount = 10m, Status = "Open", CreatedAt = DateTimeOffset.UtcNow
    });
    try
    {
        await db.SaveChangesAsync();
        Console.WriteLine("INSERT succeeded (it should not)");
    }
    catch (DbUpdateException ex) when (ex.InnerException is PostgresException pg)
    {
        Console.WriteLine($"INSERT: {pg.SqlState} {pg.MessageText}");
    }
}

Section("5. No tenant resolved");
await using (var scope = Request(App, tenantId: null, interceptor: true))
{
    var db = scope.Db;
    Console.WriteLine($"current_setting in the session: '{await db.Database.SqlQueryRaw<string>("SELECT coalesce(current_setting('app.tenant_id', true), '<null>') AS \"Value\"").SingleAsync()}'");
    Console.WriteLine($"filtered read: {await db.Invoices.CountAsync()} rows, IgnoreQueryFilters(): {await db.Invoices.IgnoreQueryFilters().CountAsync()} rows");
}

Section("6. SET at the start of the request, then let EF open and close connections");
await using (var scope = Request(App, tenant1, interceptor: false))
{
    var db = scope.Db;
    await db.Database.ExecuteSqlRawAsync($"SET app.tenant_id = '{tenant1}'");
    Console.WriteLine($"SET, then a query on a pooled connection: {await db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows");

    await db.Database.OpenConnectionAsync();
    await db.Database.ExecuteSqlRawAsync($"SET app.tenant_id = '{tenant1}'");
    Console.WriteLine($"SET on a connection EF keeps open:       {await db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows");
    await db.Database.CloseConnectionAsync();
}

Section("7. What the next request sees on the same physical connection");
const string OnePhysical = ";Maximum Pool Size=1";
const string NoReset = ";Maximum Pool Size=1;No Reset On Close=true";
foreach (var (label, cs) in new[] { ("default (reset on close)", App + OnePhysical), ("No Reset On Close=true", App + NoReset) })
{
    await using (var first = Request(cs, tenant1, interceptor: false))
    {
        await first.Db.Database.OpenConnectionAsync();
        await first.Db.Database.ExecuteSqlRawAsync($"SET app.tenant_id = '{tenant1}'");
        await first.Db.Invoices.IgnoreQueryFilters().CountAsync();
        await first.Db.Database.CloseConnectionAsync();
    }
    await using (var second = Request(cs, tenantId: null, interceptor: false))
    {
        var seen = await second.Db.Database.SqlQueryRaw<string>("SELECT coalesce(current_setting('app.tenant_id', true), '<null>') AS \"Value\"").SingleAsync();
        var rows = await second.Db.Invoices.IgnoreQueryFilters().CountAsync();
        Console.WriteLine($"{label,-26} next request, no SET: setting = '{seen}', {rows:N0} rows visible");
    }
    NpgsqlConnection.ClearPool(new NpgsqlConnection(cs));
}
await using (var second = Request(App + NoReset, tenant2, interceptor: true))
{
    var seen = await second.Db.Database.SqlQueryRaw<string>("SELECT current_setting('app.tenant_id', true) AS \"Value\"").SingleAsync();
    Console.WriteLine($"{"with the interceptor",-26} next request as tenant 2: setting = '{seen}', {await second.Db.Invoices.IgnoreQueryFilters().CountAsync():N0} rows visible");
}
NpgsqlConnection.ClearPool(new NpgsqlConnection(App + NoReset));

Section("8. Plans: the policy is a predicate like any other");
await Explain(App, tenant1, """
    SELECT id, number, amount
    FROM invoices
    ORDER BY created_at DESC
    LIMIT 20
    """);
await Explain(App, tenant1, $"""
    SELECT id, number, amount
    FROM invoices
    WHERE tenant_id = '{tenant1}'
    ORDER BY created_at DESC
    LIMIT 20
    """);

Section("9. Plans: a predicate the planner cannot push below the policy");
await Explain(App, tenant1, """
    SELECT id, number, amount
    FROM invoices
    WHERE number LIKE 'INV-00000%'
    """);
await Exec(Admin, "ALTER TABLE invoices DISABLE ROW LEVEL SECURITY");
Console.WriteLine("-- same query with row-level security disabled --");
await Explain(App, tenant1, $"""
    SELECT id, number, amount
    FROM invoices
    WHERE number LIKE 'INV-00000%' AND tenant_id = '{tenant1}'
    """);
await Exec(Admin, "ALTER TABLE invoices ENABLE ROW LEVEL SECURITY");

Section("10. A composite index does not rescue LIKE, a leakproof operator does");
Console.WriteLine(await Scalar(Admin, """
    SELECT string_agg(p.proname || ' leakproof=' || p.proleakproof, ', ' ORDER BY p.proname)
    FROM pg_proc p WHERE p.proname IN ('textlike', 'starts_with')
    """));
await Exec(Admin, "CREATE INDEX ix_invoices_tenant_number ON invoices (tenant_id, number text_pattern_ops)");
await Exec(Admin, "ANALYZE invoices");
Console.WriteLine("-- LIKE again, now with an index on (tenant_id, number text_pattern_ops) --");
await Explain(App, tenant1, """
    SELECT id, number, amount
    FROM invoices
    WHERE number LIKE 'INV-00000%'
    """);
Console.WriteLine("-- the same prefix through the leakproof ^@ operator --");
await Explain(App, tenant1, """
    SELECT id, number, amount
    FROM invoices
    WHERE number ^@ 'INV-00000'
    """);
await using (var scope = Request(App, tenant1, interceptor: true))
{
    var starts = scope.Db.Invoices.Where(i => i.Number.StartsWith("INV-00000"));
    Console.WriteLine("-- what EF Core generates for StartsWith --");
    Console.WriteLine(starts.ToQueryString().Split('\n').Last().Trim());
    Console.WriteLine($"rows: {await starts.CountAsync()}");
}
await Exec(Admin, "DROP INDEX ix_invoices_tenant_number");

Section("11. What FORCE costs the owner during maintenance");
// pg_dump runs SET row_security = off and then COPY ... TO stdout. Both steps, as the owner:
foreach (var (label, sql) in new[]
{
    ("SET row_security = off (pg_dump default)", "SET row_security = off; SELECT count(*) FROM invoices"),
    ("migration backfill, no tenant set", "UPDATE invoices SET status = status WHERE number = 'INV-000001'"),
    ("COPY FROM (Npgsql binary import)", "COPY invoices FROM STDIN")
})
{
    try
    {
        await using var connection = new NpgsqlConnection(Owner);
        await connection.OpenAsync();
        var affected = await new NpgsqlCommand(sql, connection).ExecuteNonQueryAsync();
        Console.WriteLine($"{label,-41}: {affected} rows");
    }
    catch (PostgresException ex)
    {
        Console.WriteLine($"{label,-41}: {ex.SqlState} {ex.MessageText}");
    }
}
await using (var connection = new NpgsqlConnection(Owner))
{
    // What pg_dump --enable-row-security does: the policy stays on, so the dump succeeds and is empty.
    await connection.OpenAsync();
    var lines = 0;
    await using (var reader = await connection.BeginTextExportAsync("COPY invoices TO STDOUT"))
        while (await reader.ReadLineAsync() is not null) lines++;
    Console.WriteLine($"{"COPY TO with the policy on (dump)",-41}: {lines} rows exported of 1,000,000");
}

Section("12. Cost of set_config on every open");
var someId = (Guid)(await Scalar(Admin, $"SELECT id FROM invoices WHERE tenant_id = '{tenant1}' LIMIT 1"))!;
foreach (var withInterceptor in new[] { true, false, true, false })
{
    var sw = Stopwatch.StartNew();
    const int n = 500;
    for (var i = 0; i < n; i++)
    {
        await using var scope = Request(App, tenant1, withInterceptor);
        await scope.Db.Invoices.IgnoreQueryFilters().SingleOrDefaultAsync(x => x.Id == someId);
    }
    Console.WriteLine($"{n} requests, interceptor {(withInterceptor ? "on " : "off")}: {sw.Elapsed.TotalMilliseconds / n:F3} ms per request");
}

// ---- helpers ----

static void Section(string title) => Console.WriteLine($"\n== {title}");

static string Distinct(IEnumerable<Invoice> invoices) =>
    string.Join(", ", invoices.Select(i => i.TenantId.ToString()[^2..]).Distinct());

static string FirstSentence(string message) => message.Split(". ")[0] + ".";

static async Task Exec(string cs, string sql)
{
    await using var connection = new NpgsqlConnection(cs);
    await connection.OpenAsync();
    await new NpgsqlCommand(sql, connection).ExecuteNonQueryAsync();
}

static async Task<object?> Scalar(string cs, string sql)
{
    await using var connection = new NpgsqlConnection(cs);
    await connection.OpenAsync();
    return await new NpgsqlCommand(sql, connection).ExecuteScalarAsync();
}

static async Task Explain(string cs, Guid tenantId, string sql)
{
    await using var connection = new NpgsqlConnection(cs);
    await connection.OpenAsync();
    await using (var set = new NpgsqlCommand("SELECT set_config('app.tenant_id', @t, false)", connection))
    {
        set.Parameters.AddWithValue("t", tenantId.ToString());
        await set.ExecuteNonQueryAsync();
    }
    // warm cache, then explain
    await new NpgsqlCommand(sql, connection).ExecuteNonQueryAsync();
    await using var explain = new NpgsqlCommand("EXPLAIN (ANALYZE, COSTS OFF) " + sql, connection);
    await using var reader = await explain.ExecuteReaderAsync();
    Console.WriteLine(sql.Replace("\n", "\n").Trim());
    Console.WriteLine("---");
    while (await reader.ReadAsync())
        Console.WriteLine(reader.GetString(0));
    Console.WriteLine();
}

static RequestScope Request(string cs, Guid? tenantId, bool interceptor)
{
    var services = new ServiceCollection();
    services.AddScoped<TenantContext>();
    services.AddScoped<TenantConnectionInterceptor>();
    services.AddDbContext<AppDbContext>((sp, options) =>
    {
        options.UseNpgsql(cs);
        if (interceptor)
            options.AddInterceptors(sp.GetRequiredService<TenantConnectionInterceptor>());
    });
    var provider = services.BuildServiceProvider();
    var scope = provider.CreateAsyncScope();
    scope.ServiceProvider.GetRequiredService<TenantContext>().TenantId = tenantId;
    return new RequestScope(provider, scope);
}

sealed class RequestScope(ServiceProvider provider, AsyncServiceScope scope) : IAsyncDisposable
{
    public AppDbContext Db { get; } = scope.ServiceProvider.GetRequiredService<AppDbContext>();

    public async ValueTask DisposeAsync()
    {
        await scope.DisposeAsync();
        await provider.DisposeAsync();
    }
}