# tests/test_research_findings.py
from __future__ import annotations

import json
import os
import unittest
from unittest.mock import patch

os.environ.setdefault("BORDERLESS_INFERENCE_MODE", "hub")
os.environ.setdefault("BORDERLESS_PRELOAD_MODEL", "0")

from langchain_core.messages import AIMessage, ToolMessage

from ui.agent.graph.nodes.helpers import (
    extract_assistant_text,
    extract_thinking_text,
    research_tool_calls,
)
from ui.agent.synthesis import synthesize_finding_from_tool_messages


class ExtractAssistantTextTests(unittest.TestCase):
    def test_plain_string_content(self) -> None:
        message = AIMessage(content="## Eligibility\nYou may qualify.")
        self.assertEqual(extract_assistant_text(message), "## Eligibility\nYou may qualify.")

    def test_list_content_blocks(self) -> None:
        message = AIMessage(
            content=[
                {"type": "text", "text": "Line one"},
                {"type": "text", "text": "Line two"},
            ]
        )
        self.assertEqual(extract_assistant_text(message), "Line one\nLine two")

    def test_reasoning_only_message(self) -> None:
        message = AIMessage(content="", additional_kwargs={"reasoning_content": "Final report"})
        self.assertEqual(extract_assistant_text(message), "Final report")

    def test_strips_think_wrappers(self) -> None:
        open_tag = "<" + "think" + ">"
        close_tag = "</" + "think" + ">"
        message = AIMessage(
            content=(
                f"{open_tag}internal reasoning{close_tag}\n\n"
                "## Findings\nCanada Express Entry."
            )
        )
        extracted = extract_assistant_text(message)
        self.assertIn("## Findings", extracted)
        self.assertNotIn("internal reasoning", extracted)

    def test_prefers_content_over_reasoning(self) -> None:
        message = AIMessage(
            content="Visible answer",
            additional_kwargs={"reasoning_content": "Hidden reasoning"},
        )
        self.assertEqual(extract_assistant_text(message), "Visible answer")
        self.assertEqual(extract_thinking_text(message), "Hidden reasoning")

    def test_extract_thinking_text_from_think_tags(self) -> None:
        open_tag = "<" + "think" + ">"
        close_tag = "</" + "think" + ">"
        message = AIMessage(
            content=(
                f"{open_tag}Compare Canada and Germany pathways.{close_tag}\n\n"
                "## Recommended Countries\n- Canada"
            )
        )
        self.assertIn("Canada", extract_assistant_text(message))
        self.assertEqual(
            extract_thinking_text(message),
            "Compare Canada and Germany pathways.",
        )


class ResearchToolCallsTests(unittest.TestCase):
    def test_parses_tool_calls_from_reasoning(self) -> None:
        reasoning = 'search_immigration_info{"query": "Canada Express Entry"}'
        message = AIMessage(content="", additional_kwargs={"reasoning_content": reasoning})
        calls = research_tool_calls(message)
        self.assertEqual(len(calls), 1)
        self.assertEqual(calls[0][0], "search_immigration_info")
        self.assertEqual(calls[0][1]["query"], "Canada Express Entry")


class SynthesizeFindingTests(unittest.TestCase):
    def test_builds_fallback_from_search_results(self) -> None:
        todo = {"id": 1, "country": "Canada", "methods": "Express Entry"}
        payload = {
            "query": "Canada Express Entry",
            "num_results": 1,
            "results": [
                {
                    "title": "Express Entry",
                    "url": "https://www.canada.ca/en/immigration/express-entry.html",
                    "highlights": ["Skilled workers may apply through Express Entry."],
                }
            ],
        }
        tool_messages = [ToolMessage(content=json.dumps(payload), tool_call_id="abc")]
        summary = synthesize_finding_from_tool_messages(todo, tool_messages)
        self.assertIn("Canada", summary)
        self.assertIn("Express Entry", summary)
        self.assertIn("canada.ca", summary)
        self.assertNotEqual(summary, "")

    def test_returns_empty_without_tool_evidence(self) -> None:
        todo = {"id": 1, "country": "Canada", "methods": "Express Entry"}
        self.assertEqual(
            synthesize_finding_from_tool_messages(todo, []),
            "",
        )


class ResearcherFallbackTests(unittest.TestCase):
    @patch("ui.agent.graph.nodes.researcher.build_llm")
    @patch("ui.agent.graph.nodes.researcher.get_stream_writer")
    @patch("ui.agent.graph.nodes.researcher.execute_tool_call")
    def test_uses_tool_fallback_when_llm_summary_empty(
        self,
        mock_execute_tool_call: unittest.mock.Mock,
        mock_get_stream_writer: unittest.mock.Mock,
        mock_build_llm: unittest.mock.Mock,
    ) -> None:
        from ui.agent.graph.nodes.researcher import researcher_node

        mock_get_stream_writer.return_value = lambda _event: None
        mock_execute_tool_call.return_value = json.dumps(
            {
                "query": "Germany Blue Card",
                "num_results": 1,
                "results": [
                    {
                        "title": "EU Blue Card",
                        "url": "https://www.make-it-in-germany.com/en/visa-residence/eu-blue-card",
                        "highlights": ["The EU Blue Card is for qualified professionals."],
                    }
                ],
            }
        )

        search_response = AIMessage(
            content="",
            tool_calls=[
                {
                    "name": "search_immigration_info",
                    "args": {"query": "Germany Blue Card"},
                    "id": "call-1",
                }
            ],
        )
        empty_summary = AIMessage(content="")

        mock_llm = unittest.mock.Mock()
        mock_bound = unittest.mock.Mock()
        mock_bound.invoke.side_effect = [search_response, empty_summary]
        mock_llm.bind_tools.return_value = mock_bound
        mock_build_llm.return_value = mock_llm

        task = {
            "todo": {"id": 2, "country": "Germany", "methods": "EU Blue Card"},
            "profile_summary": "Software engineer from India.",
        }
        config = {"configurable": {"hf_token": "test-token"}}

        result = researcher_node(task, config)

        summary = result["findings"][0]["summary"]
        self.assertNotEqual(summary, "No findings could be produced for this to-do.")
        self.assertIn("Germany", summary)


if __name__ == "__main__":
    unittest.main()