import types
from asyncio import CancelledError
from contextlib import AsyncExitStack
from io import StringIO
from unittest import IsolatedAsyncioTestCase, TestCase
from unittest.mock import AsyncMock, MagicMock, patch

import numpy as np

from avalan.entities import (
    Message,
    MessageRole,
    TextPartition,
    ToolCallContext,
)
from avalan.tool.browser import (
    BrowserTool,
    BrowserToolSet,
    BrowserToolSettings,
)


class BrowserToolCallTestCase(IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.page = MagicMock()
        self.page.goto = AsyncMock()
        self.page.content = AsyncMock(return_value="<html>hi</html>")
        self.page.close = AsyncMock()

        self.browser = MagicMock()
        self.browser.new_page = AsyncMock(return_value=self.page)
        self.browser.close = AsyncMock()

        self.browser_type = MagicMock()
        self.browser_type.launch = AsyncMock(return_value=self.browser)

        self.client = MagicMock()
        self.client.firefox = self.browser_type

        result = types.SimpleNamespace(text_content="parsed")
        converter = MagicMock(return_value=result)

        self.mark_patch = patch(
            "avalan.tool.browser.MarkItDown",
            return_value=MagicMock(convert=converter),
        )
        self.mark_patch.start()
        self.tool = BrowserTool(BrowserToolSettings(), self.client)

    async def asyncTearDown(self):
        self.mark_patch.stop()


class BrowserToolSetTestCase(IsolatedAsyncioTestCase):
    async def test_init_and_enter(self):
        dummy_stack = AsyncMock(spec=AsyncExitStack)
        dummy_stack.enter_async_context = AsyncMock(return_value="client2")
        dummy_stack.__aexit__ = AsyncMock(return_value=False)

        dummy_tool = MagicMock()
        dummy_tool.with_client = MagicMock(return_value=dummy_tool)

        async def dummy_aenter(self):
            return self

        with (
            patch(
                "avalan.tool.browser.async_playwright", return_value="client1"
            ),
            patch("avalan.tool.browser.BrowserTool", return_value=dummy_tool),
            patch("avalan.tool.browser.ToolSet.__aenter__", dummy_aenter),
        ):
            toolset = BrowserToolSet(
                settings=BrowserToolSettings(),
                exit_stack=dummy_stack,
                namespace="b",
            )
            self.assertEqual(toolset.namespace, "b")
            self.assertEqual(toolset._client, "client1")
            self.assertEqual(toolset.tools, [dummy_tool])

            result = await toolset.__aenter__()

        self.assertIs(result, toolset)
        dummy_stack.enter_async_context.assert_awaited_once_with("client1")
        dummy_tool.with_client.assert_called_once_with("client2")

    async def test_init_uses_default_exit_stack(self):
        dummy_tool = MagicMock()

        with (
            patch(
                "avalan.tool.browser.async_playwright", return_value="client1"
            ),
            patch("avalan.tool.browser.BrowserTool", return_value=dummy_tool),
        ):
            toolset = BrowserToolSet(settings=BrowserToolSettings())

        self.assertIsInstance(toolset._exit_stack, AsyncExitStack)
        self.assertEqual(toolset._client, "client1")
        self.assertEqual(toolset.tools, [dummy_tool])


class BrowserToolWithClientTestCase(TestCase):
    def test_with_client(self):
        client1 = MagicMock()
        client2 = MagicMock()
        tool = BrowserTool(BrowserToolSettings(), client1)
        result = tool.with_client(client2)
        self.assertIs(result, tool)
        self.assertIs(tool._client, client2)


class BrowserToolCallSearchTestCase(IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.client = MagicMock()
        self.mark_patch = patch("avalan.tool.browser.MarkItDown")
        self.mark_patch.start()

    async def asyncTearDown(self):
        self.mark_patch.stop()

    async def test_call_search_with_partitioner(self):
        partitions = [
            TextPartition(
                data="a", embeddings=np.array([0.1, 0.2]), total_tokens=1
            ),
            TextPartition(
                data="b", embeddings=np.array([0.2, 0.3]), total_tokens=1
            ),
        ]

        class DummyPartitioner:
            def __init__(self) -> None:
                self.call_mock = AsyncMock(return_value=partitions)
                self.sentence_model = AsyncMock(
                    return_value=np.array([[0.1, 0.2]])
                )

            async def __call__(self, text: str):
                return await self.call_mock(text)

        partitioner = DummyPartitioner()

        settings = BrowserToolSettings(search=True)
        tool = BrowserTool(settings, self.client, partitioner=partitioner)
        tool._read = AsyncMock(return_value="html")

        index = MagicMock()
        index.add = MagicMock()
        index.search = MagicMock(
            return_value=(np.array([[0.05]]), np.array([[0]]))
        )

        with patch(
            "avalan.tool.browser.IndexFlatL2", return_value=index
        ) as idx_patch:
            ctx = ToolCallContext(
                input=Message(role=MessageRole.USER, content="q")
            )
            result = await tool("http://t", context=ctx)

        partitioner.call_mock.assert_awaited_once_with("html")
        partitioner.sentence_model.assert_awaited_once_with(["q"])
        idx_patch.assert_called_once_with(partitions[0].embeddings.shape[0])
        index.add.assert_called_once()
        index.search.assert_called_once()
        self.assertEqual(result, "a\nb")


class BrowserToolCallDebugTestCase(IsolatedAsyncioTestCase):
    async def test_call_debug_source(self):
        debug_io = StringIO("debug")
        settings = BrowserToolSettings(
            debug=True, debug_url="http://t", debug_source=debug_io
        )
        tool = BrowserTool(settings, MagicMock())
        result = await tool("http://t", context=ToolCallContext())
        self.assertEqual(result, "debug")


class BrowserToolReadTestCase(IsolatedAsyncioTestCase):
    async def test_read_launches_browser_and_page(self):
        settings = BrowserToolSettings()
        page = MagicMock()
        page.goto = AsyncMock(
            return_value=MagicMock(
                headers={"content-type": "text/html; charset=utf-8"}
            )
        )
        page.content = AsyncMock(return_value="<html>hi</html>")
        page.close = AsyncMock()

        browser = MagicMock()
        browser.new_page = AsyncMock(return_value=page)
        browser.close = AsyncMock()

        browser_type = MagicMock()
        browser_type.launch = AsyncMock(return_value=browser)

        client = MagicMock()
        client.firefox = browser_type

        with patch("avalan.tool.browser.MarkItDown") as md_patch:
            md_patch.return_value.convert_stream.return_value = (
                types.SimpleNamespace(text_content="parsed")
            )
            tool = BrowserTool(settings, client)
        tool._md.convert_stream = MagicMock(
            return_value=types.SimpleNamespace(text_content="parsed")
        )

        content = await tool._read("http://t")

        browser_type.launch.assert_awaited_once()
        browser.new_page.assert_awaited_once()
        page.goto.assert_awaited_once_with("http://t")
        page.content.assert_awaited_once()
        self.assertEqual(content, "parsed")

    async def test_read_uses_markdownify_when_markitdown_is_missing(self):
        settings = BrowserToolSettings()
        page = MagicMock()
        page.goto = AsyncMock(
            return_value=MagicMock(headers={"content-type": "text/html"})
        )
        page.content = AsyncMock(return_value="<html><body>hi</body></html>")
        browser = MagicMock()
        browser.new_page = AsyncMock(return_value=page)
        browser_type = MagicMock()
        browser_type.launch = AsyncMock(return_value=browser)
        client = MagicMock()
        client.firefox = browser_type

        with (
            patch("avalan.tool.browser.MarkItDown", None),
            patch(
                "avalan.tool.browser.markdownify_html",
                return_value="hi",
            ) as markdownify_patch,
        ):
            tool = BrowserTool(settings, client)
            content = await tool._read("http://t")

        markdownify_patch.assert_called_once_with(
            "<html><body>hi</body></html>"
        )
        self.assertEqual(content, "hi")

    async def test_aexit_closes_resources(self):
        tool = BrowserTool(BrowserToolSettings(), MagicMock())
        page = AsyncMock()
        browser = AsyncMock()
        tool._page = page
        tool._browser = browser
        with patch(
            "avalan.tool.browser.Tool.__aexit__",
            AsyncMock(return_value=True),
        ) as base_exit:
            result = await tool.__aexit__(None, None, None)
        page.close.assert_awaited_once()
        browser.close.assert_awaited_once()
        base_exit.assert_awaited_once_with(None, None, None)
        self.assertTrue(result)

    async def test_aexit_suppresses_close_errors_on_interrupt(self):
        tool = BrowserTool(BrowserToolSettings(), MagicMock())
        page = AsyncMock()
        page.close.side_effect = Exception("connection closed")
        browser = AsyncMock()
        tool._page = page
        tool._browser = browser

        with patch(
            "avalan.tool.browser.Tool.__aexit__",
            AsyncMock(return_value=False),
        ) as base_exit:
            result = await tool.__aexit__(
                CancelledError, CancelledError(), None
            )

        page.close.assert_awaited_once()
        browser.close.assert_awaited_once()
        base_exit.assert_awaited_once()
        self.assertFalse(result)

    async def test_aexit_raises_close_errors_without_interrupt(self):
        tool = BrowserTool(BrowserToolSettings(), MagicMock())
        tool._page = AsyncMock()
        tool._page.close.side_effect = Exception("connection closed")

        with self.assertRaisesRegex(Exception, "connection closed"):
            await tool.__aexit__(None, None, None)

    async def test_close_resource_reraises_timeout_without_interrupt(self):
        tool = BrowserTool(BrowserToolSettings(), MagicMock())
        resource = MagicMock()
        resource.close = AsyncMock(side_effect=TimeoutError)

        with self.assertRaises(TimeoutError):
            await tool._close_resource(resource, interrupted_exit=False)


class BrowserToolCallSearchSkipTestCase(IsolatedAsyncioTestCase):
    async def test_call_search_skips_empty_match(self):
        partitions = [
            TextPartition(
                data="", embeddings=np.array([0.1, 0.2]), total_tokens=1
            ),
            TextPartition(
                data="b", embeddings=np.array([0.2, 0.3]), total_tokens=1
            ),
        ]

        class DummyPartitioner:
            def __init__(self) -> None:
                self.call_mock = AsyncMock(return_value=partitions)
                self.sentence_model = AsyncMock(
                    return_value=np.array([[0.1, 0.2]])
                )

            async def __call__(self, text: str):
                return await self.call_mock(text)

        partitioner = DummyPartitioner()

        settings = BrowserToolSettings(search=True)
        tool = BrowserTool(settings, MagicMock(), partitioner=partitioner)
        tool._read = AsyncMock(return_value="html")

        index = MagicMock()
        index.add = MagicMock()
        index.search = MagicMock(
            return_value=(np.array([[0.05]]), np.array([[0]]))
        )

        with patch("avalan.tool.browser.IndexFlatL2", return_value=index):
            ctx = ToolCallContext(
                input=Message(role=MessageRole.USER, content="q")
            )
            result = await tool("http://t", context=ctx)

        self.assertEqual(result, "html")
        index.add.assert_called_once()
        index.search.assert_called_once()
        partitioner.call_mock.assert_awaited_once_with("html")
        partitioner.sentence_model.assert_awaited_once_with(["q"])
