from sqlalchemy import create_engine

from app.db import Base, create_session_factory
from app.models.memory import MemoryGraph
from app.services.local_graph_builder import LocalGraphBuilderService


class FakeExtractionClient:
    def __init__(self):
        self.calls = []

    def chat_json(self, messages, temperature=0.3, max_tokens=4096):
        self.calls.append(messages)
        return {
            "entities": [
                {
                    "mention": "Alice",
                    "canonical_name": "Alice",
                    "labels": ["Person"],
                    "aliases": [],
                    "attributes": {"role": "founder"},
                    "summary": "Alice founded the local project.",
                    "confidence": 0.95,
                },
                {
                    "mention": "Orbit",
                    "canonical_name": "Orbit",
                    "labels": ["Project"],
                    "aliases": [],
                    "attributes": {},
                    "summary": "Orbit is the project discussed in the episode.",
                    "confidence": 0.9,
                },
            ],
            "edges": [
                {
                    "source_entity_ref": "Alice",
                    "target_entity_ref": "Orbit",
                    "relation": "FOUNDED",
                    "fact": "Alice founded Orbit.",
                    "attributes": {},
                    "valid_at": None,
                    "invalid_at": None,
                    "expired_at": None,
                    "confidence": 0.92,
                    "evidence": ["episode-0"],
                }
            ],
            "episode_summary": "Alice founded Orbit.",
            "unresolved_mentions": [],
        }


def make_session_factory():
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    return engine, create_session_factory(engine)


def test_local_graph_builder_persists_scoped_graph_and_extracted_memory():
    engine, session_factory = make_session_factory()
    client = FakeExtractionClient()
    builder = LocalGraphBuilderService(
        session_factory,
        organization_id="org-a",
        project_id="project-a",
        extraction_client=client,
    )

    graph_id = builder.create_graph(name="Orbit")
    builder.set_ontology(graph_id, {"entity_types": ["Person", "Project"], "relations": ["FOUNDED"]})
    episode_ids = builder.add_text_batches(
        graph_id,
        ["Alice founded Orbit."],
        batch_size=3,
    )

    assert len(episode_ids) == 1
    assert episode_ids[0].startswith("episode_")
    assert client.calls
    data = builder.get_graph_data(graph_id)
    assert data["graph_id"] == graph_id
    assert data["node_count"] == 2
    assert data["edge_count"] == 1

    with session_factory() as session:
        graph = session.get(MemoryGraph, graph_id)
        assert graph.organization_id == "org-a"
        assert graph.project_id == "project-a"
        assert graph.ontology["relations"] == ["FOUNDED"]


def test_local_graph_builder_reprocessing_is_idempotent_for_episode_and_edge():
    engine, session_factory = make_session_factory()
    builder = LocalGraphBuilderService(
        session_factory,
        organization_id="org-a",
        project_id="project-a",
        extraction_client=FakeExtractionClient(),
    )
    graph_id = builder.create_graph(name="Orbit")
    builder.set_ontology(graph_id, {"entity_types": ["Person", "Project"]})

    builder.add_text_batches(graph_id, ["Alice founded Orbit."])
    builder.add_text_batches(graph_id, ["Alice founded Orbit."])

    data = builder.get_graph_data(graph_id)
    assert data["node_count"] == 2
    assert data["edge_count"] == 1


def test_local_graph_builder_fails_closed_for_wrong_organization():
    engine, session_factory = make_session_factory()
    builder = LocalGraphBuilderService(
        session_factory,
        organization_id="org-a",
        project_id="project-a",
        extraction_client=FakeExtractionClient(),
    )
    graph_id = builder.create_graph(name="Orbit")

    other_builder = LocalGraphBuilderService(
        session_factory,
        organization_id="org-b",
        project_id="project-b",
        extraction_client=FakeExtractionClient(),
    )

    try:
        other_builder.get_graph_data(graph_id)
    except ValueError as exc:
        assert str(exc) == "memory_graph_not_found"
    else:
        raise AssertionError("wrong organization must not read graph")