from __future__ import annotations

import json
import uuid
from typing import Any

import pytest

from amsdal_ml.ml_retrievers.exceptions import MultipleRecordsError
from amsdal_ml.ml_retrievers.query_updater import AmbiguousUpdateIntentError
from amsdal_ml.ml_retrievers.query_updater import NLQueryUpdater
from tests.agents_tests.test_fakes import FakeModel


def get_base_id() -> int:
    return uuid.uuid4().int % (10**9)


BASE_ID = get_base_id()

EXPECTED_TWO = 2


@pytest.mark.asyncio
async def test_nl_query_updater_applies_not_deleted_base_filter(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    llm = FakeModel(async_mode=True, responses={})
    llm.setup()

    updater = NLQueryUpdater(llm=llm, queryset=Author.objects.all())

    assert updater.executor.base_filters.get('_metadata__is_deleted') is False


@pytest.mark.asyncio
async def test_nl_query_updater_updates_single_record(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    author = Author(id=BASE_ID + 1, name='Alice', email='alice@example.com')
    await author.asave(force_insert=True)

    query = 'update author Alice: set email to alice2@example.com'

    split_response = json.dumps(
        {
            'find_query': 'find author with name Alice',
            'update_query': 'set email to alice2@example.com',
        }
    )

    filter_response = json.dumps(
        {
            'filters': [
                {
                    'field': 'name',
                    'lookup': 'eq',
                    'value': 'Alice',
                }
            ]
        }
    )

    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID + 1,
                'name': 'Alice',
                'email': 'alice2@example.com',
            },
            'missing_required': [],
        }
    )

    llm = FakeModel(
        async_mode=True,
        scripted=[split_response, filter_response, payload_response, payload_response],
    )
    llm.setup()

    updater = NLQueryUpdater(llm=llm, queryset=Author.objects.all())
    analysis = await updater.analyze(query, preview_limit=5)
    docs = await updater.invoke(analysis)
    assert len(docs) == 1
    updated = json.loads(docs[0].page_content)

    assert updated['id'] == str(BASE_ID + 1)
    assert updated['email'] == 'alice2@example.com'

    saved = await Author.objects.filter(id__eq=BASE_ID + 1).aexecute()
    assert len(saved) == 1
    assert saved[0].email == 'alice2@example.com'


@pytest.mark.asyncio
async def test_nl_query_updater_no_matches(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    query = 'update author Bob: set email to bob@example.com'

    split_response = json.dumps(
        {
            'find_query': 'find author with name Bob',
            'update_query': 'set email to bob@example.com',
        }
    )

    filter_response = json.dumps(
        {
            'filters': [
                {
                    'field': 'name',
                    'lookup': 'eq',
                    'value': 'Bob',
                }
            ]
        }
    )

    payload_response_1 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    payload_response_2 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    llm = FakeModel(
        async_mode=True,
        scripted=[split_response, filter_response, payload_response_1, payload_response_2],
    )
    llm.setup()

    updater = NLQueryUpdater(llm=llm, queryset=Author.objects.all())

    with pytest.raises(AmbiguousUpdateIntentError):
        await updater.analyze(query, preview_limit=5)


@pytest.mark.asyncio
async def test_nl_query_updater_too_many_matches(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    author1 = Author(id=BASE_ID + 10, name='Alice', email='alice1@example.com')
    author2 = Author(id=BASE_ID + 11, name='Alice', email='alice2@example.com')
    await author1.asave(force_insert=True)
    await author2.asave(force_insert=True)

    query = 'update author Alice: set email to alice3@example.com'

    split_response = json.dumps(
        {
            'find_query': 'find author with name Alice',
            'update_query': 'set email to alice3@example.com',
        }
    )

    filter_response = json.dumps(
        {
            'filters': [
                {
                    'field': 'name',
                    'lookup': 'eq',
                    'value': 'Alice',
                }
            ]
        }
    )

    payload_response_1 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )
    payload_response_2 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    llm = FakeModel(
        async_mode=True,
        scripted=[split_response, filter_response, payload_response_1, payload_response_2],
    )
    llm.setup()

    updater = NLQueryUpdater(llm=llm, queryset=Author.objects.all())

    analysis = await updater.analyze(query, preview_limit=5)
    with pytest.raises(MultipleRecordsError) as exc:
        await updater.invoke_one(analysis)

    assert exc.value.total == EXPECTED_TWO
    assert exc.value.records


@pytest.mark.asyncio
async def test_nl_query_updater_analyze_and_execute_multiple_records(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    author1 = Author(id=BASE_ID + 20, name='Alice', email='alice1@example.com')
    author2 = Author(id=BASE_ID + 21, name='Alice', email='alice2@example.com')
    await author1.asave(force_insert=True)
    await author2.asave(force_insert=True)

    query = 'update author Alice: set email to alice3@example.com'

    split_response = json.dumps(
        {
            'find_query': 'find author with name Alice',
            'update_query': 'set email to alice3@example.com',
        }
    )

    filter_response = json.dumps(
        {
            'filters': [
                {
                    'field': 'name',
                    'lookup': 'eq',
                    'value': 'Alice',
                }
            ]
        }
    )

    payload_response_1 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    payload_response_2 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    payload_response_3 = json.dumps(
        {
            'fields': {
                'email': 'alice3@example.com',
            },
            'missing_required': [],
        }
    )

    llm = FakeModel(
        async_mode=True,
        scripted=[
            split_response,
            filter_response,
            payload_response_1,
            payload_response_2,
            payload_response_3,
        ],
    )
    llm.setup()

    updater = NLQueryUpdater(llm=llm, queryset=Author.objects.all())

    preview = await updater.analyze(query, preview_limit=1)
    assert preview.total == EXPECTED_TWO
    assert len(preview.records) == 1

    updated_docs = await updater.invoke(preview)
    assert len(updated_docs) == EXPECTED_TWO

    saved = await Author.objects.filter(id__eq=BASE_ID + 20).aexecute()
    assert saved[0].email == 'alice3@example.com'
    saved = await Author.objects.filter(id__eq=BASE_ID + 21).aexecute()
    assert saved[0].email == 'alice3@example.com'
