from neurovault.cognitive.entity_extractor import EntityExtractor

extractor = EntityExtractor()


def _type_names(entities):
    return [e["type"].value if hasattr(e["type"], "value") else e["type"] for e in entities]


def test_extract_technology():
    entities = extractor.extract("We use Python and React for the frontend")
    names = [e["name"] for e in entities]
    assert "python" in names
    assert "react" in names


def test_extract_url():
    entities = extractor.extract("Check https://github.com/foo/bar")
    assert "URL" in _type_names(entities)


def test_extract_date():
    entities = extractor.extract("Meeting on March 15th")
    assert "DATE" in _type_names(entities)


def test_extract_file():
    entities = extractor.extract("Edit the config.toml file")
    names = [e["name"] for e in entities]
    assert "config.toml" in names


def test_deduplication():
    entities = extractor.extract("Python is great. I love Python.")
    python_count = sum(1 for e in entities if e["name"] == "python")
    assert python_count == 1


def test_single_capitalized_sentence_starters_are_not_person_entities():
    entities = extractor.extract("The system is stable. However we should monitor errors.")
    person_like = [e for e in entities if _type_names([e])[0] == "PERSON"]
    assert all(e["name"] not in {"The", "However"} for e in person_like)
