from unittest.mock import patch

import pytest

from notebooklm_tools.core.client import NotebookLMClient


@pytest.fixture
def mock_client():
    cookies = {"SID": "test_sid"}
    with patch.object(NotebookLMClient, "_refresh_auth_tokens"):
        client = NotebookLMClient(cookies=cookies, csrf_token="old_token", session_id="old_sid")
        return client


class TestNotebookLMClientFiltering:
    def test_poll_research_filtering(self, mock_client):
        """Test poll_research filters by task_id and handles status codes correctly."""

        # Mock API response for poll_research (RPC e3bVqc)
        # Structure: [[task_id, [query, [text, type], mode, [sources], status_code]], ...]

        task1 = [
            "task_uuid_1",
            [
                None,
                ["Query 1", 1],
                1,
                [[], "Summary 1"],
                2,  # Status 2 = Completed
            ],
        ]

        task2 = [
            "task_uuid_2",
            [
                None,
                ["Query 2", 1],
                1,
                [[], "Summary 2"],
                6,  # Status 6 = Imported (also Completed)
            ],
        ]

        task3 = [
            "task_uuid_3",
            [
                None,
                ["Query 3", 1],
                1,
                [[], "Summary 3"],
                1,  # Status 1 = In Progress
            ],
        ]

        # Outer wrapper: [[tasks], [timestamps]]
        mock_raw_result = [[task3, task2, task1]]  # Tasks returned in some order

        with (
            patch.object(mock_client, "_get_client") as mock_get_client,
            patch.object(mock_client, "_extract_rpc_result", return_value=mock_raw_result),
        ):
            # Mock the http client to avoid network calls
            mock_http = mock_get_client.return_value
            mock_http.post.return_value.status_code = 200
            mock_http.post.return_value.text = (
                ")]}'\n10\n[]"  # Dummy content, ignored since we mock extract
            )

            # 1. Test filtering by task_id (task 1)
            result1 = mock_client.poll_research("nb_id", target_task_id="task_uuid_1")
            assert result1 is not None
            assert result1["task_id"] == "task_uuid_1"
            assert result1["status"] == "completed"  # Status 2 -> completed

            # 2. Test filtering by task_id (task 2 - status 6)
            result2 = mock_client.poll_research("nb_id", target_task_id="task_uuid_2")
            assert result2 is not None
            assert result2["task_id"] == "task_uuid_2"
            assert result2["status"] == "completed"  # Status 6 -> completed (Fix)

            # 3. Test filtering by task_id (task 3 - in progress)
            result3 = mock_client.poll_research("nb_id", target_task_id="task_uuid_3")
            assert result3 is not None
            assert result3["task_id"] == "task_uuid_3"
            assert result3["status"] == "in_progress"  # Status 1 -> in_progress

            # 4. Test unknown task_id — Issue #106 fix: should fallback to
            # in-progress task instead of returning None
            result_unknown = mock_client.poll_research("nb_id", target_task_id="unknown_uuid")
            assert result_unknown is not None
            assert result_unknown["task_id"] == "task_uuid_3"  # in-progress task preferred
            assert result_unknown["status"] == "in_progress"

            # 5. Test default behavior (no task_id) - returns first valid task (task3 in our mock list)
            # Logic: "Return the most recent (first) task"
            result_default = mock_client.poll_research("nb_id")
            assert result_default is not None
            assert result_default["task_id"] == "task_uuid_3"