"""Tests for entity resolution wiring in the extraction pipeline.

Verifies that:
1. EntityResolver is called for each entity mention during wake_pipeline step 2
2. Resolved entity IDs are stamped as scope_entity_id on claims
3. scope_entity_id is carried through to CandidateClaim conversion (step 2d)
4. Reconciliation engine reuses scope_entity_id when available
5. Reconciliation engine handles entity_resolver=None gracefully
"""
import asyncio
from dataclasses import dataclass, field
from typing import List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.extraction.models import ClaimExtraction, EntityMention, ExtractionResult
from cortex.types import CandidateClaim, ResolvedEntity


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_claim(subject="Andrew", predicate="likes", obj="Python",
                entity_mentions=None):
    """Build a ClaimExtraction with optional entity mentions."""
    if entity_mentions is None:
        entity_mentions = [EntityMention(name=subject, entity_type="person")]
    return ClaimExtraction(
        action="ADD",
        subject=subject,
        predicate=predicate,
        object_value=obj,
        confidence=0.9,
        entities_mentioned=entity_mentions,
    )


def _make_extraction_result(claims=None):
    if claims is None:
        claims = [_make_claim()]
    return ExtractionResult(
        claims=claims,
        noise_filtered=0,
        raw_message_count=1,
    )


def _make_resolver(entity_id="ent-001"):
    """Mock EntityResolver that returns a ResolvedEntity."""
    resolver = MagicMock()

    async def _resolve(mention, owner_id="default", entity_type="unknown", **kw):
        return ResolvedEntity(
            entity_id=entity_id,
            canonical_name=mention,
            entity_type=entity_type,
            confidence=1.0,
            is_new=True,
            matched_alias=None,
        )

    resolver.resolve = AsyncMock(side_effect=_resolve)
    return resolver


def _noop_budget_check(n):
    return None

def _noop_budget_consume(n):
    pass


# ---------------------------------------------------------------------------
# Tests: wake_pipeline entity resolution (step 2)
# ---------------------------------------------------------------------------

class TestWakePipelineEntityResolution:
    """Tests for entity resolution in cortex.capture.wake_pipeline."""

    @pytest.mark.asyncio
    async def test_entity_resolver_called_for_each_mention(self):
        """EntityResolver.resolve() is called once per entity mention."""
        from cortex.capture.wake_pipeline import run_extraction

        resolver = _make_resolver()
        claims = [_make_claim(
            entity_mentions=[
                EntityMention(name="Andrew", entity_type="person"),
                EntityMention(name="Python", entity_type="technology"),
            ]
        )]
        extractor = MagicMock()
        extractor.extract = AsyncMock(return_value=_make_extraction_result(claims))

        reconciler = MagicMock()
        reconciler.reconcile = AsyncMock(return_value=MagicMock(
            results=[MagicMock(action=MagicMock(value="ADD"), resulting_claim_id="c1")],
        ))

        outcome = await run_extraction(
            extractor=extractor,
            storage=MagicMock(),
            llm=MagicMock(),
            entity_resolver=resolver,
            cornerstone_guardian=None,
            reconciler=reconciler,
            event_emitter=None,
            config=MagicMock(extraction_dedup_enabled=False),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "test"}],
            sid="test-session",
            db_session_pk="pk-1",
            effective_owner_id="eva",
            episode_event_id=None,
            budget_check=_noop_budget_check,
            budget_consume=_noop_budget_consume,
            errors=[],
            warnings=[],
            zero_claim_outcomes=0,
        )

        assert resolver.resolve.call_count == 2
        assert outcome.entities_resolved == 2

    @pytest.mark.asyncio
    async def test_scope_entity_id_stamped_on_matching_claim(self):
        """When a mention matches the claim subject, scope_entity_id is set."""
        from cortex.capture.wake_pipeline import run_extraction

        resolver = _make_resolver(entity_id="ent-andrew")
        claim = _make_claim(subject="Andrew")
        extractor = MagicMock()
        extractor.extract = AsyncMock(
            return_value=_make_extraction_result([claim])
        )

        reconciler = MagicMock()
        reconciler.reconcile = AsyncMock(return_value=MagicMock(
            results=[MagicMock(action=MagicMock(value="ADD"), resulting_claim_id="c1")],
        ))

        await run_extraction(
            extractor=extractor,
            storage=MagicMock(),
            llm=MagicMock(),
            entity_resolver=resolver,
            cornerstone_guardian=None,
            reconciler=reconciler,
            event_emitter=None,
            config=MagicMock(extraction_dedup_enabled=False),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "test"}],
            sid="test-session",
            db_session_pk="pk-1",
            effective_owner_id="eva",
            episode_event_id=None,
            budget_check=_noop_budget_check,
            budget_consume=_noop_budget_consume,
            errors=[],
            warnings=[],
            zero_claim_outcomes=0,
        )

        # The claim should have scope_entity_id stamped
        assert getattr(claim, "scope_entity_id", None) == "ent-andrew"

    @pytest.mark.asyncio
    async def test_scope_entity_id_carried_to_candidate_claim(self):
        """scope_entity_id is carried through ClaimExtraction→CandidateClaim."""
        from cortex.capture.wake_pipeline import run_extraction

        resolver = _make_resolver(entity_id="ent-andrew")
        claim = _make_claim(subject="Andrew")
        extractor = MagicMock()
        extractor.extract = AsyncMock(
            return_value=_make_extraction_result([claim])
        )

        # Capture what reconciler receives
        received_candidates = []

        async def capture_reconcile(candidates, **kw):
            received_candidates.extend(candidates)
            return MagicMock(
                results=[MagicMock(action=MagicMock(value="ADD"), resulting_claim_id="c1")],
            )

        reconciler = MagicMock()
        reconciler.reconcile = AsyncMock(side_effect=capture_reconcile)

        await run_extraction(
            extractor=extractor,
            storage=MagicMock(),
            llm=MagicMock(),
            entity_resolver=resolver,
            cornerstone_guardian=None,
            reconciler=reconciler,
            event_emitter=None,
            config=MagicMock(extraction_dedup_enabled=False),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "test"}],
            sid="test-session",
            db_session_pk="pk-1",
            effective_owner_id="eva",
            episode_event_id=None,
            budget_check=_noop_budget_check,
            budget_consume=_noop_budget_consume,
            errors=[],
            warnings=[],
            zero_claim_outcomes=0,
        )

        assert len(received_candidates) == 1
        cand = received_candidates[0]
        assert isinstance(cand, CandidateClaim)
        assert cand.scope_entity_id == "ent-andrew"

    @pytest.mark.asyncio
    async def test_entity_resolver_none_skips_resolution_gracefully(self):
        """When entity_resolver is None, resolution is skipped, not crashed."""
        from cortex.capture.wake_pipeline import run_extraction

        claim = _make_claim(subject="Andrew")
        extractor = MagicMock()
        extractor.extract = AsyncMock(
            return_value=_make_extraction_result([claim])
        )

        reconciler = MagicMock()
        reconciler.reconcile = AsyncMock(return_value=MagicMock(
            results=[MagicMock(action=MagicMock(value="ADD"), resulting_claim_id="c1")],
        ))

        errors: list = []
        outcome = await run_extraction(
            extractor=extractor,
            storage=MagicMock(),
            llm=MagicMock(),
            entity_resolver=None,  # <-- None
            cornerstone_guardian=None,
            reconciler=reconciler,
            event_emitter=None,
            config=MagicMock(extraction_dedup_enabled=False),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "test"}],
            sid="test-session",
            db_session_pk="pk-1",
            effective_owner_id="eva",
            episode_event_id=None,
            budget_check=_noop_budget_check,
            budget_consume=_noop_budget_consume,
            errors=errors,
            warnings=[],
            zero_claim_outcomes=0,
        )

        assert outcome.entities_resolved == 0
        assert any("EntityResolver not available" in e for e in errors)
        # Reconciliation should still proceed
        reconciler.reconcile.assert_called_once()


