import pg from 'pg';

const { Pool } = pg;

const databaseUrl = process.env.DATABASE_URL || 'postgresql://postgres:postgres@127.0.0.1:54322/postgres';

const ids = {
  mainTenant: '00000000-0000-0000-0000-000000000001',
  partnerTenant: '00000000-0000-0000-0000-000000000901',
  mainUser: '00000000-0000-0000-0000-000000000101',
  partnerAdminUser: '00000000-0000-0000-0000-000000000907',
};

const readChecks = [
  {
    name: 'tenant_branding',
    sql: 'select tenant_id::text as tenant_id, brand_name as label from public.tenant_branding order by tenant_id',
    expectMain: true,
    expectPartner: true,
  },
  {
    name: 'tenant_settings',
    sql: 'select tenant_id::text as tenant_id, public_config::text as label from public.tenant_settings order by tenant_id',
    expectMain: true,
    expectPartner: true,
  },
  {
    name: 'tenant_domains',
    sql: 'select tenant_id::text as tenant_id, host as label from public.tenant_domains order by tenant_id',
    expectMain: true,
    expectPartner: true,
  },
  {
    name: 'tenant_memberships',
    sql: 'select tenant_id::text as tenant_id, role as label from public.tenant_memberships order by tenant_id, role',
    expectMain: true,
    expectPartner: true,
  },
  {
    name: 'regions',
    sql: 'select tenant_id::text as tenant_id, name as label from public.regions order by tenant_id',
    expectMain: true,
    expectPartner: true,
  },
  {
    name: 'questions',
    sql: 'select tenant_id::text as tenant_id, legacy_id as label from public.questions order by tenant_id, id',
    expectMain: true,
    expectPartner: false,
  },
  {
    name: 'student_profiles',
    sql: 'select tenant_id::text as tenant_id, user_id::text as label from public.student_profiles order by tenant_id, user_id',
    expectMain: true,
    expectPartner: false,
  },
  {
    name: 'orders',
    sql: 'select tenant_id::text as tenant_id, order_no as label from public.orders order by tenant_id',
    expectMain: true,
    expectPartner: false,
  },
  {
    name: 'content_assets',
    sql: 'select tenant_id::text as tenant_id, asset_key as label from public.content_assets order by tenant_id',
    expectMain: true,
    expectPartner: false,
  },
  {
    name: 'tenant_subscriptions',
    sql: 'select tenant_id::text as tenant_id, plan_code as label from public.tenant_subscriptions order by tenant_id',
    expectMain: false,
    expectPartner: true,
  },
  {
    name: 'tenant_invoices',
    sql: 'select tenant_id::text as tenant_id, invoice_no as label from public.tenant_invoices order by tenant_id',
    expectMain: false,
    expectPartner: true,
  },
  {
    name: 'tenant_usage_records',
    sql: 'select tenant_id::text as tenant_id, metric_key as label from public.tenant_usage_records order by tenant_id, metric_key',
    expectMain: false,
    expectPartner: true,
  },
];

const writeChecks = [
  {
    name: 'tenant_branding cross-tenant update',
    setupRole: 'authenticated',
    tenantId: ids.mainTenant,
    roleClaim: 'authenticated',
    sql: `
      update public.tenant_branding
         set slogan = slogan
       where tenant_id = $1::uuid
      returning tenant_id::text
    `,
    params: [ids.partnerTenant],
  },
  {
    name: 'regions cross-tenant insert',
    setupRole: 'authenticated',
    tenantId: ids.mainTenant,
    roleClaim: 'authenticated',
    sql: `
      insert into public.regions (tenant_id, legacy_id, name, code, sort_order, is_active)
      values ($1::uuid, 'rls-cross-tenant-probe', 'RLS 越权探针', 'RLS-X', 9999, false)
      returning tenant_id::text
    `,
    params: [ids.partnerTenant],
    expectErrorCode: '42501',
  },
];

const probeGrantStatements = [
  'grant usage on schema app to authenticated, anon',
  'grant usage on schema public to authenticated, anon',
  `grant select on
    public.tenant_branding,
    public.tenant_settings,
    public.tenant_domains,
    public.tenant_memberships,
    public.regions,
    public.questions,
    public.student_profiles,
    public.orders,
    public.content_assets,
    public.tenant_subscriptions,
    public.tenant_invoices,
    public.tenant_usage_records
   to authenticated, anon`,
  'grant update on public.tenant_branding to authenticated',
  'grant insert on public.regions to authenticated',
];

const pool = new Pool({ connectionString: databaseUrl });
const results = [];

function pass(name, detail = {}) {
  results.push({ status: 'pass', name, detail });
  console.log(`PASS ${name}`);
}

function fail(name, message, detail = {}) {
  results.push({ status: 'fail', name, message, detail });
  console.error(`FAIL ${name}: ${message}`);
  if (Object.keys(detail).length > 0) console.error(JSON.stringify(detail, null, 2));
}

function assert(condition, name, message, detail = {}) {
  if (condition) pass(name, detail);
  else fail(name, message, detail);
}

function onlyTenantRows(rows, tenantId) {
  return rows.every(row => row.tenant_id === tenantId);
}

function hasTenantRows(rows, tenantId) {
  return rows.some(row => row.tenant_id === tenantId);
}

