from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import patch

import pytest
from amsdal.models.core.file import File
from amsdal_utils.models.enums import ModuleType

from amsdal_ml.mcp_server.server_model_explorer.services.models import ErrorArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import ExistingRefArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import ExplicitIdArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import LiteralArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import NewObjectArgAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import PlannedArgument
from amsdal_ml.mcp_server.server_model_explorer.services.models import PlannedArguments
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionParameterSpec
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionRunAnalysis
from amsdal_ml.mcp_server.server_model_explorer.services.models import TransactionSpec
from amsdal_ml.mcp_server.server_model_explorer.services.types import ModelInfo
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionArgumentMode
from amsdal_ml.mcp_server.server_model_explorer.services.types import TransactionParameterKind
from amsdal_ml.mcp_server.server_model_explorer.tools.execution import _analyze_transaction_request
from amsdal_ml.mcp_server.server_model_explorer.tools.execution import _materialize_transaction_args
from amsdal_ml.mcp_server.server_model_explorer.tools.execution import _plan_arguments
from amsdal_ml.mcp_server.server_model_explorer.tools.execution import perform_transaction_operation
from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.exceptions import MultipleRecordsError
from amsdal_ml.ml_retrievers.query_creator import CreateAnalysis
from amsdal_ml.ml_retrievers.query_retriever import RetrieveAnalysis
from tests.fixtures.models.book import Book

_PRIMITIVE_VALUE = 7


def _tx_spec(params: list[TransactionParameterSpec]) -> TransactionSpec:
    return TransactionSpec(
        name='Tx',
        description='d',
        scope='user',
        parameters=params,
    )


@pytest.mark.asyncio
async def test_analyze_transaction_disallows_new_object_for_file_model():
    transaction = TransactionSpec(
        name='AttachFile',
        description='',
        scope='user',
        parameters=[
            TransactionParameterSpec(
                name='file',
                annotation='File',
                kind=TransactionParameterKind.MODEL,
                model_name='File',
            )
        ],
    )

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._plan_arguments',
            new_callable=AsyncMock,
        ) as mock_plan,
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
        ) as mock_registry,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
    ):
        mock_plan.return_value = PlannedArguments(
            items=[PlannedArgument(name='file', mode='new_object', query='create file from document')]
        )
        mock_registry.return_value = {'File': ModelInfo(model=File, scope=ModuleType.USER)}

        analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='attach a new file',
            preview_limit=2,
        )

        assert analysis.is_executable is False
        assert (
            analysis.args[0].error
            == 'new_object mode is not supported for File model. Use explicit_id or existing_ref.'
        )


@pytest.mark.asyncio
async def test_analyze_transaction_allows_nullable_model_literal_null():
    transaction = TransactionSpec(
        name='ProbeOptionalModel',
        description='',
        scope='user',
        parameters=[
            TransactionParameterSpec(
                name='book',
                annotation='Book | None',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                nullable=True,
            )
        ],
    )

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._plan_arguments',
            new_callable=AsyncMock,
        ) as mock_plan,
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
        ) as mock_registry,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
    ):
        mock_plan.return_value = PlannedArguments(
            items=[PlannedArgument(name='book', mode='literal', literal_value=None)]
        )
        mock_registry.return_value = {'Book': ModelInfo(model=Book, scope=ModuleType.USER)}

        analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='book: null',
            preview_limit=2,
        )

        assert analysis.is_executable is True
        assert analysis.args[0].error is None
        assert isinstance(analysis.args[0], LiteralArgAnalysis)
        assert analysis.args[0].value is None


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ('kind', 'mode', 'literal_value', 'expected_summary', 'expected_error'),
    [
        (TransactionParameterKind.PRIMITIVE, TransactionArgumentMode.LITERAL, 42, 'ready', None),
        (
            TransactionParameterKind.PRIMITIVE,
            TransactionArgumentMode.EXISTING_REF,
            None,
            'needs_clarification',
            'Expected literal mode for parameter value',
        ),
        (TransactionParameterKind.DICT, TransactionArgumentMode.LITERAL, {'a': 1}, 'ready', None),
        (
            TransactionParameterKind.DICT,
            TransactionArgumentMode.NEW_OBJECT,
            None,
            'needs_clarification',
            'Expected literal mode for parameter value',
        ),
    ],
)
async def test_analyze_transaction_supported_literal_kinds(
    kind: TransactionParameterKind,
    mode: TransactionArgumentMode,
    literal_value: object,
    expected_summary: str,
    expected_error: str | None,
):
    transaction = TransactionSpec(
        name='ProbeKinds',
        description='',
        scope='user',
        parameters=[
            TransactionParameterSpec(
                name='value',
                annotation='Any',
                kind=kind,
            )
        ],
    )

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._plan_arguments',
            new_callable=AsyncMock,
        ) as mock_plan,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
    ):
        mock_plan.return_value = PlannedArguments(
            items=[PlannedArgument(name='value', mode=mode, literal_value=literal_value)]
        )

        analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='probe',
            preview_limit=2,
        )

        assert analysis.summary == expected_summary
        assert analysis.args[0].error == expected_error


