from unittest.mock import Mock, patch

import pytest

from tradingagents.dataflows.tavily import (
    _search_with_retry,
    get_api_key,
    get_bulk_news_tavily,
)


class TestGetApiKey:
    def test_get_api_key_success(self):
        from tradingagents import config as main_config

        main_config._settings = None
        with patch.dict(
            "os.environ", {"TRADINGAGENTS_TAVILY_API_KEY": "test_key_123"}, clear=False
        ):
            result = get_api_key()
            assert result == "test_key_123"

    def test_get_api_key_missing(self):
        with patch("tradingagents.config.get_settings") as mock_get_settings:
            mock_settings = Mock()
            mock_settings.require_api_key.side_effect = ValueError(
                "tavily API key not configured"
            )
            mock_get_settings.return_value = mock_settings
            with pytest.raises(ValueError, match="tavily API key not configured"):
                get_api_key()


class TestSearchWithRetry:
    def test_successful_search(self):
        mock_client = Mock()
        mock_client.search.return_value = {"results": []}

        result = _search_with_retry(
            client=mock_client,
            query="test query",
            search_depth="advanced",
            topic="news",
            time_range="day",
            max_results=10,
        )

        assert result == {"results": []}
        mock_client.search.assert_called_once()

    @patch("tradingagents.dataflows.tavily.time.sleep")
    def test_retry_on_rate_limit(self, mock_sleep):
        mock_client = Mock()
        mock_client.search.side_effect = [
            RuntimeError("Rate limit exceeded"),
            {"results": []},
        ]

        result = _search_with_retry(
            client=mock_client,
            query="test query",
            search_depth="advanced",
            topic="news",
            time_range="day",
            max_results=10,
        )

        assert result == {"results": []}
        assert mock_client.search.call_count == 2
        assert mock_sleep.call_count == 1

    @patch("tradingagents.dataflows.tavily.time.sleep")
    def test_retry_on_timeout(self, mock_sleep):
        mock_client = Mock()
        mock_client.search.side_effect = [
            TimeoutError("Request timed out"),
            {"results": []},
        ]

        result = _search_with_retry(
            client=mock_client,
            query="test query",
            search_depth="advanced",
            topic="news",
            time_range="day",
            max_results=10,
        )

        assert result == {"results": []}
        assert mock_client.search.call_count == 2

    @patch("tradingagents.dataflows.tavily.time.sleep")
    def test_retry_on_connection_error(self, mock_sleep):
        mock_client = Mock()
        mock_client.search.side_effect = [
            ConnectionError("Connection error occurred"),
            {"results": []},
        ]

        result = _search_with_retry(
            client=mock_client,
            query="test query",
            search_depth="advanced",
            topic="news",
            time_range="day",
            max_results=10,
        )

        assert result == {"results": []}
        assert mock_client.search.call_count == 2

    @patch("tradingagents.dataflows.tavily.time.sleep")
    def test_max_retries_exceeded(self, mock_sleep):
        mock_client = Mock()
        mock_client.search.side_effect = RuntimeError("Rate limit 429")

        with pytest.raises(RuntimeError, match="Rate limit 429"):
            _search_with_retry(
                client=mock_client,
                query="test query",
                search_depth="advanced",
                topic="news",
                time_range="day",
                max_results=10,
                max_retries=3,
            )

        assert mock_client.search.call_count == 3

    def test_non_retryable_error(self):
        mock_client = Mock()
        mock_client.search.side_effect = ValueError("Invalid API key")

        with pytest.raises(ValueError, match="Invalid API key"):
            _search_with_retry(
                client=mock_client,
                query="test query",
                search_depth="advanced",
                topic="news",
                time_range="day",
                max_results=10,
            )

        assert mock_client.search.call_count == 1


