from __future__ import annotations

import json
import uuid
from typing import Any

import pytest

from amsdal_ml.ml_models.utils import ResponseFormat
from amsdal_ml.ml_retrievers.query_creator import CreateIntent
from amsdal_ml.ml_retrievers.query_creator import NLQueryCreator
from tests.agents_tests.test_fakes import FakeModel


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


CASE_NAMES = ['flat', 'nested']


async def seed_category(base_id: int) -> int:
    from tests.fixtures.models.category import Category

    category = Category(id=base_id + 10, name='Electronics', description='Tech')
    await category.asave(force_insert=True)
    return category.id


def build_flat_payload(base_id: int) -> dict[str, Any]:
    return {
        'fields': {
            'id': base_id + 1,
            'make': 'Toyota',
            'model': 'Camry',
            'year': 2022,
            'engine_type': 'gasoline',
            'price': 25000.0,
        },
        'missing_required': [],
    }


def build_nested_payload(base_id: int, category_id: int) -> dict[str, Any]:
    return {
        'fields': {
            'id': base_id + 20,
            'name': 'Smartphone',
            'price': 499.0,
            'category': {
                '$ref': {
                    'class_name': 'Category',
                    'object_id': category_id,
                    'resource': 'categories',
                }
            },
            'description': 'Mobile',
        },
        'missing_required': [],
    }


@pytest.mark.asyncio
@pytest.mark.parametrize('case_name', CASE_NAMES, ids=CASE_NAMES)
async def test_nl_query_creator_analyze(case_name: str, async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class: type[Any]

    if case_name == 'flat':
        from tests.fixtures.models.vehicle import Vehicle

        query = 'create vehicle Toyota Camry 2022 gasoline with price 25000'
        payload_response = json.dumps(build_flat_payload(base_id))
        scripted = [payload_response]
        model_class = Vehicle
    else:
        from tests.fixtures.models.catalog_item import CatalogItem

        category_id = await seed_category(base_id)
        query = 'create catalog item smartphone in Electronics'
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        payload_response = json.dumps(build_nested_payload(base_id, category_id))
        scripted = [resolution_response, filter_response, payload_response]
        model_class = CatalogItem

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

    creator = NLQueryCreator(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    analysis = await creator.analyze(query, preview_limit=5)

    assert analysis.total == 1
    assert analysis.intent == CreateIntent.CREATE
    assert analysis.target_count == 1
    assert analysis.summary is not None


@pytest.mark.asyncio
@pytest.mark.parametrize('case_name', CASE_NAMES, ids=CASE_NAMES)
async def test_nl_query_creator_invoke(case_name: str, async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class: type[Any]
    expected_id: int

    if case_name == 'flat':
        from tests.fixtures.models.vehicle import Vehicle

        query = 'create vehicle Toyota Camry 2022 gasoline with price 25000'
        payload_response = json.dumps(build_flat_payload(base_id))
        scripted = [payload_response]
        model_class = Vehicle
        expected_id = base_id + 1
    else:
        from tests.fixtures.models.catalog_item import CatalogItem

        category_id = await seed_category(base_id)
        query = 'create catalog item smartphone in Electronics'
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        payload_response = json.dumps(build_nested_payload(base_id, category_id))
        scripted = [resolution_response, filter_response, payload_response]
        model_class = CatalogItem
        expected_id = base_id + 20

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

    creator = NLQueryCreator(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    analysis = await creator.analyze(query, preview_limit=5)
    docs = await creator.invoke(analysis)

    assert len(docs) == 1

    saved = await model_class.objects.filter(id__eq=expected_id).aexecute()
    assert len(saved) == 1


@pytest.mark.asyncio
@pytest.mark.parametrize('case_name', CASE_NAMES, ids=CASE_NAMES)
async def test_nl_query_creator_invoke_one(case_name: str, async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class: type[Any]
    expected_id: int

    if case_name == 'flat':
        from tests.fixtures.models.vehicle import Vehicle

        query = 'create vehicle Toyota Camry 2022 gasoline with price 25000'
        payload_response = json.dumps(build_flat_payload(base_id))
        scripted = [payload_response]
        model_class = Vehicle
        expected_id = base_id + 1
    else:
        from tests.fixtures.models.catalog_item import CatalogItem

        category_id = await seed_category(base_id)
        query = 'create catalog item smartphone in Electronics'
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        payload_response = json.dumps(build_nested_payload(base_id, category_id))
        scripted = [resolution_response, filter_response, payload_response]
        model_class = CatalogItem
        expected_id = base_id + 20

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

    creator = NLQueryCreator(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    doc = await creator.invoke_one(query)
    parsed = json.loads(doc.page_content)

    assert parsed['id'] == str(expected_id)


@pytest.mark.asyncio
@pytest.mark.parametrize('case_name', CASE_NAMES, ids=CASE_NAMES)
async def test_nl_query_creator_execute(case_name: str, async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class: type[Any]
    expected_id: int

    if case_name == 'flat':
        from tests.fixtures.models.vehicle import Vehicle

        query = 'create vehicle Toyota Camry 2022 gasoline with price 25000'
        payload_response = json.dumps(build_flat_payload(base_id))
        scripted = [payload_response]
        model_class = Vehicle
        expected_id = base_id + 1
    else:
        from tests.fixtures.models.catalog_item import CatalogItem

        category_id = await seed_category(base_id)
        query = 'create catalog item smartphone in Electronics'
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        payload_response = json.dumps(build_nested_payload(base_id, category_id))
        scripted = [resolution_response, filter_response, payload_response]
        model_class = CatalogItem
        expected_id = base_id + 20

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

    creator = NLQueryCreator(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    analysis = await creator.executor.analyze(query, preview_limit=5)
    results: list[Any] = await creator.executor.execute(analysis)

    assert len(results) == 1

    saved = await model_class.objects.filter(id__eq=expected_id).aexecute()
    assert len(saved) == 1


@pytest.mark.asyncio
@pytest.mark.parametrize('case_name', CASE_NAMES, ids=CASE_NAMES)
async def test_nl_query_creator_execute_one(case_name: str, async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class: type[Any]
    expected_id: int

    if case_name == 'flat':
        from tests.fixtures.models.vehicle import Vehicle

        query = 'create vehicle Toyota Camry 2022 gasoline with price 25000'
        payload_response = json.dumps(build_flat_payload(base_id))
        scripted = [payload_response]
        model_class = Vehicle
        expected_id = base_id + 1
    else:
        from tests.fixtures.models.catalog_item import CatalogItem

        category_id = await seed_category(base_id)
        query = 'create catalog item smartphone in Electronics'
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        payload_response = json.dumps(build_nested_payload(base_id, category_id))
        scripted = [resolution_response, filter_response, payload_response]
        model_class = CatalogItem
        expected_id = base_id + 20

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

    creator = NLQueryCreator(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    result: Any = await creator.executor.execute_one(query)

    assert result.id == expected_id
