"""Tests for entity extraction"""

import pytest
from sentiment_engine.nlp.entity_extraction import AssetMapper, EntityExtractor


class TestAssetMapper:
    """Test asset mapping"""

    def test_map_known_ticker(self):
        mapper = AssetMapper()
        asset_id, confidence = mapper.map_ticker("BTC")
        assert asset_id == "BTC"
        assert confidence >= 0.9

    def test_map_alias(self):
        mapper = AssetMapper()
        asset_id, confidence = mapper.map_ticker("VITALIK")
        assert asset_id == "ETH"
        assert confidence >= 0.7

    def test_map_unknown_ticker(self):
        mapper = AssetMapper()
        asset_id, confidence = mapper.map_ticker("UNKNOWNTICKER")
        assert asset_id == "UNKNOWNTICKER"
        assert confidence == 0.5

    def test_map_contract(self):
        mapper = AssetMapper()
        asset_id, confidence, chain = mapper.map_contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2")
        assert asset_id == "ETH"
        assert confidence >= 0.9
        assert chain == "ethereum"

    def test_resolve_aliases(self):
        mapper = AssetMapper()
        results = mapper.resolve_alias("Vitalik Buterin says ETH will moon")
        assert any(r[1] == "ETH" for r in results)


class TestEntityExtractor:
    """Test entity extraction"""

    @pytest.fixture
    def extractor(self):
        return EntityExtractor()

    def test_extract_tickers(self, extractor):
        text = "BTC and ETH are pumping hard"
        mentions = extractor.extract_tickers(text)
        assert len(mentions) == 2
        asset_ids = [m.asset_id for m in mentions]
        assert "BTC" in asset_ids
        assert "ETH" in asset_ids

    def test_extract_contracts(self, extractor):
        text = "Send to 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
        mentions = extractor.extract_contracts(text)
        assert len(mentions) == 1
        assert mentions[0].asset_id == "ETH"

    def test_extract_aliases(self, extractor):
        text = "Vitalik says ETH to the moon"
        mentions = extractor.extract_aliases(text)
        assert len(mentions) >= 1
        assert mentions[0].asset_id == "ETH"

    def test_deduplication(self, extractor):
        text = "BTC BTC BTC"
        mentions = extractor.extract_tickers(text)
        assert len(mentions) == 1
        assert mentions[0].asset_id == "BTC"

    @pytest.mark.asyncio
    async def test_extract_all(self, extractor):
        text = "BTC surges. Vitalik buys ETH. Send to 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
        entities = await extractor.extract_all(text)
        asset_ids = [e.asset_id for e in entities]
        assert "BTC" in asset_ids
        assert "ETH" in asset_ids
        # Contract should also map to ETH
        assert asset_ids.count("ETH") >= 1