"""Tests for entity and relation extraction."""

from memotrail.extractors.entities import extract_entities, extract_relations


class TestExtractEntities:
    def test_empty_messages(self):
        assert extract_entities([]) == []

    def test_extract_library_from_import(self):
        messages = [
            {"role": "assistant", "content": "import pandas\nfrom flask import Flask"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "pandas" in names
        assert "flask" in names

    def test_extract_library_from_pip_install(self):
        messages = [
            {"role": "assistant", "content": "pip install requests"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "requests" in names

    def test_extract_framework(self):
        messages = [
            {"role": "assistant", "content": "We chose React for the frontend"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "react" in names

    def test_extract_language(self):
        messages = [
            {"role": "user", "content": "The project is written in Python"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "python" in names

    def test_extract_service(self):
        messages = [
            {"role": "assistant", "content": "We're using PostgreSQL for the database"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "postgresql" in names

    def test_extract_file(self):
        messages = [
            {"role": "assistant", "content": "I updated the `server.py` file"},
        ]
        entities = extract_entities(messages)
        names = [e.name for e in entities]
        assert "server.py" in names

    def test_deduplication(self):
        messages = [
            {"role": "assistant", "content": "import React from 'react'\nimport React from 'react'"},
        ]
        entities = extract_entities(messages)
        react_count = sum(1 for e in entities if e.name.lower() == "react")
        assert react_count <= 1

    def test_skips_generic_words(self):
        messages = [
            {"role": "user", "content": "the file from this path using data"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "the" not in names
        assert "this" not in names

    def test_entity_type_classification(self):
        messages = [
            {"role": "assistant", "content": "Using Redis for caching"},
        ]
        entities = extract_entities(messages)
        redis_entities = [e for e in entities if e.name.lower() == "redis"]
        if redis_entities:
            assert redis_entities[0].entity_type == "service"

    def test_npm_install(self):
        messages = [
            {"role": "assistant", "content": "npm install express"},
        ]
        entities = extract_entities(messages)
        names = [e.name.lower() for e in entities]
        assert "express" in names

    def test_cap_at_50(self):
        content = " ".join(f"import module{i}" for i in range(100))
        messages = [{"role": "assistant", "content": content}]
        entities = extract_entities(messages)
        assert len(entities) <= 50

    def test_multiple_entity_types(self):
        messages = [
            {"role": "assistant", "content": (
                "We're using Python with FastAPI framework. "
                "The database is PostgreSQL. We installed pytest for testing."
            )},
        ]
        entities = extract_entities(messages)
        types = {e.entity_type for e in entities}
        assert len(types) >= 2  # Should have multiple types


class TestExtractRelations:
    def test_empty_messages(self):
        assert extract_relations([]) == []

    def test_uses_relation(self):
        messages = [
            {"role": "assistant", "content": "FastAPI uses Starlette under the hood"},
        ]
        relations = extract_relations(messages)
        if relations:
            assert any(
                r.source.lower() == "fastapi" and r.relation_type == "uses"
                for r in relations
            )

    def test_depends_on_relation(self):
        messages = [
            {"role": "assistant", "content": "MemoTrail depends on ChromaDB for vector storage"},
        ]
        relations = extract_relations(messages)
        if relations:
            assert any(r.relation_type == "depends_on" for r in relations)

    def test_replaces_relation(self):
        messages = [
            {"role": "assistant", "content": "We replaced Memcached with Redis"},
        ]
        relations = extract_relations(messages)
        if relations:
            assert any(r.relation_type == "replaces" for r in relations)

    def test_relation_has_context(self):
        messages = [
            {"role": "assistant", "content": "The project uses React for building the UI components"},
        ]
        relations = extract_relations(messages)
        for r in relations:
            assert r.context is not None

    def test_deduplication(self):
        messages = [
            {"role": "assistant", "content": (
                "FastAPI uses Starlette. FastAPI uses Starlette internally."
            )},
        ]
        relations = extract_relations(messages)
        keys = [(r.source.lower(), r.relation_type, r.target.lower()) for r in relations]
        assert len(keys) == len(set(keys))

    def test_cap_at_30(self):
        content = " ".join(f"tool{i} uses lib{i}" for i in range(50))
        messages = [{"role": "assistant", "content": content}]
        relations = extract_relations(messages)
        assert len(relations) <= 30
