from app.services.memory_extraction import MemoryExtractionResult
from app.services.memory_repository import SqlAlchemyMemoryRepository
from app.services.memory_service import MemoryExtractionService
from app.db import Base, create_session_factory
from sqlalchemy import create_engine


class FakeLLM:
    def __init__(self, payload):
        self.payload = payload
        self.messages = None

    def chat_json(self, messages, temperature=0.3, max_tokens=4096):
        self.messages = messages
        assert temperature <= 0.3
        assert max_tokens >= 4096
        return self.payload


def test_memory_extraction_service_calls_json_llm_and_persists_validated_result():
    payload = {
        "entities": [
            {
                "mention": "Alice",
                "canonical_name": "Alice",
                "labels": ["Person"],
                "aliases": [],
                "attributes": {},
                "summary": "Founder.",
                "confidence": 0.9,
            },
            {
                "mention": "Bob",
                "canonical_name": "Bob",
                "labels": ["Person"],
                "aliases": [],
                "attributes": {},
                "summary": "Partner.",
                "confidence": 0.8,
            },
        ],
        "edges": [
            {
                "source_entity_ref": "Alice",
                "target_entity_ref": "Bob",
                "relation": "KNOWS",
                "fact": "Alice knows Bob.",
                "attributes": {},
                "valid_at": None,
                "invalid_at": None,
                "expired_at": None,
                "confidence": 0.8,
                "evidence": ["episode-1:0-15"],
            }
        ],
        "episode_summary": "Relationship.",
        "unresolved_mentions": [],
    }
    client = FakeLLM(payload)
    service = MemoryExtractionService(client)

    result = service.extract(
        language="en",
        ontology={"entity_types": ["Person"], "edge_types": ["KNOWS"]},
        episode_text="Alice knows Bob.",
    )
    assert isinstance(result, MemoryExtractionResult)
    assert client.messages[0]["content"].startswith("IMPORTANT:")

    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    try:
        with create_session_factory(engine)() as session:
            repository = SqlAlchemyMemoryRepository(session, organization_id="org-a", graph_id="graph-a")
            repository.create_graph(project_id="project-a")
            ingest = service.persist(
                repository,
                result,
                source_type="document",
                source_ref="episode-1",
                episode_text="Alice knows Bob.",
            )
            session.commit()
            assert ingest.entity_count == 2
            assert ingest.edge_count == 1
            assert repository.search("Alice").total_count == 1
    finally:
        engine.dispose()