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_updater import NLQueryUpdater
from amsdal_ml.ml_retrievers.query_updater import UpdateIntent
from amsdal_ml.utils.reference_utils import build_reference_payload
from tests.agents_tests.test_fakes import FakeModel


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


def build_vehicle_payload(base_id: int, price: float) -> dict[str, Any]:
    return {
        'fields': {
            'id': base_id + 1,
            'make': 'Toyota',
            'model': 'Camry',
            'year': 2022,
            'engine_type': 'gasoline',
            'fuel_efficiency': 30.0,
            'features': ['airbags'],
            'specs': {'engine_size': '2.5'},
            'is_available': True,
            'price': price,
            'warranty_years': 5,
            'maintenance_schedule': {'5000': {'oil_change': 'yes'}},
            'safety_ratings': {'nhtsa': 5},
            'options': ['navigation'],
        },
        'missing_required': [],
    }


def build_catalog_item_payload(base_id: int, category_id: int, description: str) -> 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': description,
        },
        'missing_required': [],
    }


def build_customer_payload(base_id: int, email: str) -> dict[str, Any]:
    return {
        'fields': {
            'id': base_id + 30,
            'name': 'Bob',
            'contact': {
                'email': email,
                'phone': '+380123456789',
                'address': {
                    'city': 'Kyiv',
                    'street': 'Shevchenko',
                },
            },
        },
        'missing_required': [],
    }


def build_category_payload(description: str) -> dict[str, Any]:
    return {
        'fields': {
            'description': description,
        },
        'missing_required': [],
    }


async def seed_vehicle(base_id: int) -> tuple[type[Any], list[int]]:
    from tests.fixtures.models.vehicle import Vehicle

    vehicle_one = Vehicle(
        id=base_id + 1,
        make='Toyota',
        model='Camry',
        year=2022,
        engine_type='gasoline',
        fuel_efficiency=30.0,
        features=['airbags'],
        specs={'engine_size': '2.5'},
        is_available=True,
        price=25000.0,
        warranty_years=5,
        maintenance_schedule={'5000': {'oil_change': 'yes'}},
        safety_ratings={'nhtsa': 5},
        options=['navigation'],
    )
    vehicle_two = Vehicle(
        id=base_id + 2,
        make='Tesla',
        model='Model 3',
        year=2023,
        engine_type='electric',
        fuel_efficiency=None,
        features=['autopilot'],
        specs={'battery_range': '358'},
        is_available=False,
        price=45000.0,
        warranty_years=8,
        maintenance_schedule=None,
        safety_ratings={'nhtsa': 5},
        options=['sunroof'],
    )

    await vehicle_one.asave(force_insert=True)
    await vehicle_two.asave(force_insert=True)

    return Vehicle, [vehicle_one.id]


async def seed_catalog_item(base_id: int) -> tuple[type[Any], list[int]]:
    from amsdal_utils.models.data_models.reference import Reference

    from tests.fixtures.models.catalog_item import CatalogItem
    from tests.fixtures.models.category import Category

    category_one = Category(id=base_id + 10, name='Electronics', description='Tech')
    category_two = Category(id=base_id + 11, name='Books', description='Reading')

    await category_one.asave(force_insert=True)
    await category_two.asave(force_insert=True)

    category_one_ref = Reference(
        ref={
            'class_name': 'Category',
            'object_id': category_one.id,
            'resource': Category.__table_name__,
            'class_version': 'LATEST',
            'object_version': 'LATEST',
        }
    )
    category_two_ref = Reference(
        ref={
            'class_name': 'Category',
            'object_id': category_two.id,
            'resource': Category.__table_name__,
            'class_version': 'LATEST',
            'object_version': 'LATEST',
        }
    )

    item_one = CatalogItem(
        id=base_id + 20,
        name='Smartphone',
        price=499.0,
        category=category_one_ref,  # type: ignore[arg-type]
        description='Mobile',
    )
    item_two = CatalogItem(
        id=base_id + 21,
        name='Novel',
        price=19.0,
        category=category_two_ref,  # type: ignore[arg-type]
        description='Fiction',
    )

    await item_one.asave(force_insert=True)
    await item_two.asave(force_insert=True)

    return CatalogItem, [item_one.id]


