"""
tests/test_split_extraction.py — Tests for issue #321: User vs Agent Memory Extraction.

Covers:
  1. extraction_mode="unified" (default): behavior unchanged, source_role=None
  2. extraction_mode="split": user messages use USER_EXTRACTION_SYSTEM_PROMPT
  3. extraction_mode="split": agent messages use AGENT_EXTRACTION_SYSTEM_PROMPT
  4. extraction_mode="split": claims tagged with correct source_role
  5. extraction_mode="split": extract_assistant_turns=False skips agent pass
  6. source_role preserved through to_candidate() / CandidateClaim
  7. CortexConfig.extraction_mode field defaults to "unified"
  8. dict_to_claim parses source_role from LLM output
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import asdict
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch, call

import pytest

from cortex.extraction.pipeline import ExtractionPipeline
from cortex.extraction.models import ClaimExtraction, ExtractionResult
from cortex.extraction.utils import dict_to_claim
from cortex.extraction.prompts import (
    USER_EXTRACTION_SYSTEM_PROMPT,
    AGENT_EXTRACTION_SYSTEM_PROMPT,
    EXTRACTION_SYSTEM_PROMPT,
)
from cortex.types.claims import CandidateClaim


# ---------------------------------------------------------------------------
# Fixtures & helpers
# ---------------------------------------------------------------------------

def _make_llm(response_json: str) -> MagicMock:
    """Create a mock LLM that always returns the given JSON string."""
    mock = MagicMock()
    resp = MagicMock()
    resp.text = response_json
    mock.complete = AsyncMock(return_value=resp)
    return mock


def _claim_json(subject="user", predicate="prefers", obj="Python", source_role=None):
    d = {
        "action": "ADD",
        "subject": subject,
        "predicate": predicate,
        "object_value": obj,
        "confidence": 0.9,
        "temporal": "current",
        "temporal_marker": None,
        "entities_mentioned": [],
        "sensitivity": "normal",
        "category": "preferences",
        "memory_class": "semantic",
        "salience": "high",
        "event_date": None,
        "sequence_order": None,
        "cause": None,
        "effect": None,
        "state_before": None,
        "state_after": None,
    }
    if source_role is not None:
        d["source_role"] = source_role
    return d


MESSAGES_MIXED = [
    {"role": "user", "content": "I love Python for backend work."},
    {"role": "assistant", "content": "I tend to explain things step by step."},
]

MESSAGES_USER_ONLY = [
    {"role": "user", "content": "I love Python for backend work."},
]

MESSAGES_AGENT_ONLY = [
    {"role": "assistant", "content": "I tend to explain things step by step."},
]


# ---------------------------------------------------------------------------
# 1. Unified mode: behavior unchanged, source_role=None on claims
# ---------------------------------------------------------------------------

class TestUnifiedModeUnchanged:
    """extraction_mode='unified' must produce identical behavior to current default."""

    def _make_pipeline(self, mode="unified"):
        llm = _make_llm(json.dumps({"claims": [_claim_json()]}))
        p = ExtractionPipeline(llm=llm, extraction_mode=mode)
        p.extract_assistant_turns = True  # allow all roles
        return p, llm

    def test_unified_mode_is_default(self):
        llm = _make_llm(json.dumps({"claims": [_claim_json()]}))
        p = ExtractionPipeline(llm=llm)
        assert p.extraction_mode == "unified"

    def test_unified_mode_claims_have_no_source_role(self):
        p, _ = self._make_pipeline("unified")
        result = asyncio.run(
            p.extract(MESSAGES_USER_ONLY, session_id="test-unified")
        )
        assert isinstance(result, ExtractionResult)
        for claim in result.claims:
            # unified mode: source_role is None (not set by the pipeline)
            assert claim.source_role is None

    def test_unified_mode_calls_system_prompt_not_split_prompts(self):
        p, llm = self._make_pipeline("unified")
        asyncio.run(
            p.extract(MESSAGES_USER_ONLY, session_id="test-unified-prompt")
        )
        # The system prompt used must NOT be the split-mode prompts
        called_system = llm.complete.call_args[1].get("system") or llm.complete.call_args[0][0] if llm.complete.call_args else None
        # check via the pipeline's _system_prompt attribute
        assert p._system_prompt not in (USER_EXTRACTION_SYSTEM_PROMPT, AGENT_EXTRACTION_SYSTEM_PROMPT)


# ---------------------------------------------------------------------------
# 2 & 3. Split mode: correct prompts used per role
# ---------------------------------------------------------------------------

class TestSplitModePromptRouting:
    """extraction_mode='split' must use role-specific prompts."""

    def test_user_messages_use_user_prompt(self):
        """When split mode processes user messages, USER_EXTRACTION_SYSTEM_PROMPT is used."""
        captured_calls = []

        async def mock_complete(**kwargs):
            captured_calls.append(kwargs.get("system", ""))
            resp = MagicMock()
            resp.text = json.dumps({"claims": [_claim_json()]})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = False  # skip agent pass

        asyncio.run(
            p.extract(MESSAGES_USER_ONLY, session_id="test-split-user")
        )

        assert len(captured_calls) >= 1
        assert captured_calls[0] == USER_EXTRACTION_SYSTEM_PROMPT

    def test_agent_messages_use_agent_prompt(self):
        """When split mode processes assistant messages, AGENT_EXTRACTION_SYSTEM_PROMPT is used."""
        captured_calls = []

        async def mock_complete(**kwargs):
            captured_calls.append(kwargs.get("system", ""))
            resp = MagicMock()
            resp.text = json.dumps({"claims": [_claim_json(subject="Eva", predicate="has_trait", obj="explains step by step")]})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True

        # Only agent messages
        asyncio.run(
            p.extract(MESSAGES_AGENT_ONLY, session_id="test-split-agent")
        )

        # user pass returns empty (no user messages), agent pass should be called
        assert any(c == AGENT_EXTRACTION_SYSTEM_PROMPT for c in captured_calls)


# ---------------------------------------------------------------------------
# 4. Split mode: claims tagged with correct source_role
# ---------------------------------------------------------------------------

class TestSplitModeSourceRoleTagging:
    """Claims extracted in split mode must have source_role set."""

    def test_user_claims_tagged_user(self):
        llm = _make_llm(json.dumps({"claims": [_claim_json()]}))
        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = False

        result = asyncio.run(
            p.extract(MESSAGES_USER_ONLY, session_id="test-user-tag")
        )

        assert len(result.claims) > 0
        for claim in result.claims:
            assert claim.source_role == "user", f"Expected source_role='user', got {claim.source_role!r}"

    def test_agent_claims_tagged_agent(self):
        agent_claim = _claim_json(subject="Eva", predicate="has_trait", obj="explains step by step")
        llm = _make_llm(json.dumps({"claims": [agent_claim]}))
        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True

        result = asyncio.run(
            p.extract(MESSAGES_AGENT_ONLY, session_id="test-agent-tag")
        )

        assert len(result.claims) > 0
        for claim in result.claims:
            assert claim.source_role == "agent", f"Expected source_role='agent', got {claim.source_role!r}"

    def test_mixed_messages_both_roles_tagged(self):
        """Mixed conversation: user claims get 'user', agent claims get 'agent'."""
        call_count = 0

        async def mock_complete(**kwargs):
            nonlocal call_count
            call_count += 1
            resp = MagicMock()
            system = kwargs.get("system", "")
            if system == USER_EXTRACTION_SYSTEM_PROMPT:
                resp.text = json.dumps({"claims": [_claim_json(subject="user", predicate="prefers", obj="Python")]})
            elif system == AGENT_EXTRACTION_SYSTEM_PROMPT:
                resp.text = json.dumps({"claims": [_claim_json(subject="Eva", predicate="has_trait", obj="explains step by step")]})
            else:
                resp.text = json.dumps({"claims": []})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True

        result = asyncio.run(
            p.extract(MESSAGES_MIXED, session_id="test-mixed-tag")
        )

        # Should have made 2 LLM calls (user pass + agent pass)
        assert call_count == 2

        source_roles = {c.source_role for c in result.claims}
        assert "user" in source_roles, f"Expected 'user' in source roles, got {source_roles}"
        assert "agent" in source_roles, f"Expected 'agent' in source roles, got {source_roles}"


# ---------------------------------------------------------------------------
# 5. Split mode with extract_assistant_turns=False: skip agent pass
# ---------------------------------------------------------------------------

class TestSplitModeAgentTurnsGate:

    def test_agent_pass_skipped_when_extract_assistant_turns_false(self):
        call_count = 0

        async def mock_complete(**kwargs):
            nonlocal call_count
            call_count += 1
            resp = MagicMock()
            resp.text = json.dumps({"claims": [_claim_json()]})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = False  # gate: skip agent pass

        asyncio.run(
            p.extract(MESSAGES_MIXED, session_id="test-no-agent")
        )

        # Only 1 call (user pass); agent pass skipped
        assert call_count == 1

    def test_agent_pass_runs_when_extract_assistant_turns_true(self):
        call_count = 0

        async def mock_complete(**kwargs):
            nonlocal call_count
            call_count += 1
            resp = MagicMock()
            resp.text = json.dumps({"claims": []})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True  # enable agent pass

        asyncio.run(
            p.extract(MESSAGES_MIXED, session_id="test-with-agent")
        )

        # 2 calls: user pass + agent pass
        assert call_count == 2


# ---------------------------------------------------------------------------
# 6. source_role preserved through to_candidate()
# ---------------------------------------------------------------------------

class TestSourceRoleToCandidatePropagation:

    def test_source_role_preserved_in_to_candidate(self):
        claim = ClaimExtraction(
            action="ADD",
            subject="user",
            predicate="prefers",
            object_value="Python",
            confidence=0.9,
            source_role="user",
        )
        candidate = claim.to_candidate()
        assert isinstance(candidate, CandidateClaim)
        assert candidate.source_role == "user"

    def test_agent_source_role_preserved_in_to_candidate(self):
        claim = ClaimExtraction(
            action="ADD",
            subject="Eva",
            predicate="has_trait",
            object_value="explains step by step",
            confidence=0.8,
            source_role="agent",
        )
        candidate = claim.to_candidate()
        assert candidate.source_role == "agent"

    def test_none_source_role_preserved_in_to_candidate(self):
        """Unified mode: source_role=None must not be mutated."""
        claim = ClaimExtraction(
            action="ADD",
            subject="user",
            predicate="prefers",
            object_value="Python",
            confidence=0.9,
            source_role=None,
        )
        candidate = claim.to_candidate()
        assert candidate.source_role is None


# ---------------------------------------------------------------------------
# 7. CortexConfig.extraction_mode field
# ---------------------------------------------------------------------------

class TestCortexConfigExtractionMode:

    def test_extraction_mode_defaults_to_unified(self):
        from cortex.config.model import CortexConfig
        config = CortexConfig()
        assert config.extraction_mode == "unified"

    def test_extraction_mode_can_be_set_to_split(self):
        from cortex.config.model import CortexConfig
        config = CortexConfig(extraction_mode="split")
        assert config.extraction_mode == "split"


# ---------------------------------------------------------------------------
# 8. dict_to_claim parses source_role from LLM output
# ---------------------------------------------------------------------------

class TestDictToClaimSourceRole:

    def test_dict_to_claim_parses_user_source_role(self):
        d = _claim_json(source_role="user")
        claim = dict_to_claim(d)
        assert claim.source_role == "user"

    def test_dict_to_claim_parses_agent_source_role(self):
        d = _claim_json(source_role="agent")
        claim = dict_to_claim(d)
        assert claim.source_role == "agent"

    def test_dict_to_claim_parses_system_source_role(self):
        d = _claim_json(source_role="system")
        claim = dict_to_claim(d)
        assert claim.source_role == "system"

    def test_dict_to_claim_ignores_invalid_source_role(self):
        """Invalid source_role values are ignored (set to None)."""
        d = _claim_json(source_role="unknown_role")
        claim = dict_to_claim(d)
        assert claim.source_role is None

    def test_dict_to_claim_source_role_none_when_missing(self):
        """When source_role is absent in LLM output, it defaults to None."""
        d = _claim_json()  # no source_role key
        claim = dict_to_claim(d)
        assert claim.source_role is None


# ---------------------------------------------------------------------------
# 9. Empty inputs handled gracefully in split mode
# ---------------------------------------------------------------------------

class TestSplitModeEdgeCases:

    def test_empty_messages_returns_empty_result(self):
        llm = _make_llm(json.dumps({"claims": []}))
        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True

        result = asyncio.run(
            p.extract([], session_id="test-empty")
        )
        assert result.claims == []
        assert result.batches_processed == 0

    def test_no_agent_messages_skips_agent_pass(self):
        call_count = 0

        async def mock_complete(**kwargs):
            nonlocal call_count
            call_count += 1
            resp = MagicMock()
            resp.text = json.dumps({"claims": [_claim_json()]})
            return resp

        llm = MagicMock()
        llm.complete = AsyncMock(side_effect=mock_complete)

        p = ExtractionPipeline(llm=llm, extraction_mode="split")
        p.extract_assistant_turns = True  # enabled but no agent messages in input

        result = asyncio.run(
            p.extract(MESSAGES_USER_ONLY, session_id="test-no-agent-msgs")
        )

        # Only 1 LLM call (user pass); agent pass has nothing to process
        assert call_count == 1

    def test_unified_mode_string_input_still_works(self):
        """Auto-wrap string input (#602) works in unified mode."""
        llm = _make_llm(json.dumps({"claims": [_claim_json()]}))
        p = ExtractionPipeline(llm=llm, extraction_mode="unified")

        result = asyncio.run(
            p.extract("I love Python", session_id="test-str-input")
        )
        assert len(result.claims) > 0
