from backend.relation_extractor import extract_patterns, extract_relations, _find_mentions


CHUNK = (
    "Foxconn is a major supplier to Apple Inc. for iPhone assembly. "
    "In contrast, Apple competes with Samsung in the smartphone market. "
    "Beats Electronics, a subsidiary of Apple, also contributed to revenue."
)
ORGS = ["Foxconn", "Apple Inc.", "Samsung", "Beats Electronics", "Apple"]


def _rels(pairs):
    return {(a.lower(), b.lower()): r for a, b, r in pairs}


def test_supplier_pattern_found():
    rels = _rels(extract_patterns(CHUNK, ORGS))
    assert rels.get(("foxconn", "apple inc.")) == "SUPPLIER_TO"


def test_competitor_pattern_found():
    rels = _rels(extract_patterns(CHUNK, ORGS))
    assert rels.get(("apple", "samsung")) == "COMPETITOR_OF"


def test_subsidiary_pattern_found():
    rels = _rels(extract_patterns(CHUNK, ORGS))
    assert rels.get(("beats electronics", "apple")) == "SUBSIDIARY_OF"


def test_cross_sentence_pairs_not_connected():
    # Foxconn (sentence 1) and Samsung (sentence 2) co-occur closely but in
    # different sentences — must NOT be related
    rels = _rels(extract_patterns(CHUNK, ORGS))
    assert ("foxconn", "samsung") not in rels
    assert ("samsung", "foxconn") not in rels


def test_overlapping_mentions_collapse_to_longest():
    spans = _find_mentions("Apple Inc. was mentioned.", ["Apple", "Apple Inc."])
    assert len(spans) == 1
    assert spans[0][2] == "Apple Inc."


def test_single_org_returns_empty():
    assert extract_relations(CHUNK, ["Apple"], use_llm_fallback=False) == []


def test_pattern_only_mode_needs_no_llm():
    results = extract_relations(CHUNK, ORGS, use_llm_fallback=False)
    assert all(r["source"] == "pattern" for r in results)
    assert len(results) >= 3