@pytest.mark.asyncio
async def test_analyze_transaction_model_explicit_id_supported():
    transaction = TransactionSpec(
        name='ProbeModelExplicitId',
        description='',
        scope='user',
        parameters=[
            TransactionParameterSpec(
                name='book',
                annotation='Book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
            )
        ],
    )

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._plan_arguments',
            new_callable=AsyncMock,
        ) as mock_plan,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry') as mock_registry,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
    ):
        mock_plan.return_value = PlannedArguments(
            items=[PlannedArgument(name='book', mode='explicit_id', explicit_object_id='abc-123')]
        )
        mock_registry.return_value = {'Book': ModelInfo(model=Book, scope=ModuleType.USER)}

        analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='book id abc-123',
            preview_limit=2,
        )

        assert analysis.is_executable is True
        assert isinstance(analysis.args[0], ExplicitIdArgAnalysis)
        assert analysis.args[0].value == 'abc-123'


@pytest.mark.asyncio
async def test_analyze_transaction_model_existing_ref_and_new_object_supported():
    transaction = TransactionSpec(
        name='ProbeModelModes',
        description='',
        scope='user',
        parameters=[
            TransactionParameterSpec(
                name='book',
                annotation='Book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
            )
        ],
    )
    retrieval = RetrieveAnalysis(query='book', filters={}, total=1, records=[])
    create = CreateAnalysis(
        query='book',
        payloads=[{'fields': {}, 'missing_required': []}],
        total=1,
        records=[],
    )

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._plan_arguments',
            new_callable=AsyncMock,
        ) as mock_plan,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry') as mock_registry,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.authorize_class', new_callable=AsyncMock),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.NLQueryExecutor') as mock_retriever,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.NLQueryCreatorExecutor') as mock_creator,
    ):
        mock_registry.return_value = {'Book': ModelInfo(model=Book, scope=ModuleType.USER)}
        mock_retriever.return_value.analyze = AsyncMock(return_value=retrieval)
        mock_creator.return_value.analyze = AsyncMock(return_value=create)

        mock_plan.return_value = PlannedArguments(
            items=[PlannedArgument(name='book', mode='existing_ref', query='book')]
        )
        existing_analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='book existing',
            preview_limit=2,
        )

        mock_plan.return_value = PlannedArguments(items=[PlannedArgument(name='book', mode='new_object', query='book')])
        new_object_analysis = await _analyze_transaction_request(
            transaction=transaction,
            query='book new',
            preview_limit=2,
        )

        assert existing_analysis.is_executable is True
        assert isinstance(existing_analysis.args[0], ExistingRefArgAnalysis)
        assert existing_analysis.args[0].existing_analysis is not None
        assert existing_analysis.args[0].existing_analysis.total == 1
        assert new_object_analysis.is_executable is True
        assert isinstance(new_object_analysis.args[0], NewObjectArgAnalysis)
        assert new_object_analysis.args[0].create_analysis is not None
        assert new_object_analysis.args[0].create_analysis.total == 1


