"""
Tests targeting specific uncovered lines in xbench_deepsearch.py.

Covers:
- Lines 133-138: Decrypt failure in load_data (exception in base64 decode/XOR)
- Lines 151-160: Sampling path in load_data when num_examples < total
- Lines 211-214: Encrypted prompt in _load_from_url
- Lines 221-226: Decrypt failure in _load_from_url
- Line 244: Sampling with seed in _load_from_url
"""

import base64
from unittest.mock import MagicMock, Mock, patch

MODULE = "local_deep_research.benchmarks.datasets.xbench_deepsearch"


def _make_dataset():
    from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
        XBenchDeepSearchDataset,
    )

    return XBenchDeepSearchDataset()


# ---------------------------------------------------------------------------
# Lines 133-138: Decrypt failure in load_data
# ---------------------------------------------------------------------------
class TestLoadDataDecryptFailure:
    def test_decrypt_failure_logs_warning_and_keeps_original(self):
        """When XOR decrypt produces invalid UTF-8, exception is caught
        and original (encrypted) value is used."""
        ds = _make_dataset()

        # Create a prompt that looks like base64 but when decrypted gives invalid UTF-8
        # Use bytes that are valid base64 but XOR to non-UTF-8 sequences
        bad_bytes = bytes([0xFF, 0xFE, 0xFD, 0xFC])  # Invalid UTF-8
        key = "k"
        key_bytes = key.encode("utf-8")
        # XOR to create "encrypted" form
        encrypted = bytes(
            [
                bad_bytes[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(bad_bytes))
            ]
        )
        encoded_prompt = base64.b64encode(encrypted).decode("utf-8")

        mock_item = {
            "id": "fail_1",
            "prompt": encoded_prompt,
            "answer": "plain answer",
            "canary": key,
            "reference_steps": "",
        }
        mock_dataset = [mock_item]
        mock_load_dataset = Mock(return_value=mock_dataset)

        with patch.dict(
            "sys.modules", {"datasets": Mock(load_dataset=mock_load_dataset)}
        ):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds.load_data()

        assert len(result) == 1
        # The warning should have been called
        mock_logger.warning.assert_called()


# ---------------------------------------------------------------------------
# Lines 151-160: Sampling in load_data when num_examples < total
# ---------------------------------------------------------------------------
class TestLoadDataSampling:
    def test_sampling_with_seed(self):
        """When num_examples < len(formatted_questions), sampling occurs."""
        ds = _make_dataset()

        items = [
            {
                "id": f"q{i}",
                "prompt": f"Question {i}",
                "answer": f"Answer {i}",
                "canary": "",
                "reference_steps": "",
            }
            for i in range(20)
        ]
        mock_dataset = items
        mock_load_dataset = Mock(return_value=mock_dataset)

        with patch.dict(
            "sys.modules", {"datasets": Mock(load_dataset=mock_load_dataset)}
        ):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds.load_data(num_examples=5, seed=42)

        assert len(result) == 5
        # Check the sampling log message was called
        sample_calls = [
            c for c in mock_logger.info.call_args_list if "ampled" in str(c)
        ]
        assert len(sample_calls) >= 1

    def test_no_sampling_when_num_examples_ge_total(self):
        """When num_examples >= total, no sampling occurs (else branch)."""
        ds = _make_dataset()

        items = [
            {
                "id": f"q{i}",
                "prompt": f"Q{i}",
                "answer": f"A{i}",
                "canary": "",
                "reference_steps": "",
            }
            for i in range(3)
        ]
        mock_load_dataset = Mock(return_value=items)

        with patch.dict(
            "sys.modules", {"datasets": Mock(load_dataset=mock_load_dataset)}
        ):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds.load_data(num_examples=10)

        assert len(result) == 3
        # "Loaded" message, not "Sampled"
        load_calls = [
            c for c in mock_logger.info.call_args_list if "oaded" in str(c)
        ]
        assert len(load_calls) >= 1


# ---------------------------------------------------------------------------
# Lines 211-214: Encrypted prompt in _load_from_url
# ---------------------------------------------------------------------------
class TestLoadFromUrlEncrypted:
    def test_url_decrypts_base64_prompt(self):
        """_load_from_url decrypts base64-encoded prompts via XOR."""
        ds = _make_dataset()

        plaintext = b"Real question from URL"
        key = "urlkey"
        key_bytes = key.encode("utf-8")
        xored = bytes(
            [
                plaintext[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(plaintext))
            ]
        )
        encoded_prompt = base64.b64encode(xored).decode("utf-8")

        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda k, default="": {
            "id": "url_enc_1",
            "prompt": encoded_prompt,
            "answer": "plain answer",
            "canary": key,
            "reference_steps": "",
        }.get(k, default)
        mock_df.iterrows.return_value = [(0, row)]

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url()

        assert len(result) == 1
        assert result[0]["problem"] == "Real question from URL"


# ---------------------------------------------------------------------------
# Lines 221-226: Decrypt failure in _load_from_url
# ---------------------------------------------------------------------------
class TestLoadFromUrlDecryptFailure:
    def test_url_decrypt_failure_keeps_original(self):
        """When decrypt fails in _load_from_url, original value is used."""
        ds = _make_dataset()

        bad_bytes = bytes([0xFF, 0xFE, 0xFD])
        key = "x"
        key_bytes = key.encode("utf-8")
        encrypted = bytes(
            [
                bad_bytes[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(bad_bytes))
            ]
        )
        encoded_prompt = base64.b64encode(encrypted).decode("utf-8")

        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda k, default="": {
            "id": "url_fail_1",
            "prompt": encoded_prompt,
            "answer": "plain",
            "canary": key,
            "reference_steps": "",
        }.get(k, default)
        mock_df.iterrows.return_value = [(0, row)]

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds._load_from_url()

        assert len(result) == 1
        mock_logger.warning.assert_called()


# ---------------------------------------------------------------------------
# Line 244: Sampling with seed in _load_from_url
# ---------------------------------------------------------------------------
class TestLoadFromUrlSampling:
    def test_url_sampling_with_seed(self):
        """_load_from_url samples when num_examples < total."""
        ds = _make_dataset()

        mock_df = MagicMock()
        rows = []
        for i in range(10):
            row = MagicMock()
            row.get = lambda k, default="", idx=i: {
                "id": f"url_q{idx}",
                "prompt": f"Question {idx}",
                "answer": f"Answer {idx}",
                "canary": "",
                "reference_steps": "",
            }.get(k, default)
            rows.append((i, row))
        mock_df.iterrows.return_value = rows

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url(num_examples=3, seed=42)

        assert len(result) == 3

    def test_url_no_sampling_when_enough(self):
        """_load_from_url does not sample when num_examples >= total."""
        ds = _make_dataset()

        mock_df = MagicMock()
        rows = []
        for i in range(3):
            row = MagicMock()
            row.get = lambda k, default="", idx=i: {
                "id": f"url_q{idx}",
                "prompt": f"Q{idx}",
                "answer": f"A{idx}",
                "canary": "",
                "reference_steps": "",
            }.get(k, default)
            rows.append((i, row))
        mock_df.iterrows.return_value = rows

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url(num_examples=10)

        assert len(result) == 3