async function withRlsContext(client, { dbRole = 'authenticated', tenantId = '', roleClaim = 'authenticated', sub = '' }, action) {
  await client.query('begin');
  try {
    for (const statement of probeGrantStatements) await client.query(statement);
    await client.query(`set local role ${dbRole}`);
    if (tenantId) {
      await client.query("select set_config('request.jwt.claim.tenant_id', $1, true)", [tenantId]);
    } else {
      await client.query("select set_config('request.jwt.claim.tenant_id', '', true)");
    }
    await client.query("select set_config('request.jwt.claim.role', $1, true)", [roleClaim]);
    await client.query("select set_config('request.jwt.claim.app_role', $1, true)", [roleClaim]);
    if (sub) await client.query("select set_config('request.jwt.claim.sub', $1, true)", [sub]);

    const value = await action();
    await client.query('rollback');
    return value;
  } catch (error) {
    await client.query('rollback').catch(() => {});
    throw error;
  }
}

async function queryAs(client, context, sql, params = []) {
  return withRlsContext(client, context, async () => {
    const result = await client.query(sql, params);
    return result.rows;
  });
}

async function runReadIsolationChecks(client) {
  for (const check of readChecks) {
    const mainRows = await queryAs(client, { tenantId: ids.mainTenant }, check.sql);
    assert(
      onlyTenantRows(mainRows, ids.mainTenant),
      `rls.read.${check.name}.main_no_leak`,
      '主租户上下文不应看到其它租户数据',
      { rows: mainRows },
    );
    assert(
      hasTenantRows(mainRows, ids.mainTenant) === check.expectMain,
      `rls.read.${check.name}.main_expected_seed`,
      '主租户 seed 数据存在性不符合预期',
      { expected: check.expectMain, rows: mainRows },
    );

    const partnerRows = await queryAs(client, { tenantId: ids.partnerTenant }, check.sql);
    assert(
      onlyTenantRows(partnerRows, ids.partnerTenant),
      `rls.read.${check.name}.partner_no_leak`,
      '伙伴租户上下文不应看到其它租户数据',
      { rows: partnerRows },
    );
    assert(
      hasTenantRows(partnerRows, ids.partnerTenant) === check.expectPartner,
      `rls.read.${check.name}.partner_expected_seed`,
      '伙伴租户 seed 数据存在性不符合预期',
      { expected: check.expectPartner, rows: partnerRows },
    );

    const anonymousRows = await queryAs(client, { dbRole: 'anon', tenantId: '', roleClaim: 'anon' }, check.sql);
    assert(
      anonymousRows.length === 0,
      `rls.read.${check.name}.no_tenant_claim_empty`,
      '没有 tenant_id claim 的上下文不应看到租户数据',
      { rows: anonymousRows },
    );
  }
}

async function runPlatformAdminChecks(client) {
  for (const check of readChecks) {
    const rows = await queryAs(client, { tenantId: '', roleClaim: 'platform_admin', sub: ids.mainUser }, check.sql);
    const hasMain = hasTenantRows(rows, ids.mainTenant);
    const hasPartner = hasTenantRows(rows, ids.partnerTenant);
    assert(
      hasMain === check.expectMain && hasPartner === check.expectPartner,
      `rls.platform_admin.${check.name}.cross_tenant_visibility`,
      '平台管理员 RLS 旁路应只暴露当前表已有的多租户 seed 数据',
      { expectedMain: check.expectMain, expectedPartner: check.expectPartner, rows },
    );
  }
}

async function runWriteIsolationChecks(client) {
  for (const check of writeChecks) {
    try {
      const rows = await queryAs(
        client,
        { dbRole: check.setupRole, tenantId: check.tenantId, roleClaim: check.roleClaim },
        check.sql,
        check.params,
      );

      assert(
        rows.length === 0,
        `rls.write.${check.name}`,
        '跨租户写入不应成功返回任何行',
        { rows },
      );
    } catch (error) {
      if (check.expectErrorCode && error.code === check.expectErrorCode) {
        pass(`rls.write.${check.name}`, { expectedErrorCode: error.code });
      } else {
        fail(`rls.write.${check.name}`, error.message, { code: error.code });
      }
    }
  }
}

async function verifySeed(client) {
  const result = await client.query(
    `
    select tenant_id::text, count(*)::int as count
    from public.tenant_branding
    where tenant_id in ($1::uuid, $2::uuid)
    group by tenant_id
    `,
    [ids.mainTenant, ids.partnerTenant],
  );
  const counts = new Map(result.rows.map(row => [row.tenant_id, Number(row.count)]));
  assert(
    counts.get(ids.mainTenant) === 1 && counts.get(ids.partnerTenant) === 1,
    'rls.seed.main_and_partner_tenants',
    '需要先运行 npm run db:smoke-seed，确保主租户和伙伴租户 seed 都存在',
    { rows: result.rows },
  );
}

async function main() {
  const client = await pool.connect();
  try {
    await verifySeed(client);
    await runReadIsolationChecks(client);
    await runPlatformAdminChecks(client);
    await runWriteIsolationChecks(client);
  } finally {
    client.release();
    await pool.end();
  }

  const failed = results.filter(item => item.status === 'fail');
  console.log(`\nRLS tenant isolation checks: ${results.length - failed.length} passed, ${failed.length} failed.`);
  if (failed.length > 0) process.exitCode = 1;
}

main().catch(error => {
  console.error(error);
  process.exitCode = 1;
});