@pytest.mark.asyncio
async def test_perform_transaction_operation_catalog_error():
    with patch(
        'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
        new_callable=AsyncMock,
        side_effect=RuntimeError('catalog down'),
    ):
        result = await perform_transaction_operation('Tx', 'query')

    assert isinstance(result, str)
    assert result.startswith('Error loading transactions:')


@pytest.mark.asyncio
async def test_perform_transaction_operation_resolve_error():
    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name',
            side_effect=ValueError('not found'),
        ),
    ):
        result = await perform_transaction_operation('Tx', 'query')

    assert result == 'Error resolving transaction: not found'


@pytest.mark.asyncio
async def test_perform_transaction_operation_confirm_false_success():
    tx = _tx_spec([])
    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._analyze_transaction_request',
            new_callable=AsyncMock,
            return_value=TransactionRunAnalysis(
                transaction_name='Tx',
                args=[],
                is_executable=True,
                summary='ready',
            ),
        ),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=False)

    assert isinstance(result, dict)
    assert result['transaction_name'] == 'Tx'
    assert result['requires_confirmation'] is True


@pytest.mark.asyncio
async def test_perform_transaction_operation_confirm_true_non_executable():
    tx = _tx_spec([])
    analysis = {
        'transaction_name': 'Tx',
        'args': [],
        'is_executable': False,
        'summary': 'needs_clarification',
    }

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=True, analysis=analysis)

    assert isinstance(result, str)
    assert result.startswith('Error: transaction arguments are ambiguous')


@pytest.mark.asyncio
async def test_perform_transaction_operation_confirm_true_requires_analysis_payload():
    tx = _tx_spec([])

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=True, analysis=None)

    assert result == "Error: The 'analysis' payload is required when confirm=True."


@pytest.mark.asyncio
async def test_perform_transaction_operation_confirm_true_executes():
    tx = _tx_spec([])
    analysis = {
        'transaction_name': 'Tx',
        'args': [],
        'is_executable': True,
        'summary': 'ready',
    }

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._materialize_transaction_args',
            new_callable=AsyncMock,
            return_value={'a': 1},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.TransactionExecutionApi.execute_transaction',
            new_callable=AsyncMock,
            return_value={'ok': True},
        ),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=True, analysis=analysis)

    assert isinstance(result, dict)
    assert result['transaction_name'] == 'Tx'
    assert result['result'] == {'ok': True}


@pytest.mark.asyncio
async def test_perform_transaction_operation_confirm_true_validation_error():
    tx = _tx_spec([])

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=True, analysis={})

    assert isinstance(result, str)
    assert result.startswith('Error: invalid analysis payload:')


@pytest.mark.asyncio
async def test_perform_transaction_operation_generic_error():
    tx = _tx_spec([])
    analysis = {
        'transaction_name': 'Tx',
        'args': [],
        'is_executable': True,
        'summary': 'ready',
    }

    with (
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.collect_transaction_catalog',
            new_callable=AsyncMock,
            return_value={'Tx': tx},
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_transaction_name', return_value='Tx'),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._materialize_transaction_args',
            new_callable=AsyncMock,
            side_effect=RuntimeError('explode'),
        ),
    ):
        result = await perform_transaction_operation('Tx', 'query', confirm=True, analysis=analysis)

    assert isinstance(result, str)
    assert result.startswith('Error executing transaction:')


