import unittest
from unittest.mock import MagicMock, patch

from agents.research_graph import (
    plan_search_node, 
    execute_search_node, 
    route_research,
    save_note_node,
    finalize_research_node
)


class ResearchAgentNodesTests(unittest.TestCase):
    @patch("agents.research_graph.generate_sub_queries")
    def test_plan_search_node_extracts_queries(self, mock_generate_sub_queries):
        state = {"topic": "AI agents", "debug_mode": False}
        llm = MagicMock()
        mock_generate_sub_queries.return_value = [
            "what are ai agents",
            "benefits of ai agents",
            "future of ai agents",
        ]
        
        updated_state = plan_search_node(state, llm=llm, max_queries=3)
        
        expected_queries = ["what are ai agents", "benefits of ai agents", "future of ai agents"]
        self.assertEqual(updated_state["search_queries"], expected_queries)

    @patch("agents.research_graph.generate_sub_queries")
    def test_plan_search_node_respects_max_queries(self, mock_generate_sub_queries):
        state = {"topic": "AI agents", "debug_mode": False}
        llm = MagicMock()
        mock_generate_sub_queries.return_value = ["q1", "q2"]
        
        updated_state = plan_search_node(state, llm=llm, max_queries=2)
        
        self.assertEqual(len(updated_state["search_queries"]), 2)
        self.assertEqual(updated_state["search_queries"], ["q1", "q2"])

    @patch("agents.research_graph.fetch_urls_parallel")
    @patch("agents.research_graph.web_search")
    def test_execute_search_node_collects_results(self, mock_search, mock_fetch_parallel):
        mock_search.return_value = [
            {"url": "http://e1.com", "title": "T1", "snippet": "S1"},
            {"url": "http://e2.com", "title": "T2", "snippet": "S2"}
        ]
        mock_fetch_parallel.return_value = [
            {"source_url": "http://e1.com", "page_text": "Full Content"},
        ]

        state = {"search_queries": ["query 1"], "debug_mode": False}
        updated_state = execute_search_node(state)

        self.assertEqual(len(updated_state["search_results"]), 2)
        self.assertEqual(updated_state["search_results"][0]["content"], "Full Content")

    @patch("agents.research_graph.save_note")
    def test_save_note_node_calls_save_note_when_requested(self, mock_save):
        mock_save.return_value = "note_123"
        state = {
            "topic": "Research Topic",
            "synthesis": "The Synthesis",
            "save_to_notes": True,
            "search_results": [{"url": "U1", "snippet": "S1"}]
        }
        
        updated_state = save_note_node(state)
        
        self.assertEqual(updated_state["note_id"], "note_123")
        mock_save.assert_called_once()

    def test_save_note_node_skips_when_not_requested(self):
        state = {"save_to_notes": False}
        updated_state = save_note_node(state)
        self.assertNotIn("note_id", updated_state)

    def test_finalize_research_node_handles_no_results_error(self):
        state = {"topic": "Unknown topic", "error": "no_search_results"}
        updated_state = finalize_research_node(state)
        self.assertIn("unable to find", updated_state["synthesis"])

    def test_finalize_research_node_builds_lightweight_summary_when_results_exist(self):
        state = {
            "topic": "Topic",
            "search_results": [
                {"title": "Result One"},
                {"title": "Result Two"},
            ],
        }
        updated_state = finalize_research_node(state)
        self.assertIn("Collected web sources", updated_state["synthesis"])

    def test_route_research_returns_finalize_on_error(self):
        state = {"error": "no_search_results"}
        self.assertEqual(route_research(state), "finalize")

    def test_route_research_returns_finalize_when_no_error(self):
        state = {"search_results": [{"url": "http://e1.com"}]}
        self.assertEqual(route_research(state), "finalize")