"""
Unit tests for scripts/run_agent.py and scripts/build_index.py CLIs.

We mock out heavy dependencies (LangGraphAgent, MathLibRetriever) so the tests
exercise only the argparse / control-flow surface of each script.
"""
import contextlib
import importlib
import io
import os
import sys
import unittest
from unittest import mock

_HERE = os.path.dirname(__file__)
_ROOT = os.path.abspath(os.path.join(_HERE, ".."))
sys.path.insert(0, os.path.join(_ROOT, "scripts"))
sys.path.insert(0, os.path.join(_ROOT, "src"))


def _fresh_import(module_name: str):
    """(Re)import a script module so that patches applied via sys.modules
    take effect for top-level `from x import Y` statements."""
    if module_name in sys.modules:
        del sys.modules[module_name]
    return importlib.import_module(module_name)


# ---------------------------------------------------------------------------
# run_agent.py
# ---------------------------------------------------------------------------

class TestRunAgent(unittest.TestCase):
    def _run_main(self, argv, agent_success=True):
        """Patch LangGraphAgent + sys.argv, import run_agent, run main(),
        return (captured_stdout, MockAgent class, mock_instance)."""
        mock_agent_cls = mock.MagicMock(name="LangGraphAgent")
        mock_instance = mock_agent_cls.return_value
        mock_instance.solve_file.return_value = agent_success

        # Inject the mock into the source module so `from langgraph_agent import LangGraphAgent`
        # picks it up at import-time inside run_agent.
        with mock.patch.dict(sys.modules, {"langgraph_agent": mock.MagicMock(LangGraphAgent=mock_agent_cls)}):
            run_agent = _fresh_import("run_agent")
            buf = io.StringIO()
            with mock.patch.object(sys, "argv", argv), contextlib.redirect_stdout(buf):
                run_agent.main()
            return buf.getvalue(), mock_agent_cls, mock_instance

    def test_defaults(self):
        out, cls, inst = self._run_main(["run_agent.py", "problems/simple_add.lean"])
        cls.assert_called_once_with(
            model_name="qwen3-vl:4b",
            max_retries=5,
            index_dir=None,
        )
        inst.solve_file.assert_called_once_with("problems/simple_add.lean")
        self.assertIn("Success", out)

    def test_custom_model(self):
        _, cls, _ = self._run_main(
            ["run_agent.py", "problems/simple_add.lean", "--model", "custom-model"]
        )
        _, kwargs = cls.call_args
        self.assertEqual(kwargs["model_name"], "custom-model")

    def test_retries_forwarded(self):
        _, cls, _ = self._run_main(
            ["run_agent.py", "problems/simple_add.lean", "--retries", "3"]
        )
        _, kwargs = cls.call_args
        self.assertEqual(kwargs["max_retries"], 3)

    def test_index_dir_forwarded(self):
        _, cls, _ = self._run_main(
            ["run_agent.py", "problems/simple_add.lean", "--index-dir", "/tmp/idx"]
        )
        _, kwargs = cls.call_args
        self.assertEqual(kwargs["index_dir"], "/tmp/idx")

    def test_failure_prints_failed_message(self):
        out, _, _ = self._run_main(
            ["run_agent.py", "problems/simple_add.lean"], agent_success=False
        )
        self.assertIn("Failed", out)
        self.assertNotIn("Success", out)

    def test_missing_file_arg_exits(self):
        mock_agent_cls = mock.MagicMock(name="LangGraphAgent")
        with mock.patch.dict(sys.modules, {"langgraph_agent": mock.MagicMock(LangGraphAgent=mock_agent_cls)}):
            run_agent = _fresh_import("run_agent")
            with mock.patch.object(sys, "argv", ["run_agent.py"]):
                # argparse writes the error to stderr; swallow it for clean test output.
                with contextlib.redirect_stderr(io.StringIO()):
                    with self.assertRaises(SystemExit) as ctx:
                        run_agent.main()
                self.assertEqual(ctx.exception.code, 2)


# ---------------------------------------------------------------------------
# build_index.py
# ---------------------------------------------------------------------------

class TestBuildIndex(unittest.TestCase):
    def _run_main(self, argv, is_built=False):
        """Patch MathLibRetriever + sys.argv, import build_index, run main(),
        return (captured_stdout, MockRetriever class, mock_instance)."""
        mock_retr_cls = mock.MagicMock(name="MathLibRetriever")
        mock_instance = mock_retr_cls.return_value
        mock_instance.is_index_built.return_value = is_built
        mock_instance.index_dir = "/fake/index/dir"

        with mock.patch.dict(sys.modules, {"retriever": mock.MagicMock(MathLibRetriever=mock_retr_cls)}):
            build_index = _fresh_import("build_index")
            buf = io.StringIO()
            with mock.patch.object(sys, "argv", argv), contextlib.redirect_stdout(buf):
                build_index.main()
            return buf.getvalue(), mock_retr_cls, mock_instance

    def test_defaults(self):
        out, cls, inst = self._run_main(["build_index.py"])
        cls.assert_called_once_with(index_dir=None)
        inst.build.assert_called_once_with(mathlib_root=None, max_files=None)
        self.assertIn("Done", out)

    def test_mathlib_root_forwarded(self):
        _, _, inst = self._run_main(["build_index.py", "--mathlib-root", "/tmp/fake"])
        _, kwargs = inst.build.call_args
        self.assertEqual(kwargs["mathlib_root"], "/tmp/fake")

    def test_max_files_forwarded(self):
        _, _, inst = self._run_main(["build_index.py", "--max-files", "5"])
        _, kwargs = inst.build.call_args
        self.assertEqual(kwargs["max_files"], 5)

    def test_index_dir_forwarded(self):
        _, cls, _ = self._run_main(["build_index.py", "--index-dir", "/tmp/idx"])
        cls.assert_called_once_with(index_dir="/tmp/idx")

    def test_existing_index_skips_rebuild(self):
        out, _, inst = self._run_main(["build_index.py"], is_built=True)
        inst.build.assert_not_called()
        self.assertIn("Index already exists", out)

    def test_existing_index_rebuilds_when_max_files_set(self):
        out, _, inst = self._run_main(
            ["build_index.py", "--max-files", "5"], is_built=True
        )
        inst.build.assert_called_once_with(mathlib_root=None, max_files=5)
        self.assertNotIn("Index already exists", out)


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