@pytest.mark.asyncio
async def test_materialize_args_core_error_and_success_paths():
    primitive_param = TransactionParameterSpec(
        name='p',
        annotation='int',
        kind=TransactionParameterKind.PRIMITIVE,
    )
    model_param = TransactionParameterSpec(
        name='book',
        annotation='Book',
        kind=TransactionParameterKind.MODEL,
        model_name='Book',
    )
    tx = _tx_spec([primitive_param, model_param])

    with patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()):
        with pytest.raises(ValueError, match='Missing analysis for parameter p'):
            await _materialize_transaction_args(
                transaction=tx,
                run_analysis=TransactionRunAnalysis(transaction_name='Tx', args=[], is_executable=True),
            )

    with patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()):
        with pytest.raises(ValueError, match='Argument p is not executable'):
            await _materialize_transaction_args(
                transaction=tx,
                run_analysis=TransactionRunAnalysis(
                    transaction_name='Tx',
                    args=[ErrorArgAnalysis(name='p', kind=TransactionParameterKind.PRIMITIVE, error='bad')],
                    is_executable=False,
                ),
            )

    run = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            LiteralArgAnalysis(
                name='p',
                kind=TransactionParameterKind.PRIMITIVE,
                value=_PRIMITIVE_VALUE,
            ),
            ExplicitIdArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                value='id-1',
            ),
        ],
        is_executable=True,
    )
    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
            return_value={'Book': ModelInfo(model=Book, scope=ModuleType.USER)},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_reference_resource',
            return_value='test-resource',
        ),
    ):
        args = await _materialize_transaction_args(transaction=tx, run_analysis=run)

    assert args['p'] == _PRIMITIVE_VALUE
    assert 'test-resource#Book:' in str(args['book'].ref)


@pytest.mark.asyncio
async def test_materialize_args_literal_nullable_and_invalid_mode():
    model_param = TransactionParameterSpec(
        name='book',
        annotation='Book | None',
        kind=TransactionParameterKind.MODEL,
        model_name='Book',
        nullable=True,
    )
    tx = _tx_spec([model_param])

    run_null = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            LiteralArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                value=None,
            )
        ],
        is_executable=True,
    )
    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
            return_value={'Book': ModelInfo(model=Book, scope=ModuleType.USER)},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_reference_resource',
            return_value='test-resource',
        ),
    ):
        args = await _materialize_transaction_args(transaction=tx, run_analysis=run_null)
    assert args['book'] is None

    run_bad = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            LiteralArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                value='x',
            )
        ],
        is_executable=False,
    )
    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
            return_value={'Book': ModelInfo(model=Book, scope=ModuleType.USER)},
        ),
    ):
        with pytest.raises(ValueError, match='Unsupported mode for book'):
            await _materialize_transaction_args(transaction=tx, run_analysis=run_bad)


@pytest.mark.asyncio
async def test_materialize_args_existing_and_new_object_paths():
    model_param = TransactionParameterSpec(
        name='book',
        annotation='Book',
        kind=TransactionParameterKind.MODEL,
        model_name='Book',
    )
    tx = _tx_spec([model_param])

    existing_run = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            ExistingRefArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                query='q',
                existing_analysis=RetrieveAnalysis(query='q', filters={}, total=1, records=[]),
            )
        ],
        is_executable=True,
    )
    new_run = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            NewObjectArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                query='q',
                create_analysis=CreateAnalysis(
                    query='q',
                    payloads=[{'fields': {}, 'missing_required': []}],
                    total=1,
                    records=[],
                ),
            )
        ],
        is_executable=True,
    )

    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
            return_value={'Book': ModelInfo(model=Book, scope=ModuleType.USER)},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_reference_resource',
            return_value='test-resource',
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.NLQueryExecutor') as mock_retriever,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.NLQueryCreatorExecutor') as mock_creator,
    ):
        mock_retriever.return_value.execute_one = AsyncMock(return_value=SimpleNamespace(object_id='b-1'))
        mock_creator.return_value.execute_one = AsyncMock(return_value=SimpleNamespace(object_id='b-2'))

        existing_args = await _materialize_transaction_args(transaction=tx, run_analysis=existing_run)
        new_args = await _materialize_transaction_args(transaction=tx, run_analysis=new_run)

    assert 'test-resource#Book:' in str(existing_args['book'].ref)
    assert 'test-resource#Book:' in str(new_args['book'].ref)
    assert str(existing_args['book'].ref) != str(new_args['book'].ref)