async def seed_catalog_item_with_categories(base_id: int) -> tuple[type[Any], int, int]:
    from amsdal_utils.models.data_models.reference import Reference

    from tests.fixtures.models.catalog_item import CatalogItem
    from tests.fixtures.models.category import Category

    category_one = Category(id=base_id + 10, name='Electronics', description='Tech')
    category_two = Category(id=base_id + 11, name='Books', description='Reading')

    await category_one.asave(force_insert=True)
    await category_two.asave(force_insert=True)

    category_one_ref = Reference(
        ref={
            'class_name': 'Category',
            'object_id': category_one.id,
            'resource': Category.__table_name__,
            'class_version': 'LATEST',
            'object_version': 'LATEST',
        }
    )

    item_one = CatalogItem(
        id=base_id + 20,
        name='Smartphone',
        price=499.0,
        category=category_one_ref,  # type: ignore[arg-type]
        description='Mobile',
    )

    await item_one.asave(force_insert=True)

    return CatalogItem, item_one.id, category_two.id


async def seed_customer(base_id: int) -> tuple[type[Any], list[int]]:
    from tests.fixtures.models.customer import Address
    from tests.fixtures.models.customer import ContactInfo
    from tests.fixtures.models.customer import Customer

    customer = Customer(
        id=base_id + 30,
        name='Bob',
        contact=ContactInfo(
            email='old@gmail.com',
            phone='+380123456789',
            address=Address(city='Kyiv', street='Shevchenko'),
        ),
    )

    await customer.asave(force_insert=True)

    return Customer, [customer.id]


def get_nested_value(data: dict[str, Any], path: str) -> Any:
    value: Any = data
    for part in path.split('.'):
        if isinstance(value, dict):
            value = value.get(part)
        else:
            value = getattr(value, part, None)
    return value


