package com.sonymusic

import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test

import static org.junit.jupiter.api.Assertions.assertNull

class SnowcliTest extends BaseGlobalVarTest {
    def snowcli

    @BeforeEach
    void setUp() {
        super.setUp()
        snowcli = loadScript('vars/snowcli.groovy')
        helper.registerAllowedMethod('sh')
        helper.registerAllowedMethod('withEcr', [Map, Closure])
        binding.setVariable('env', [:])
    }

    @Test
    void testCallWithMissingCommand() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'command'",
            {
                snowcli.call(
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingAccount() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'account'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingUser() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'user'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    database: 'database',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingDatabase() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'database'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingSchema() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'schema'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingRole() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'role'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingWarehouse() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "snowcli: Missing required parameter: 'warehouse'",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    role: 'role'
                )
            }
        )
    }

    @Test
    void testCallWithMissingPrivateKey() {
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertThrowsWithMessage(
            IllegalStateException.class,
            "Environment variables SNOWFLAKE_PRIVATE_KEY_RAW and PRIVATE_KEY_PASSPHRASE must be set for authentication.",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithMissingPrivateKeyPassphrase() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"

        assertThrowsWithMessage(
            IllegalStateException.class,
            "Environment variables SNOWFLAKE_PRIVATE_KEY_RAW and PRIVATE_KEY_PASSPHRASE must be set for authentication.",
            {
                snowcli.call(
                    command: 'sql -f test.sql',
                    account: 'account',
                    user: 'user',
                    database: 'database',
                    schema: 'schema',
                    role: 'role',
                    warehouse: 'warehouse'
                )
            }
        )
    }

    @Test
    void testCallWithAllParameters() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertNull(
            snowcli(
                command: 'sql -f test.sql',
                account: 'account',
                user: 'user',
                database: 'database',
                schema: 'schema',
                role: 'role',
                warehouse: 'warehouse'
            )
        )

        assertMethodCalledOnceWith('sh', [
            "-e SNOWFLAKE_ACCOUNT=account",
            "-e SNOWFLAKE_USER=user",
            "-e SNOWFLAKE_DATABASE=database",
            "-e SNOWFLAKE_SCHEMA=schema",
            "-e SNOWFLAKE_ROLE=role",
            "-e SNOWFLAKE_WAREHOUSE=warehouse",
            "-e SNOWFLAKE_PRIVATE_KEY_RAW",
            "-e PRIVATE_KEY_PASSPHRASE",
            "-e SNOWFLAKE_AUTHENTICATOR=SNOWFLAKE_JWT",
            "sql -f test.sql --temporary-connection"
        ])
    }

    @Test
    void testCallWithCustomImage() {
        binding.getVariable("env")["SNOWFLAKE_PRIVATE_KEY_RAW"] = "private_key"
        binding.getVariable("env")["PRIVATE_KEY_PASSPHRASE"] = "passphrase"

        assertNull(
            snowcli(
                command: 'sql -f test.sql',
                account: 'account',
                user: 'user',
                database: 'database',
                schema: 'schema',
                role: 'role',
                warehouse: 'warehouse',
                snowcliImage: '123456789012.dkr.ecr.us-west-2.amazonaws.com/custom-snowcli:latest'
            )
        )

        assertMethodCalledOnceWith('sh', [
            "-e SNOWFLAKE_ACCOUNT=account",
            "-e SNOWFLAKE_USER=user",
            "-e SNOWFLAKE_DATABASE=database",
            "-e SNOWFLAKE_SCHEMA=schema",
            "-e SNOWFLAKE_ROLE=role",
            "-e SNOWFLAKE_WAREHOUSE=warehouse",
            "-e SNOWFLAKE_PRIVATE_KEY_RAW",
            "-e PRIVATE_KEY_PASSPHRASE",
            "-e SNOWFLAKE_AUTHENTICATOR=SNOWFLAKE_JWT",
            "123456789012.dkr.ecr.us-west-2.amazonaws.com/custom-snowcli:latest",
            "sql -f test.sql --temporary-connection"
        ])
    }
}