# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the dsql_lint tool in server.py."""

import json
import pytest
import subprocess
from unittest.mock import AsyncMock, MagicMock
from awslabs.aurora_dsql_mcp_server.server import dsql_lint


@pytest.fixture
def ctx():
    """Mock MCP Context with awaitable error()."""
    c = MagicMock()
    c.error = AsyncMock()
    return c


def _clean_output():
    return json.dumps({
        'schema_version': 1,
        'files': [{
            'file': '<stdin>', 'diagnostics': [], 'error': None,
            'output_file': None, 'fixed_sql': None,
        }],
        'summary': {'errors': 0, 'warnings': 0, 'fixed': 0},
    })


class TestDsqlLint:
    """Tests for the dsql_lint MCP tool."""

    @pytest.mark.asyncio
    async def test_empty_sql_returns_empty_diagnostics(self, ctx, mocker):
        """Empty SQL input returns empty diagnostics without calling the binary."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        result = await dsql_lint('', ctx=ctx, fix=False)
        assert result['diagnostics'] == []
        assert result['fixed_sql'] is None
        assert result['summary'] == {'errors': 0, 'warnings': 0, 'fixed': 0}
        mock_run.assert_not_called()

    @pytest.mark.asyncio
    async def test_whitespace_only_returns_empty_diagnostics(self, ctx, mocker):
        """Whitespace-only SQL returns empty diagnostics without calling the binary."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        result = await dsql_lint('   \n\t  ', ctx=ctx, fix=False)
        assert result['diagnostics'] == []
        assert result['fixed_sql'] is None
        mock_run.assert_not_called()

    @pytest.mark.asyncio
    async def test_binary_not_found_raises_runtime_error(self, ctx, mocker):
        """Raises RuntimeError when dsql-lint binary is not on PATH."""
        mocker.patch(
            'awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value=None
        )

        with pytest.raises(RuntimeError, match=r'^dsql-lint binary not found on PATH\.'):
            await dsql_lint('CREATE TABLE t (id INT);', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_lint_clean_sql(self, ctx, mocker):
        """Clean SQL returns empty diagnostics."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(
            returncode=0,
            stdout=_clean_output(),
            stderr='',
        )
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        result = await dsql_lint(
            'CREATE TABLE t (id INT PRIMARY KEY);', ctx=ctx, fix=False
        )

        assert result['diagnostics'] == []
        assert result['fixed_sql'] is None
        assert result['summary'] == {'errors': 0, 'warnings': 0, 'fixed': 0}

    @pytest.mark.asyncio
    async def test_lint_detects_serial_type(self, ctx, mocker):
        """Detects SERIAL type usage and preserves full diagnostic fields."""
        diag = {
            'rule': 'serial_type',
            'line': 1,
            'message': 'Column `id` uses SERIAL, which is not supported in DSQL.',
            'suggestion': 'Use `BIGINT GENERATED BY DEFAULT AS IDENTITY (CACHE 1)` instead.',
            'fix_result': {
                'status': 'fixed_with_warning',
                'detail': 'Replaced SERIAL with BIGINT GENERATED BY DEFAULT AS IDENTITY',
            },
            'statement_preview': 'CREATE TABLE t (id SERIAL PRIMARY KEY)',
        }
        mock_output = json.dumps({
            'schema_version': 1,
            'files': [{
                'file': '<stdin>',
                'diagnostics': [diag],
                'error': None,
                'output_file': None,
                'fixed_sql': None,
            }],
            'summary': {'errors': 0, 'warnings': 1, 'fixed': 0},
        })

        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=1, stdout=mock_output, stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        result = await dsql_lint(
            'CREATE TABLE t (id SERIAL PRIMARY KEY);', ctx=ctx, fix=False
        )

        assert len(result['diagnostics']) == 1
        assert result['diagnostics'][0] == diag

    @pytest.mark.asyncio
    async def test_fix_mode_returns_fixed_sql(self, ctx, mocker):
        """Fix mode returns corrected SQL."""
        fixed_sql = (
            'CREATE TABLE t (\n  id BIGINT PRIMARY KEY GENERATED BY '
            'DEFAULT AS IDENTITY ( CACHE 1 ),\n  name VARCHAR(255)\n);\n'
        )
        mock_output = json.dumps({
            'schema_version': 1,
            'files': [{
                'file': '<stdin>',
                'diagnostics': [
                    {
                        'rule': 'serial_type',
                        'line': 1,
                        'message': 'Column `id` uses SERIAL.',
                        'suggestion': 'Use IDENTITY.',
                        'fix_result': {
                            'status': 'fixed_with_warning',
                            'detail': 'Replaced SERIAL',
                        },
                        'statement_preview': 'CREATE TABLE t ...',
                    },
                    {
                        'rule': 'collation',
                        'line': 1,
                        'message': 'Column `name` uses COLLATE "en_US".',
                        'suggestion': 'Remove the COLLATE clause; DSQL uses C collation database-wide.',
                        'fix_result': {
                            'status': 'fixed',
                            'detail': 'Removed COLLATE "en_US" from column `name`',
                        },
                        'statement_preview': 'CREATE TABLE t ...',
                    },
                ],
                'error': None,
                'output_file': None,
                'fixed_sql': fixed_sql,
            }],
            'summary': {'errors': 0, 'warnings': 1, 'fixed': 1},
        })

        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=3, stdout=mock_output, stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        result = await dsql_lint(
            'CREATE TABLE t (id SERIAL, name VARCHAR(255) COLLATE "en_US");',
            ctx=ctx, fix=True,
        )

        assert result['fixed_sql'] == fixed_sql
        assert len(result['diagnostics']) == 2

    @pytest.mark.asyncio
    async def test_fix_false_always_returns_none_fixed_sql(self, ctx, mocker):
        """When fix=False, fixed_sql is always None even if upstream populates it."""
        mock_output = json.dumps({
            'schema_version': 1,
            'files': [{
                'file': '<stdin>', 'diagnostics': [], 'error': None,
                'output_file': None,
                'fixed_sql': 'LEAKED UPSTREAM FIXED SQL',
            }],
            'summary': {'errors': 0, 'warnings': 0, 'fixed': 0},
        })
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=0, stdout=mock_output, stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        result = await dsql_lint('SELECT 1;', ctx=ctx, fix=False)

        assert result['fixed_sql'] is None

    @pytest.mark.asyncio
    async def test_command_and_stdin_payload(self, ctx, mocker):
        """Verify argv flags and that sql is piped via stdin (input=sql)."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=0, stdout=_clean_output(), stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        sql = 'SELECT 1;'
        await dsql_lint(sql, ctx=ctx, fix=True)

        cmd = mock_run.call_args.args[0]
        assert cmd[-1] == '-'
        assert '--fix' in cmd
        assert cmd[1:3] == ['--format', 'json']
        assert mock_run.call_args.kwargs['input'] == sql
        assert mock_run.call_args.kwargs.get('encoding') == 'utf-8'

    @pytest.mark.asyncio
    async def test_no_fix_flag_when_fix_false(self, ctx, mocker):
        """Verify --fix flag is NOT included when fix=False."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=0, stdout=_clean_output(), stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        await dsql_lint('SELECT 1;', ctx=ctx, fix=False)

        cmd = mock_run.call_args.args[0]
        assert '--fix' not in cmd

    @pytest.mark.asyncio
    async def test_usage_error_raises_runtime_error(self, ctx, mocker):
        """Exit code 2 (usage error) raises RuntimeError and surfaces stderr."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(
            returncode=2, stdout='', stderr='error: missing required argument',
        )
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'^dsql-lint usage error:.*missing required'):
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_unexpected_returncode_raises_runtime_error(self, ctx, mocker):
        """Return codes outside {0,1,3} (e.g. 139 segfault) raise even with valid JSON."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(
            returncode=139, stdout=_clean_output(), stderr='segfault',
        )
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'unexpected returncode=139'):
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_timeout_raises_runtime_error(self, ctx, mocker):
        """Timeout raises RuntimeError anchored on the exact message."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.side_effect = subprocess.TimeoutExpired(cmd='dsql-lint', timeout=30)
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'^dsql-lint timed out after 30 seconds$'):
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_os_error_raises_runtime_error(self, ctx, mocker):
        """OSError at exec time (e.g. PermissionError) raises RuntimeError chained from e."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.side_effect = PermissionError('Permission denied')
        mocker.patch(
            'awslabs.aurora_dsql_mcp_server.server.shutil.which',
            return_value='/usr/bin/dsql-lint',
        )

        with pytest.raises(RuntimeError, match=r'failed to execute:.*Permission denied') as exc_info:
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        assert isinstance(exc_info.value.__cause__, PermissionError)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_invalid_json_raises_runtime_error(self, ctx, mocker):
        """Invalid JSON output raises RuntimeError chained from JSONDecodeError."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(
            returncode=0, stdout='not valid json{{{', stderr='',
        )
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'returned invalid JSON') as exc_info:
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        assert isinstance(exc_info.value.__cause__, json.JSONDecodeError)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_empty_files_array_raises_runtime_error(self, ctx, mocker):
        """Empty files array is a protocol violation — raises rather than masking."""
        mock_output = json.dumps({
            'schema_version': 1, 'files': [],
            'summary': {'errors': 0, 'warnings': 0, 'fixed': 0},
        })
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=0, stdout=mock_output, stderr='')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'no file results'):
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_empty_stdout_raises_runtime_error(self, ctx, mocker):
        """Empty stdout from dsql-lint raises RuntimeError."""
        mock_run = mocker.patch('awslabs.aurora_dsql_mcp_server.server.subprocess.run')
        mock_run.return_value = MagicMock(returncode=0, stdout='', stderr='some error')
        mocker.patch('awslabs.aurora_dsql_mcp_server.server.shutil.which', return_value='/usr/bin/dsql-lint')

        with pytest.raises(RuntimeError, match=r'produced no output'):
            await dsql_lint('SELECT 1;', ctx=ctx, fix=False)
        ctx.error.assert_awaited_once()