class TestGetBulkNewsTavily:
    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", False)
    def test_returns_empty_when_library_not_installed(self):
        result = get_bulk_news_tavily(24)
        assert result == []

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    def test_returns_empty_when_no_api_key(self, mock_get_api_key, mock_client_class):
        mock_get_api_key.side_effect = ValueError("TAVILY_API_KEY not set")

        result = get_bulk_news_tavily(24)

        assert result == []

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_basic_call(self, mock_search, mock_get_api_key, mock_client_class):
        mock_get_api_key.return_value = "test_key"
        mock_search.return_value = {"results": []}

        result = get_bulk_news_tavily(24)

        assert isinstance(result, list)
        assert mock_search.call_count == 5

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_parses_articles(self, mock_search, mock_get_api_key, mock_client_class):
        mock_get_api_key.return_value = "test_key"

        mock_article = {
            "title": "Test Stock News",
            "url": "https://reuters.com/article1",
            "published_date": "2024-01-15T10:30:00Z",
            "content": "This is a test article about stocks.",
        }

        mock_search.return_value = {"results": [mock_article]}

        result = get_bulk_news_tavily(24)

        assert len(result) >= 1
        article = result[0]
        assert article["title"] == "Test Stock News"
        assert article["source"] == "Tavily"
        assert article["url"] == "https://reuters.com/article1"
        assert "published_at" in article
        assert article["content_snippet"] == "This is a test article about stocks."

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_deduplicates_by_url(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        duplicate_article = {
            "title": "Duplicate Article",
            "url": "https://news.com/same-url",
            "published_date": "2024-01-15T10:30:00Z",
            "content": "Duplicate content.",
        }

        mock_search.return_value = {"results": [duplicate_article, duplicate_article]}

        result = get_bulk_news_tavily(24)

        urls = [a["url"] for a in result]
        assert len(urls) == len(set(urls))

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_truncates_long_content(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        long_content = "A" * 1000

        mock_article = {
            "title": "Long Article",
            "url": "https://news.com/article",
            "published_date": "2024-01-15T10:30:00Z",
            "content": long_content,
        }

        mock_search.return_value = {"results": [mock_article]}

        result = get_bulk_news_tavily(24)

        assert len(result[0]["content_snippet"]) == 500

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_time_range_day(self, mock_search, mock_get_api_key, mock_client_class):
        mock_get_api_key.return_value = "test_key"
        mock_search.return_value = {"results": []}

        get_bulk_news_tavily(24)

        call_kwargs = mock_search.call_args_list[0][1]
        assert call_kwargs["time_range"] == "day"

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_time_range_week(self, mock_search, mock_get_api_key, mock_client_class):
        mock_get_api_key.return_value = "test_key"
        mock_search.return_value = {"results": []}

        get_bulk_news_tavily(168)

        call_kwargs = mock_search.call_args_list[0][1]
        assert call_kwargs["time_range"] == "week"

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_time_range_month(self, mock_search, mock_get_api_key, mock_client_class):
        mock_get_api_key.return_value = "test_key"
        mock_search.return_value = {"results": []}

        get_bulk_news_tavily(720)

        call_kwargs = mock_search.call_args_list[0][1]
        assert call_kwargs["time_range"] == "month"

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_handles_missing_published_date(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        mock_article = {
            "title": "Article Without Date",
            "url": "https://news.com/article",
            "content": "Content",
        }

        mock_search.return_value = {"results": [mock_article]}

        result = get_bulk_news_tavily(24)

        assert len(result) == 1
        assert "published_at" in result[0]

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_handles_invalid_date_format(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        mock_article = {
            "title": "Article With Bad Date",
            "url": "https://news.com/article",
            "published_date": "invalid_date_format",
            "content": "Content",
        }

        mock_search.return_value = {"results": [mock_article]}

        result = get_bulk_news_tavily(24)

        assert len(result) == 1
        assert "published_at" in result[0]

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_continues_on_query_failure(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        mock_search.side_effect = [
            RuntimeError("Query failed"),
            {
                "results": [
                    {"title": "Article", "url": "https://test.com", "content": "test"}
                ]
            },
            {"results": []},
            {"results": []},
            {"results": []},
        ]

        result = get_bulk_news_tavily(24)

        assert len(result) > 0

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_skips_articles_without_url(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"

        mock_articles = [
            {"title": "No URL Article", "content": "test"},
            {"title": "Has URL", "url": "https://test.com", "content": "test"},
        ]

        mock_search.return_value = {"results": mock_articles}

        result = get_bulk_news_tavily(24)

        urls = [a["url"] for a in result if a.get("url")]
        assert all(url for url in urls)

    @patch("tradingagents.dataflows.tavily.TAVILY_AVAILABLE", True)
    @patch("tradingagents.dataflows.tavily.TavilyClient")
    @patch("tradingagents.dataflows.tavily.get_api_key")
    @patch("tradingagents.dataflows.tavily._search_with_retry")
    def test_uses_correct_search_parameters(
        self, mock_search, mock_get_api_key, mock_client_class
    ):
        mock_get_api_key.return_value = "test_key"
        mock_search.return_value = {"results": []}

        get_bulk_news_tavily(24)

        call_kwargs = mock_search.call_args_list[0][1]
        assert call_kwargs["search_depth"] == "advanced"
        assert call_kwargs["topic"] == "news"
        assert call_kwargs["max_results"] == 10