CASES = [
    (
        'flat',
        seed_vehicle,
        'update gasoline vehicles: set price to 30000',
        json.dumps({'find_query': 'find gasoline vehicles', 'update_query': 'set price to 30000'}),
        json.dumps({'filters': [{'field': 'engine_type', 'lookup': 'eq', 'value': 'gasoline'}]}),
        'price',
        30000.0,
        False,
        UpdateIntent.UPDATE_DIRECT,
    ),
    (
        'nested',
        seed_catalog_item,
        'update catalog items in Electronics: set description to Updated',
        json.dumps({'find_query': 'find catalog items in Electronics', 'update_query': 'set description to Updated'}),
        json.dumps({'filters': [{'field': 'category__name', 'lookup': 'eq', 'value': 'Electronics'}]}),
        'description',
        'Updated',
        True,
        UpdateIntent.UPDATE_DIRECT,
    ),
    (
        'jsonb',
        seed_customer,
        'update customer Bob: set contact email to new@gmail.com',
        json.dumps(
            {
                'find_query': 'find customer with name Bob',
                'update_query': 'set contact email to new@gmail.com',
            }
        ),
        json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Bob'}]}),
        'contact.email',
        'new@gmail.com',
        False,
        UpdateIntent.UPDATE_DIRECT,
    ),
]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'case_name,seed,query,split_response,filter_response,field_name,field_value,needs_resolution,expected_intent',
    CASES,
    ids=[c[0] for c in CASES],
)
async def test_nl_query_updater_analyze(
    case_name: str,
    seed: Any,
    query: str,
    split_response: str,
    filter_response: str,
    field_name: str,
    field_value: Any,
    needs_resolution: bool,  # noqa: FBT001
    expected_intent: UpdateIntent,
    async_test_amsdal_manager: Any,
) -> None:
    _ = async_test_amsdal_manager
    _ = case_name
    _ = field_name
    _ = field_value
    _ = needs_resolution
    base_id = get_base_id()
    model_class, expected_ids = await seed(base_id)

    if case_name == 'flat':
        payload_response = json.dumps(build_vehicle_payload(base_id, float(field_value)))
        scripted = [split_response, filter_response, payload_response]
    elif case_name == 'nested':
        payload_response = json.dumps(build_catalog_item_payload(base_id, base_id + 10, str(field_value)))
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        search_filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
        scripted = [split_response, filter_response, resolution_response, search_filter_response, payload_response]
    else:
        payload_response = json.dumps(build_customer_payload(base_id, str(field_value)))
        scripted = [split_response, filter_response, payload_response]

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert analysis.total == len(expected_ids)
    assert analysis.intent == expected_intent
    assert analysis.target_updates
    assert analysis.target_count == len(expected_ids)
    assert analysis.summary is not None


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'case_name,seed,query,split_response,filter_response,field_name,field_value,needs_resolution,expected_intent',
    CASES,
    ids=[c[0] for c in CASES],
)
async def test_nl_query_updater_invoke(
    case_name: str,
    seed: Any,
    query: str,
    split_response: str,
    filter_response: str,
    field_name: str,
    field_value: Any,
    needs_resolution: bool,  # noqa: FBT001
    expected_intent: UpdateIntent,
    async_test_amsdal_manager: Any,
) -> None:
    _ = async_test_amsdal_manager
    _ = case_name
    _ = expected_intent
    base_id = get_base_id()
    model_class, expected_ids = await seed(base_id)

    if case_name == 'flat':
        payload_response = json.dumps(build_vehicle_payload(base_id, float(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    elif case_name == 'nested':
        payload_response = json.dumps(build_catalog_item_payload(base_id, base_id + 10, str(field_value)))
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        search_filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
    else:
        payload_response = json.dumps(build_customer_payload(base_id, str(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    scripted = [split_response, filter_response]
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert len(docs) == len(expected_ids)

    saved = await model_class.objects.filter(id__eq=expected_ids[0]).aexecute()
    assert get_nested_value(saved[0].model_dump(), field_name) == field_value


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'case_name,seed,query,split_response,filter_response,field_name,field_value,needs_resolution,expected_intent',
    CASES,
    ids=[c[0] for c in CASES],
)
async def test_nl_query_updater_invoke_one(
    case_name: str,
    seed: Any,
    query: str,
    split_response: str,
    filter_response: str,
    field_name: str,
    field_value: Any,
    needs_resolution: bool,  # noqa: FBT001
    expected_intent: UpdateIntent,
    async_test_amsdal_manager: Any,
) -> None:
    _ = async_test_amsdal_manager
    _ = case_name
    _ = expected_intent
    base_id = get_base_id()
    model_class, expected_ids = await seed(base_id)

    if case_name == 'flat':
        payload_response = json.dumps(build_vehicle_payload(base_id, float(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    elif case_name == 'nested':
        payload_response = json.dumps(build_catalog_item_payload(base_id, base_id + 10, str(field_value)))
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        search_filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
    else:
        payload_response = json.dumps(build_customer_payload(base_id, str(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    scripted = [split_response, filter_response]
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert parsed['id'] == str(expected_ids[0])

    saved = await model_class.objects.filter(id__eq=expected_ids[0]).aexecute()
    assert get_nested_value(saved[0].model_dump(), field_name) == field_value


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'case_name,seed,query,split_response,filter_response,field_name,field_value,needs_resolution,expected_intent',
    CASES,
    ids=[c[0] for c in CASES],
)
async def test_nl_query_updater_execute(
    case_name: str,
    seed: Any,
    query: str,
    split_response: str,
    filter_response: str,
    field_name: str,
    field_value: Any,
    needs_resolution: bool,  # noqa: FBT001
    expected_intent: UpdateIntent,
    async_test_amsdal_manager: Any,
) -> None:
    _ = async_test_amsdal_manager
    _ = case_name
    _ = expected_intent
    base_id = get_base_id()
    model_class, expected_ids = await seed(base_id)

    if case_name == 'flat':
        payload_response = json.dumps(build_vehicle_payload(base_id, float(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    elif case_name == 'nested':
        payload_response = json.dumps(build_catalog_item_payload(base_id, base_id + 10, str(field_value)))
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        search_filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
    else:
        payload_response = json.dumps(build_customer_payload(base_id, str(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    scripted = [split_response, filter_response]
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    analysis = await updater.executor.analyze(query, preview_limit=5)
    results = await updater.executor.execute(analysis)

    assert len(results) == len(expected_ids)

    saved = await model_class.objects.filter(id__eq=expected_ids[0]).aexecute()
    assert get_nested_value(saved[0].model_dump(), field_name) == field_value


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'case_name,seed,query,split_response,filter_response,field_name,field_value,needs_resolution,expected_intent',
    CASES,
    ids=[c[0] for c in CASES],
)
async def test_nl_query_updater_execute_one(
    case_name: str,
    seed: Any,
    query: str,
    split_response: str,
    filter_response: str,
    field_name: str,
    field_value: Any,
    needs_resolution: bool,  # noqa: FBT001
    expected_intent: UpdateIntent,
    async_test_amsdal_manager: Any,
) -> None:
    _ = async_test_amsdal_manager
    _ = case_name
    _ = expected_intent
    base_id = get_base_id()
    model_class, expected_ids = await seed(base_id)

    if case_name == 'flat':
        payload_response = json.dumps(build_vehicle_payload(base_id, float(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    elif case_name == 'nested':
        payload_response = json.dumps(build_catalog_item_payload(base_id, base_id + 10, str(field_value)))
        resolution_response = json.dumps(
            {
                'candidates': [
                    {
                        'field': 'category',
                        'entity_type': 'Category',
                        'search_keywords': ['Electronics'],
                    }
                ]
            }
        )
        search_filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Electronics'}]})
    else:
        payload_response = json.dumps(build_customer_payload(base_id, str(field_value)))
        resolution_response = json.dumps({'candidates': []})
        search_filter_response = json.dumps({'filters': []})
    scripted = [split_response, filter_response]
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)
    if needs_resolution:
        scripted.append(resolution_response)
        scripted.append(search_filter_response)
    scripted.append(payload_response)

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

    result = await updater.executor.execute_one(query)

    assert result.id == expected_ids[0]

    saved = await model_class.objects.filter(id__eq=expected_ids[0]).aexecute()
    assert get_nested_value(saved[0].model_dump(), field_name) == field_value


@pytest.mark.asyncio
async def test_nl_query_updater_jsonb_analysis_changes(async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class, _ = await seed_customer(base_id)

    query = 'update customer Bob: set contact email to new@gmail.com'
    split_response = json.dumps(
        {
            'find_query': 'find customer with name Bob',
            'update_query': 'set contact email to new@gmail.com',
        }
    )
    filter_response = json.dumps({'filters': [{'field': 'name', 'lookup': 'eq', 'value': 'Bob'}]})
    payload_response = json.dumps(build_customer_payload(base_id, 'new@gmail.com'))

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert analysis.target_updates
    changes = analysis.target_updates[0].changes
    assert any(change.field_path == 'contact.email' for change in changes)
    assert any(change.field_type.value == 'jsonb' for change in changes)


@pytest.mark.asyncio
async def test_nl_query_updater_related_intent_analysis(async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class, _, category_two_id = await seed_catalog_item_with_categories(base_id)

    query = 'update catalog items in Electronics: set category description to Updated'
    split_response = json.dumps(
        {
            'find_query': 'find catalog items in Electronics',
            'update_query': 'set category description to Updated',
        }
    )
    filter_response = json.dumps({'filters': [{'field': 'category__name', 'lookup': 'eq', 'value': 'Electronics'}]})
    resolution_response = json.dumps({'candidates': []})
    payload_response = json.dumps({'fields': {'description': 'Mobile'}, 'missing_required': []})
    normalize_response = 'set description to Updated'
    related_payload_response = json.dumps(build_category_payload('Updated'))

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert analysis.intent == UpdateIntent.UPDATE_RELATED
    assert analysis.related_updates
    assert analysis.related_updates[0].entity_update.entity_type == 'Category'
    assert analysis.target_updates[0].changes == []
    assert category_two_id


@pytest.mark.asyncio
async def test_nl_query_updater_reassign_fk_intent_analysis(async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class, _, category_two_id = await seed_catalog_item_with_categories(base_id)

    query = 'update catalog items in Electronics: set category to Books'
    split_response = json.dumps(
        {
            'find_query': 'find catalog items in Electronics',
            'update_query': 'set category to Books',
        }
    )
    filter_response = json.dumps({'filters': [{'field': 'category__name', 'lookup': 'eq', 'value': 'Electronics'}]})
    resolution_response = json.dumps({'candidates': []})
    payload_response = json.dumps(build_catalog_item_payload(base_id, category_two_id, 'Mobile'))

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert analysis.intent == UpdateIntent.REASSIGN_FK
    assert analysis.target_updates
    assert analysis.related_updates == []


@pytest.mark.asyncio
async def test_nl_query_updater_mixed_intent_analysis(async_test_amsdal_manager: Any) -> None:
    _ = async_test_amsdal_manager
    base_id = get_base_id()
    model_class, _, category_two_id = await seed_catalog_item_with_categories(base_id)

    query = 'update catalog items in Electronics: set description to Updated and category description to New'
    split_response = json.dumps(
        {
            'find_query': 'find catalog items in Electronics',
            'update_query': 'set description to Updated and category description to New',
        }
    )
    filter_response = json.dumps({'filters': [{'field': 'category__name', 'lookup': 'eq', 'value': 'Electronics'}]})
    resolution_response = json.dumps({'candidates': []})
    payload_response = json.dumps({'fields': {'description': 'Updated'}, 'missing_required': []})
    normalize_response = 'set description to New'
    related_payload_response = json.dumps(build_category_payload('New'))

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

    updater = NLQueryUpdater(
        llm=llm,
        queryset=model_class.objects.all(),
        llm_response_format=ResponseFormat.JSON_OBJECT,
    )

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

    assert analysis.intent == UpdateIntent.MIXED
    assert analysis.target_updates
    assert analysis.related_updates
    assert category_two_id


@pytest.mark.asyncio
async def test_nl_query_updater_build_reference_uses_amsdal_builder(
    async_test_amsdal_manager: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _ = async_test_amsdal_manager
    from amsdal_utils.models.data_models.reference import Reference

    from tests.fixtures.models.category import Category

    llm = FakeModel(async_mode=True, scripted=[])
    llm.setup()
    updater = NLQueryUpdater(llm=llm, queryset=Category.objects.all(), llm_response_format=ResponseFormat.JSON_OBJECT)

    captured: dict[str, object] = {}

    def _fake_build_reference(
        *,
        class_name: str,
        object_id: object,
        class_version: str = 'LATEST',
        object_version: str = 'LATEST',
        resource: str | None = None,
    ) -> Reference:
        captured['class_name'] = class_name
        captured['object_id'] = object_id
        captured['class_version'] = class_version
        captured['object_version'] = object_version
        captured['resource'] = resource
        return Reference(
            ref={
                'resource': 'test_connection',
                'class_name': class_name,
                'class_version': class_version,
                'object_id': object_id,
                'object_version': object_version,
            }
        )

    monkeypatch.setattr('amsdal_ml.ml_retrievers.query_updater.build_reference', _fake_build_reference)

    category = Category(id=get_base_id() + 500, name='Electronics', description='Tech')
    ref = updater.executor._build_reference(category)  # noqa: SLF001

    assert ref is not None
    assert ref.ref.resource == 'test_connection'
    assert captured == {
        'class_name': 'Category',
        'object_id': category.object_id,
        'class_version': 'LATEST',
        'object_version': 'LATEST',
        'resource': None,
    }


@pytest.mark.asyncio
async def test_nl_query_updater_build_reference_payload_uses_connection_name(
    async_test_amsdal_manager: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _ = async_test_amsdal_manager
    monkeypatch.setattr(
        'amsdal_ml.utils.reference_utils.resolve_reference_resource',
        lambda model_name: f'conn_for_{model_name.lower()}',
    )

    payload = build_reference_payload('Category', '123')

    assert payload == {
        '$ref': {
            'class_name': 'Category',
            'object_id': '123',
            'resource': 'conn_for_category',
        }
    }
