import json
from pathlib import Path

from sqlalchemy import create_engine

from app.db import Base, create_session_factory
from app.models.memory import MemoryEdge, MemoryGraph, MemoryNode
from app.services.memory_entity_reader import LocalEntityReader


FIXTURE = json.loads(
    Path(__file__).with_name("fixtures").joinpath("memory_parity", "entity_reader_fixture.json").read_text()
)


def test_local_entity_reader_matches_entity_filter_fixture():
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    try:
        with create_session_factory(engine)() as session:
            session.add(
                MemoryGraph(
                    id=FIXTURE["graph_id"],
                    organization_id="org-a",
                    project_id="project-a",
                    ontology={"entity_types": FIXTURE["filter"]["defined_entity_types"]},
                )
            )
            session.add_all(
                [
                    MemoryNode(graph_id=FIXTURE["graph_id"], **node)
                    for node in FIXTURE["nodes"]
                ]
            )
            session.add_all(
                [
                    MemoryEdge(graph_id=FIXTURE["graph_id"], **edge)
                    for edge in FIXTURE["edges"]
                ]
            )
            session.commit()

            reader = LocalEntityReader(session, organization_id="org-a", graph_id=FIXTURE["graph_id"])
            result = reader.filter_defined_entities(
                defined_entity_types=FIXTURE["filter"]["defined_entity_types"],
                enrich_with_edges=True,
            )
            assert result.total_count == FIXTURE["filter"]["total_count"]
            assert result.filtered_count == FIXTURE["filter"]["filtered_count"]
            assert sorted(result.entity_types) == FIXTURE["filter"]["entity_types"]
            alice = next(entity for entity in result.entities if entity.name == "Alice")
            assert alice.related_edges[0]["fact"] == "Alice works for Acme."

            detail = reader.get_entity_with_context("node-alice")
            assert detail is not None
            assert detail.get_entity_type() == "Person"
            assert detail.related_nodes[0]["name"] == "Acme"
    finally:
        engine.dispose()