# ---------------------------------------------------------------------------
# Tests: ReconciliationEngine entity resolution (step 1)
# ---------------------------------------------------------------------------

class TestReconciliationEngineEntityResolution:
    """Tests for entity resolution in ReconciliationEngine._reconcile_one."""

    @pytest.mark.asyncio
    async def test_reuses_scope_entity_id_from_upstream(self):
        """When scope_entity_id is set, resolver.resolve() is NOT called."""
        from cortex.core.reconciliation import ReconciliationEngine

        resolver = _make_resolver(entity_id="ent-from-resolver")
        storage = MagicMock()
        storage.get_claims_by_subject = AsyncMock(return_value=[])
        storage.get_claims_by_entity = AsyncMock(return_value=[])
        storage.create_claim = AsyncMock(return_value=MagicMock(id="claim-1"))
        storage.log_changelog = AsyncMock()
        storage.link_claim_to_episode = AsyncMock()
        storage.increment_entity_mentions = AsyncMock()

        config = MagicMock()
        config.session_confidence_default = 0.5

        engine = ReconciliationEngine(
            storage=storage,
            llm=MagicMock(),
            entity_resolver=resolver,
            config=config,
        )

        candidate = CandidateClaim(
            subject="Andrew",
            predicate="likes",
            object_value="Python",
            confidence=0.9,
            scope_entity_id="ent-from-upstream",  # pre-resolved
        )

        report = await engine.reconcile(
            candidates=[candidate],
            session_id="s1",
            owner_id="eva",
        )

        # resolver.resolve should NOT be called — scope_entity_id was provided
        resolver.resolve.assert_not_called()
        # The claim should still be stored
        assert report.added >= 1

    @pytest.mark.asyncio
    async def test_entity_resolver_none_does_not_crash(self):
        """When entity_resolver is None, reconciliation continues without entities."""
        from cortex.core.reconciliation import ReconciliationEngine

        storage = MagicMock()
        storage.get_claims_by_subject = AsyncMock(return_value=[])
        storage.get_claims_by_entity = AsyncMock(return_value=[])
        storage.create_claim = AsyncMock(return_value=MagicMock(id="claim-1"))
        storage.log_changelog = AsyncMock()
        storage.link_claim_to_episode = AsyncMock()
        storage.increment_entity_mentions = AsyncMock()

        config = MagicMock()
        config.session_confidence_default = 0.5

        engine = ReconciliationEngine(
            storage=storage,
            llm=MagicMock(),
            entity_resolver=None,  # <-- None
            config=config,
        )

        candidate = CandidateClaim(
            subject="Andrew",
            predicate="likes",
            object_value="Python",
            confidence=0.9,
        )

        report = await engine.reconcile(
            candidates=[candidate],
            session_id="s1",
            owner_id="eva",
        )

        # Should not crash, claim should be stored without entity link
        assert report.added >= 1
