"""Tests for JSON repair logic in the extraction pipeline (#639, #643).

Covers:
  - Markdown fence stripping
  - Trailing comma removal
  - Truncated JSON repair (missing closing brackets)
  - Retry-with-split behavior (mock LLM that fails on large batches)
"""

from __future__ import annotations

import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from cortex.extraction.extractor import (
    ExtractionPipeline,
    _repair_json,
    _repair_truncated_json,
)


# ---------------------------------------------------------------------------
# _repair_json unit tests
# ---------------------------------------------------------------------------

class TestMarkdownFenceStripping:
    """Test stripping of markdown code fences from LLM output."""

    def test_json_fence(self):
        raw = '```json\n[{"subject": "A", "predicate": "is", "object_value": "B"}]\n```'
        result = _repair_json(raw)
        data = json.loads(result)
        assert isinstance(data, list)
        assert data[0]["subject"] == "A"

    def test_plain_fence(self):
        raw = '```\n[{"subject": "A", "predicate": "is", "object_value": "B"}]\n```'
        result = _repair_json(raw)
        data = json.loads(result)
        assert isinstance(data, list)

    def test_no_fence(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B"}]'
        result = _repair_json(raw)
        assert json.loads(result) == json.loads(raw)

    def test_fence_with_extra_whitespace(self):
        raw = '  ```json\n  {"claims": [{"subject": "X", "predicate": "Y", "object_value": "Z"}]}  \n```  '
        result = _repair_json(raw)
        data = json.loads(result)
        assert "claims" in data


class TestTrailingCommaRemoval:
    """Test removal of trailing commas before ] or }."""

    def test_trailing_comma_in_array(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B"},]'
        result = _repair_json(raw)
        data = json.loads(result)
        assert len(data) == 1

    def test_trailing_comma_in_object(self):
        raw = '{"claims": [{"subject": "A", "predicate": "is", "object_value": "B",}]}'
        result = _repair_json(raw)
        data = json.loads(result)
        assert len(data["claims"]) == 1

    def test_multiple_trailing_commas(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B",}, {"subject": "C", "predicate": "is", "object_value": "D",},]'
        result = _repair_json(raw)
        data = json.loads(result)
        assert len(data) == 2

    def test_no_trailing_comma(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B"}]'
        result = _repair_json(raw)
        assert json.loads(result) == json.loads(raw)


class TestTruncatedJsonRepair:
    """Test repair of truncated JSON (output cut off mid-object)."""

    def test_truncated_array_mid_object(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B"}, {"subject": "C", "predi'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert isinstance(data, list)
        assert len(data) == 1
        assert data[0]["subject"] == "A"

    def test_truncated_wrapped_claims(self):
        raw = '{"claims": [{"subject": "A", "predicate": "is", "object_value": "B"}, {"subject": "C"'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert "claims" in data
        assert len(data["claims"]) == 1

    def test_truncated_single_complete_object(self):
        raw = '[{"subject": "X", "predicate": "likes", "object_value": "Y"}, {"sub'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert len(data) == 1
        assert data[0]["subject"] == "X"

    def test_truncated_no_complete_objects(self):
        raw = '[{"subject": "A", "predica'
        result = _repair_json(raw, "test")
        # No complete objects → returns original (still won't parse)
        assert result == raw

    def test_truncated_with_nested_braces_in_strings(self):
        raw = '[{"subject": "func(){}", "predicate": "is", "object_value": "code"}, {"subject": "trunc'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert len(data) == 1
        assert data[0]["subject"] == "func(){}"

    def test_valid_json_passes_through(self):
        raw = '[{"subject": "A", "predicate": "is", "object_value": "B"}]'
        result = _repair_json(raw)
        assert json.loads(result) == json.loads(raw)

    def test_multiple_complete_objects_before_truncation(self):
        obj1 = '{"subject": "A", "predicate": "is", "object_value": "1"}'
        obj2 = '{"subject": "B", "predicate": "is", "object_value": "2"}'
        raw = f'[{obj1}, {obj2}, {{"subject": "C", "predi'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert len(data) == 2
        assert data[0]["subject"] == "A"
        assert data[1]["subject"] == "B"


class TestCombinedRepair:
    """Test combinations of fence + trailing comma + truncation."""

    def test_fence_plus_trailing_comma(self):
        raw = '```json\n[{"subject": "A", "predicate": "is", "object_value": "B"},]\n```'
        result = _repair_json(raw)
        data = json.loads(result)
        assert len(data) == 1

    def test_fence_plus_truncation(self):
        raw = '```json\n[{"subject": "A", "predicate": "is", "object_value": "B"}, {"sub'
        result = _repair_json(raw, "test")
        data = json.loads(result)
        assert len(data) == 1


# ---------------------------------------------------------------------------
# Retry-with-split integration tests
# ---------------------------------------------------------------------------

class TestRetryWithSplit:
    """Test retry-with-split behavior when extraction returns 0 claims."""

    @pytest.mark.asyncio
    async def test_split_retry_on_parse_failure(self):
        """When a large batch fails JSON parse, it should be split and retried."""
        call_count = 0

        async def mock_complete(system, user, **kwargs):
            nonlocal call_count
            call_count += 1
            turns = user.count("Turn")
            if turns > 2:
                # Simulate truncated output for large batches (no complete objects)
                return MagicMock(text='[{"subject": "A", "predicate": "is", "object_val')
            # Small batch succeeds
            return MagicMock(text=json.dumps([{
                "action": "ADD",
                "subject": "TestUser",
                "predicate": "likes",
                "object_value": "testing",
                "confidence": 0.9,
                "entities_mentioned": [],
            }]))

        mock_llm = AsyncMock()
        mock_llm.complete = mock_complete

        pipeline = ExtractionPipeline(
            llm=mock_llm,
            model="test-model",
            batch_size=4,
        )
        # Enable assistant turn extraction so all 4 messages are processed (#1277).
        pipeline.extract_assistant_turns = True

        messages = [
            {"role": "user", "content": "Turn 1: I like testing"},
            {"role": "assistant", "content": "Turn 1: That's great!"},
            {"role": "user", "content": "Turn 2: I like coding"},
            {"role": "assistant", "content": "Turn 2: Awesome!"},
        ]

        with patch("cortex.llm.router.resolve_routed_model", return_value="test-model"):
            result = await pipeline.extract(messages, session_id="test-split")

        # Should have retried with split and gotten claims from the halves
        assert len(result.claims) > 0
        # Original call (1) + 2 retry halves = 3 calls
        assert call_count == 3

    @pytest.mark.asyncio
    async def test_no_split_for_single_message_batch(self):
        """Single-message batches should NOT be split further."""
        async def mock_complete(system, user, **kwargs):
            return MagicMock(text='not json at all')

        mock_llm = AsyncMock()
        mock_llm.complete = mock_complete

        pipeline = ExtractionPipeline(
            llm=mock_llm,
            model="test-model",
            batch_size=1,
        )

        messages = [{"role": "user", "content": "Hello"}]

        with patch("cortex.llm.router.resolve_routed_model", return_value="test-model"):
            result = await pipeline.extract(messages, session_id="test-no-split")

        assert len(result.claims) == 0

    @pytest.mark.asyncio
    async def test_successful_batch_skips_retry(self):
        """Batches that parse successfully should not trigger retry."""
        call_count = 0

        async def mock_complete(system, user, **kwargs):
            nonlocal call_count
            call_count += 1
            return MagicMock(text=json.dumps([{
                "action": "ADD",
                "subject": "User",
                "predicate": "likes",
                "object_value": "music",
                "confidence": 0.8,
                "entities_mentioned": [],
            }]))

        mock_llm = AsyncMock()
        mock_llm.complete = mock_complete

        pipeline = ExtractionPipeline(
            llm=mock_llm,
            model="test-model",
            batch_size=5,
        )

        messages = [
            {"role": "user", "content": "I like music"},
            {"role": "assistant", "content": "Nice!"},
        ]

        with patch("cortex.llm.router.resolve_routed_model", return_value="test-model"):
            result = await pipeline.extract(messages, session_id="test-ok")

        assert len(result.claims) == 1
        assert call_count == 1