@pytest.mark.asyncio
async def test_materialize_args_multiple_records_errors_wrapped():
    model_param = TransactionParameterSpec(
        name='book',
        annotation='Book',
        kind=TransactionParameterKind.MODEL,
        model_name='Book',
    )
    tx = _tx_spec([model_param])

    existing_run = TransactionRunAnalysis(
        transaction_name='Tx',
        args=[
            ExistingRefArgAnalysis(
                name='book',
                kind=TransactionParameterKind.MODEL,
                model_name='Book',
                query='q',
                existing_analysis=RetrieveAnalysis(query='q', filters={}, total=2, records=[{'id': 1}]),
            )
        ],
        is_executable=False,
    )

    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=MagicMock()),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_model_registry',
            return_value={'Book': ModelInfo(model=Book, scope=ModuleType.USER)},
        ),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.resolve_reference_resource',
            return_value='test-resource',
        ),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.NLQueryExecutor') as mock_retriever,
    ):
        mock_retriever.return_value.execute_one = AsyncMock(
            side_effect=MultipleRecordsError('too many', total=2, records=[{'id': 1}, {'id': 2}])
        )
        with pytest.raises(ValueError, match='Expected 1 resolved object'):
            await _materialize_transaction_args(transaction=tx, run_analysis=existing_run)


@pytest.mark.asyncio
async def test_plan_arguments_json_schema_and_plain_text_paths():
    tx = _tx_spec(
        [
            TransactionParameterSpec(
                name='a',
                annotation='int',
                kind=TransactionParameterKind.PRIMITIVE,
            )
        ]
    )

    llm_schema = SimpleNamespace(
        supported_formats=[ResponseFormat.JSON_SCHEMA],
        ainvoke=AsyncMock(return_value='{"items": [{"name": "a", "mode": "literal", "literal_value": 1}]}'),
    )
    llm_plain = SimpleNamespace(
        supported_formats=[ResponseFormat.PLAIN_TEXT],
        ainvoke=AsyncMock(return_value='raw'),
    )

    prompt = SimpleNamespace(render_text=lambda **_: 'prompt')
    adapter = SimpleNamespace(get_response_schema=lambda schema: {'schema': schema})

    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=llm_schema),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.get_prompt', return_value=prompt),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.get_retriever_adapter', return_value=adapter),
    ):
        out_schema = await _plan_arguments(transaction=tx, query='q')

    assert out_schema.items[0].name == 'a'

    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=llm_plain),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.get_prompt', return_value=prompt),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.clean_json_response',
            return_value='{"items": []}',
        ),
    ):
        out_plain = await _plan_arguments(transaction=tx, query='q')

    assert out_plain.items == []


@pytest.mark.asyncio
async def test_plan_arguments_parse_error_wrapped():
    tx = _tx_spec(
        [
            TransactionParameterSpec(
                name='a',
                annotation='int',
                kind=TransactionParameterKind.PRIMITIVE,
            )
        ]
    )
    llm_plain = SimpleNamespace(
        supported_formats=[ResponseFormat.PLAIN_TEXT],
        ainvoke=AsyncMock(return_value='not-json'),
    )
    prompt = SimpleNamespace(render_text=lambda **_: 'prompt')

    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution._get_llm', return_value=llm_plain),
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.execution.get_prompt', return_value=prompt),
        patch(
            'amsdal_ml.mcp_server.server_model_explorer.tools.execution.clean_json_response',
            return_value='not-json',
        ),
    ):
        with pytest.raises(ValueError, match='Failed to parse transaction argument plan'):
            await _plan_arguments(transaction=tx, query='q')
