package com.ruoyi.shadow.startup;

import org.junit.jupiter.api.Test;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.HashSet;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

class SqlServerPermissionProbeTest {
    private final SqlServerPermissionProbe probe =
            new SqlServerPermissionProbe();

    @Test
    void sourceProbeRejectsAnyProductionWritePermission() throws Exception {
        Map<String, Integer> evidence = acceptedSourcePermissions();
        evidence.put("update_database", 1);
        Connection connection = connectionWith(evidence,
                new ArrayList<String>());

        PermissionCheckResult result = probe.verifySource(connection);

        assertFalse(result.isAccepted());
        assertEquals("SOURCE_HAS_UPDATE_PERMISSION", result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsMissingSelectOnAnApprovedView()
            throws Exception {
        Map<String, Integer> evidence = acceptedSourcePermissions();
        evidence.put("select_v_device_srgv", 0);

        PermissionCheckResult result = probe.verifySource(
                connectionWith(evidence, new ArrayList<String>()));

        assertFalse(result.isAccepted());
        assertEquals("SOURCE_MISSING_SELECT_ON_VIEW:v_device_srgv",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsDirectSelectOnAnyProductionBaseTable()
            throws Exception {
        Map<String, Integer> evidence = acceptedSourcePermissions();
        evidence.put("select_device_srgv", 1);

        PermissionCheckResult result = probe.verifySource(
                connectionWith(evidence, new ArrayList<String>()));

        assertFalse(result.isAccepted());
        assertEquals("SOURCE_HAS_DIRECT_SELECT_PERMISSION:device_srgv",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsPrivilegedDatabaseRoleMembership()
            throws Exception {
        Map<String, Integer> evidence = acceptedSourcePermissions();
        evidence.put("role_db_datawriter", 1);

        PermissionCheckResult result = probe.verifySource(
                connectionWith(evidence, new ArrayList<String>()));

        assertFalse(result.isAccepted());
        assertEquals("SOURCE_IS_MEMBER_OF_FORBIDDEN_ROLE:db_datawriter",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsDbDataReaderMembership() throws Exception {
        Map<String, Integer> evidence = acceptedSourcePermissions();
        evidence.put("role_db_datareader", 1);

        PermissionCheckResult result = probe.verifySource(
                connectionWith(evidence, new ArrayList<String>()));

        assertFalse(result.isAccepted());
        assertEquals(
                "SOURCE_IS_MEMBER_OF_FORBIDDEN_ROLE:db_datareader",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsUnresolvedPermissionEvidence() throws Exception {
        Set<String> unresolved = new HashSet<String>();
        unresolved.add("update_database");

        PermissionCheckResult result = probe.verifySource(connectionWith(
                acceptedSourcePermissions(), new ArrayList<String>(),
                unresolved, null));

        assertFalse(result.isAccepted());
        assertEquals(
                "SOURCE_PERMISSION_PROBE_UNRESOLVED:update_database",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsEffectiveWritePermissionOnAnyOtherObject()
            throws Exception {
        EffectivePermission permission = new EffectivePermission(
                "OBJECT", "dbo.unapproved_table", "UPDATE");

        PermissionCheckResult result = probe.verifySource(connectionWith(
                acceptedSourcePermissions(), new ArrayList<String>(),
                new HashSet<String>(), permission));

        assertFalse(result.isAccepted());
        assertEquals(
                "SOURCE_HAS_EFFECTIVE_UPDATE_PERMISSION:OBJECT:"
                        + "dbo.unapproved_table",
                result.getReasonCode());
    }

    @Test
    void sourceProbeRejectsUnexpectedServerPermission() throws Exception {
        EffectivePermission serverPermission = new EffectivePermission(
                "SERVER", "SERVER", "CONTROL SERVER");

        PermissionCheckResult result = probe.verifySource(connectionWith(
                acceptedSourcePermissions(), new ArrayList<String>(),
                new HashSet<String>(), null, serverPermission));

        assertFalse(result.isAccepted());
        assertEquals("SOURCE_HAS_EFFECTIVE_SERVER_PERMISSION:CONTROL_SERVER",
                result.getReasonCode());
    }

    @Test
    void sourceProbeSetsReadOnlyIsolationAndExecutesSelectOnly()
            throws Exception {
        List<String> executedSql = new ArrayList<String>();
        Connection connection = connectionWith(acceptedSourcePermissions(),
                executedSql);

        PermissionCheckResult result = probe.verifySource(connection);

        assertTrue(result.isAccepted());
        assertEquals("SOURCE_PERMISSIONS_VERIFIED", result.getReasonCode());
        verify(connection).setReadOnly(true);
        verify(connection).setTransactionIsolation(
                Connection.TRANSACTION_READ_COMMITTED);
        for (String statement : executedSql) {
            assertTrue(statement.trim().toUpperCase().startsWith("SELECT"),
                    statement);
        }
    }

    @Test
    void sourceProbeExecutesEveryApprovedViewToValidateTheReadPath()
            throws Exception {
        List<String> executedSql = new ArrayList<String>();

        PermissionCheckResult result = probe.verifySource(connectionWith(
                acceptedSourcePermissions(), executedSql));

        assertTrue(result.isAccepted());
        for (String view : SqlServerPermissionProbe.REQUIRED_SOURCE_VIEWS) {
            assertTrue(executedSql.contains(
                    "SELECT TOP (0) * FROM [shadow_ro].[" + view + "]"),
                    "approved view was not executed: " + view);
        }
    }

    @Test
    void evidenceProbeRequiresWritePermissionOnEveryEvidenceTable()
            throws Exception {
        Map<String, Integer> permissions = acceptedEvidencePermissions();
        permissions.put("insert_shadow_cycle", 0);

        PermissionCheckResult result = probe.verifyEvidence(
                connectionWith(permissions, new ArrayList<String>()));

        assertFalse(result.isAccepted());
        assertEquals("EVIDENCE_MISSING_INSERT_ON_TABLE:shadow_cycle",
                result.getReasonCode());
    }

    @Test
    void evidenceProbeRejectsAnyExplicitDatabaseRoleMembership()
            throws Exception {
        PermissionCheckResult result = probe.verifyEvidence(
                connectionWithEvidenceExtras(acceptedEvidencePermissions(),
                        null, null, "custom_shadow_writer"));

        assertFalse(result.isAccepted());
        assertEquals("EVIDENCE_IS_MEMBER_OF_DATABASE_ROLE:custom_shadow_writer",
                result.getReasonCode());
    }

    @Test
    void evidenceProbeRejectsEffectivePermissionOnAnyOtherObject()
            throws Exception {
        EffectivePermission permission = new EffectivePermission(
                "OBJECT", "dbo.unapproved_table", "SELECT");

        PermissionCheckResult result = probe.verifyEvidence(
                connectionWithEvidenceExtras(acceptedEvidencePermissions(),
                        permission, null, null));

        assertFalse(result.isAccepted());
        assertEquals(
                "EVIDENCE_HAS_EFFECTIVE_SELECT_PERMISSION:OBJECT:"
                        + "dbo.unapproved_table",
                result.getReasonCode());
    }

    @Test
    void evidenceProbeRejectsUnexpectedServerPermission() throws Exception {
        EffectivePermission permission = new EffectivePermission(
                "SERVER", "SERVER", "CONTROL SERVER");

        PermissionCheckResult result = probe.verifyEvidence(
                connectionWithEvidenceExtras(acceptedEvidencePermissions(),
                        null, permission, null));

        assertFalse(result.isAccepted());
        assertEquals(
                "EVIDENCE_HAS_EFFECTIVE_SERVER_PERMISSION:CONTROL_SERVER",
                result.getReasonCode());
    }

    private static Connection connectionWith(
            Map<String, Integer> values,
            List<String> executedSql) throws Exception {
        return connectionWith(values, executedSql,
                new HashSet<String>(), null, null);
    }

    private static Connection connectionWith(
            Map<String, Integer> values,
            List<String> executedSql,
            Set<String> unresolvedAliases,
            EffectivePermission effectivePermission) throws Exception {
        return connectionWith(values, executedSql, unresolvedAliases,
                effectivePermission, null);
    }

    private static Connection connectionWith(
            Map<String, Integer> values,
            List<String> executedSql,
            Set<String> unresolvedAliases,
            EffectivePermission effectivePermission,
            EffectivePermission serverPermission) throws Exception {
        return connectionWith(values, executedSql, unresolvedAliases,
                effectivePermission, serverPermission, null);
    }

    private static Connection connectionWithEvidenceExtras(
            Map<String, Integer> values,
            EffectivePermission effectivePermission,
            EffectivePermission serverPermission,
            String databaseRole) throws Exception {
        return connectionWith(values, new ArrayList<String>(),
                new HashSet<String>(), effectivePermission, serverPermission,
                databaseRole);
    }

    private static Connection connectionWith(
            Map<String, Integer> values,
            List<String> executedSql,
            Set<String> unresolvedAliases,
            EffectivePermission effectivePermission,
            EffectivePermission serverPermission,
            String databaseRole) throws Exception {
        Connection connection = mock(Connection.class);
        DatabaseMetaData metadata = mock(DatabaseMetaData.class);
        Statement statement = mock(Statement.class);
        ResultSet permissionRow = mock(ResultSet.class);
        ResultSet databasePermissionRows = mock(ResultSet.class);
        ResultSet effectivePermissionRows = mock(ResultSet.class);
        ResultSet serverPermissionRows = mock(ResultSet.class);
        ResultSet databaseRoleRows = mock(ResultSet.class);

        when(connection.getMetaData()).thenReturn(metadata);
        when(metadata.getDatabaseProductName())
                .thenReturn("Microsoft SQL Server");
        when(metadata.getURL()).thenReturn(
                "jdbc:sqlserver://prod-db:1433;databaseName=FT546-wcs");
        when(metadata.getUserName()).thenReturn("ft546_shadow_ro");
        when(connection.createStatement()).thenReturn(statement);
        when(statement.executeQuery(anyString())).thenAnswer(invocation -> {
            String sql = invocation.getArgument(0);
            executedSql.add(sql);
            if (sql.contains("fn_my_permissions(NULL, 'SERVER')")) {
                return serverPermissionRows;
            }
            if (sql.contains("sys.database_role_members")) {
                return databaseRoleRows;
            }
            if (sql.contains("fn_my_permissions")) {
                if (sql.contains("scope_type")) {
                    return effectivePermissionRows;
                }
                return databasePermissionRows;
            }
            return permissionRow;
        });
        when(permissionRow.next()).thenReturn(true, false);
        when(permissionRow.getInt(anyString())).thenAnswer(invocation ->
                values.getOrDefault(invocation.getArgument(0), 0));
        when(permissionRow.getObject(anyString())).thenAnswer(invocation -> {
            String alias = invocation.getArgument(0);
            return unresolvedAliases.contains(alias)
                    ? null : values.getOrDefault(alias, 0);
        });
        when(databasePermissionRows.next()).thenReturn(false);
        when(effectivePermissionRows.next()).thenReturn(
                effectivePermission != null, false);
        if (effectivePermission != null) {
            when(effectivePermissionRows.getString("scope_type"))
                    .thenReturn(effectivePermission.scopeType);
            when(effectivePermissionRows.getString("securable_name"))
                    .thenReturn(effectivePermission.securableName);
            when(effectivePermissionRows.getString("permission_name"))
                    .thenReturn(effectivePermission.permissionName);
        }
        when(serverPermissionRows.next()).thenReturn(
                serverPermission != null, false);
        if (serverPermission != null) {
            when(serverPermissionRows.getString("permission_name"))
                    .thenReturn(serverPermission.permissionName);
        }
        when(databaseRoleRows.next()).thenReturn(databaseRole != null, false);
        if (databaseRole != null) {
            when(databaseRoleRows.getString("role_name"))
                    .thenReturn(databaseRole);
        }
        return connection;
    }

    private static final class EffectivePermission {
        private final String scopeType;
        private final String securableName;
        private final String permissionName;

        private EffectivePermission(String scopeType, String securableName,
                                    String permissionName) {
            this.scopeType = scopeType;
            this.securableName = securableName;
            this.permissionName = permissionName;
        }
    }

    private static Map<String, Integer> acceptedSourcePermissions() {
        Map<String, Integer> values = new HashMap<String, Integer>();
        for (String view : SqlServerPermissionProbe.REQUIRED_SOURCE_VIEWS) {
            values.put("select_" + view, 1);
        }
        return values;
    }

    private static Map<String, Integer> acceptedEvidencePermissions() {
        Map<String, Integer> values = new HashMap<String, Integer>();
        for (String table : SqlServerPermissionProbe.REQUIRED_EVIDENCE_TABLES) {
            values.put("select_" + table, 1);
            values.put("insert_" + table, 1);
            values.put("update_" + table, 1);
            values.put("delete_" + table, 1);
        }
        return values;
    }
}