from __future__ import annotations

import json
import uuid
from typing import Any
from unittest.mock import AsyncMock

import pytest

from amsdal_ml.ml_retrievers.query_creator import NLQueryCreatorExecutor
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_creator_creates_record(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.author import Author

    query = 'create author Alice with email alice@example.com'
    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID,
                'name': 'Alice',
                'email': 'alice@example.com',
            },
            'missing_required': [],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=Author.objects.all(),
    )

    analysis = await executor.analyze(query, preview_limit=1)
    created: Author = await executor.execute_one(analysis)

    assert created.id == BASE_ID
    saved = await Author.objects.filter(id__eq=BASE_ID).aexecute()
    assert len(saved) == 1
    assert saved[0].name == 'Alice'
    assert saved[0].email == 'alice@example.com'


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

    query = 'create two authors Alice and Bob'
    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': [
                {
                    'id': BASE_ID + 1,
                    'name': 'Alice',
                    'email': 'alice@example.com',
                },
                {
                    'id': BASE_ID + 2,
                    'name': 'Bob',
                    'email': 'bob@example.com',
                },
            ],
            'missing_required': [],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=Author.objects.all(),
    )

    analysis = await executor.analyze(query, preview_limit=EXPECTED_TWO)
    assert analysis.total == EXPECTED_TWO
    assert len(analysis.payloads) == EXPECTED_TWO

    created: list[Author] = await executor.execute(analysis)
    assert len(created) == EXPECTED_TWO

    saved = await Author.objects.filter(id__eq=BASE_ID + 1).aexecute()
    assert len(saved) == 1
    assert saved[0].name == 'Alice'

    saved = await Author.objects.filter(id__eq=BASE_ID + 2).aexecute()
    assert len(saved) == 1
    assert saved[0].name == 'Bob'


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

    query = 'create author Bob without email'
    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': {},
            'missing_required': ['name', 'email'],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=Author.objects.all(),
    )

    analysis = await executor.analyze(query, preview_limit=1)
    with pytest.raises(ValueError):
        await executor.execute_one(analysis)

    results = await Author.objects.filter(name__eq='Bob').aexecute()
    assert results == []


@pytest.mark.asyncio
async def test_nl_query_creator_nested_reference(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.catalog_item import CatalogItem
    from tests.fixtures.models.category import Category

    category = Category(id=BASE_ID + 50, name='Electronics', description='Tech')
    await category.asave(force_insert=True)

    query = 'create catalog item smartphone in electronics'

    plan_response = json.dumps(
        {
            'candidates': [
                {
                    'field': 'category',
                    'entity_type': 'Category',
                    'search_keywords': ['Electronics'],
                }
            ]
        }
    )

    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID + 60,
                'name': 'Smartphone',
                'price': 499.0,
                'category': {
                    '$ref': {
                        'class_name': 'Category',
                        'object_id': category.id,
                        'resource': Category.__table_name__,
                    }
                },
            },
            'missing_required': [],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=CatalogItem.objects.all(),
    )

    analysis = await executor.analyze(query, preview_limit=1)
    created: CatalogItem = await executor.execute_one(analysis)

    assert created.id == BASE_ID + 60
    saved = await CatalogItem.objects.filter(id__eq=BASE_ID + 60).aexecute()
    assert len(saved) == 1
    item = saved[0]
    category_obj = await item.category
    assert category_obj is not None
    assert category_obj.id == category.id


@pytest.mark.asyncio
async def test_nl_query_creator_nested_missing_required(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from tests.fixtures.models.catalog_item import CatalogItem

    query = 'create catalog item without category'

    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID + 70,
                'name': 'Gadget',
                'price': 100.0,
            },
            'missing_required': ['category'],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=CatalogItem.objects.all(),
    )

    analysis = await executor.analyze(query, preview_limit=1)
    with pytest.raises(ValueError):
        await executor.execute_one(analysis)

    results = await CatalogItem.objects.filter(id__eq=BASE_ID + 70).aexecute()
    assert results == []


@pytest.mark.asyncio
async def test_nl_query_creator_file_download_success(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from amsdal.models.core.file import File
    from amsdal_models.classes.model import Model

    class TestFileModel(Model):
        __table_name__ = 'test_files'
        id: int
        file_field: File | None

    TestFileModel.model_rebuild()

    query = 'create test file with brochure'
    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID + 80,
                'file_field': {
                    'source_url': 'https://example.com/brochure.pdf',
                },
            },
            'missing_required': [],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=TestFileModel.objects.all(),
    )

    created = TestFileModel(id=BASE_ID + 80, file_field=File(filename='brochure.pdf', data=b'PDF content here'))
    executor._phase_persistence = AsyncMock(return_value=created)  # type: ignore[method-assign]

    analysis = await executor.analyze(query, preview_limit=1)
    created = await executor.execute_one(analysis)

    assert created.id == BASE_ID + 80
    assert created.file_field is not None
    assert created.file_field.filename == 'brochure.pdf'
    assert created.file_field.data == b'PDF content here'


@pytest.mark.asyncio
async def test_nl_query_creator_file_download_failure(async_test_amsdal_manager: Any) -> None:  # noqa: ARG001
    from amsdal.models.core.file import File
    from amsdal_models.classes.model import Model

    class TestFileModel(Model):
        __table_name__ = 'test_files'
        id: int
        file_field: File | None

    TestFileModel.model_rebuild()

    query = 'create test file with invalid brochure'
    plan_response = json.dumps({'candidates': []})
    payload_response = json.dumps(
        {
            'fields': {
                'id': BASE_ID + 90,
                'file_field': {
                    'source_url': 'https://invalid-url.com/missing.pdf',
                },
            },
            'missing_required': [],
        }
    )

    responses = {
        'Identify which nested entities': plan_response,
        'Generate a JSON payload': payload_response,
    }

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

    executor = NLQueryCreatorExecutor(
        llm=llm,
        queryset=TestFileModel.objects.all(),
    )

    created = TestFileModel(id=BASE_ID + 90, file_field=None)
    executor._phase_persistence = AsyncMock(return_value=created)  # type: ignore[method-assign]

    analysis = await executor.analyze(query, preview_limit=1)
    created = await executor.execute_one(analysis)

    assert created.id == BASE_ID + 90
    assert created.file_field is None
