"""Tests for Jupyter/IPython integration."""

from __future__ import annotations

from pathlib import Path
from typing import Any
from unittest.mock import patch

import pytest

from llmdebug.jupyter_plugin import (
    _extract_jupyter_context,
    _format_full_snapshot_html,
    _format_full_snapshot_text,
    _format_snapshot_banner_html,
    _format_snapshot_banner_text,
    _is_notebook,
    install,
    load_ipython_extension,
    uninstall,
    unload_ipython_extension,
)

from .conftest import read_latest_snapshot

# ---------------------------------------------------------------------------
# Fixtures and helpers
# ---------------------------------------------------------------------------

SAMPLE_SNAPSHOT: dict[str, Any] = {
    "name": "jupyter_cell",
    "exception": {
        "type": "ValueError",
        "message": "invalid literal for int(): 'abc'",
        "error_category": {
            "category": "value_error",
            "suggestion": "Check the input type before conversion.",
        },
    },
    "frames": [
        {
            "file": "/tmp/ipykernel/cell.py",
            "file_rel": "cell.py",
            "line": 3,
            "function": "<module>",
            "code": "x = int('abc')",
            "locals": {
                "data": [1, 2, 3],
                "x": "abc",
            },
        }
    ],
    "jupyter": {
        "execution_count": 5,
        "cell_source": "x = int('abc')",
    },
}


class _FakeTerminalShell:
    """Minimal fake IPython shell for terminal mode."""

    def __init__(self) -> None:
        self.execution_count = 42
        self._original_showtraceback = self.showtraceback
        self._registered_magics: list[Any] = []
        self.history_manager = _FakeHistoryManager()

    def showtraceback(self, *args: Any, **kwargs: Any) -> None:
        pass

    def register_magics(self, cls: Any) -> None:
        self._registered_magics.append(cls)


class _FakeNotebookShell(_FakeTerminalShell):
    """Fake shell whose class name is ZMQInteractiveShell."""

    pass


# Rename the class so type().__name__ returns the right thing
_FakeNotebookShell.__name__ = "ZMQInteractiveShell"
_FakeNotebookShell.__qualname__ = "ZMQInteractiveShell"


class _FakeHistoryManager:
    def get_tail(self, n: int = 1, raw: bool = True) -> list[tuple[int, int, str]]:
        return [(1, 1, "x = 1/0")]


def _make_ipython(*, notebook: bool = False) -> Any:
    """Create a fake IPython shell."""
    if notebook:
        return _FakeNotebookShell()
    return _FakeTerminalShell()


@pytest.fixture(autouse=True)
def _cleanup_jupyter():
    """Ensure jupyter plugin is uninstalled after each test."""
    yield
    uninstall()


@pytest.fixture()
def snap_dir(tmp_path: Path) -> Path:
    """Create a temp snapshot directory."""
    d = tmp_path / ".llmdebug"
    d.mkdir()
    return d


# ---------------------------------------------------------------------------
# Installation tests
# ---------------------------------------------------------------------------


class TestInstall:
    def test_install_wraps_showtraceback(self):
        ip = _make_ipython()
        original = ip.showtraceback
        install(ip, out_dir="/tmp/test_llmdebug")

        assert ip.showtraceback is not original
        assert len(ip._registered_magics) > 0

    def test_install_idempotent(self):
        ip = _make_ipython()
        install(ip, out_dir="/tmp/test_llmdebug")
        first_wrapper = ip.showtraceback

        install(ip, out_dir="/tmp/test_llmdebug")
        assert ip.showtraceback is first_wrapper

    def test_uninstall_restores_original(self):
        ip = _make_ipython()
        # Bound methods are created fresh each access; compare via __func__
        original_func = _FakeTerminalShell.showtraceback
        install(ip, out_dir="/tmp/test_llmdebug")

        # After install, showtraceback should be a closure (not the class method)
        assert ip.showtraceback is not original_func

        uninstall()
        # After uninstall, the stored bound method is restored as instance attr
        # It should point to the same underlying function
        restored = ip.showtraceback
        assert getattr(restored, "__func__", restored) is original_func

    def test_uninstall_idempotent(self):
        uninstall()
        uninstall()

    def test_install_explicit_ipython(self):
        ip = _make_ipython()
        install(ip, out_dir="/tmp/test_llmdebug")
        assert len(ip._registered_magics) > 0


# ---------------------------------------------------------------------------
# Exception capture tests
# ---------------------------------------------------------------------------


class TestExceptionCapture:
    def test_calls_original_showtraceback_first(self):
        ip = _make_ipython()
        original_called = []

        def tracking_showtraceback(*args: Any, **kwargs: Any) -> None:
            original_called.append(True)

        ip.showtraceback = tracking_showtraceback
        install(ip, out_dir="/tmp/test_llmdebug_capture")

        wrapper = ip.showtraceback

        try:
            raise ValueError("test error")
        except ValueError:
            with (
                patch("llmdebug.jupyter_plugin._capture_safely"),
                patch("llmdebug.jupyter_plugin._display_snapshot_banner"),
            ):
                wrapper()

        assert len(original_called) == 1

    def test_captures_exception_with_jupyter_context(self, snap_dir: Path):
        ip = _make_ipython()
        install(ip, out_dir=str(snap_dir))

        wrapper = ip.showtraceback

        try:
            raise ValueError("test capture")
        except ValueError:
            with patch("llmdebug.jupyter_plugin._display_snapshot_banner"):
                wrapper()

        snapshot = read_latest_snapshot(snap_dir)
        assert snapshot["exception"]["type"] == "ValueError"
        assert "jupyter" in snapshot
        assert snapshot["jupyter"]["execution_count"] == 42

    def test_rate_limiting_prevents_flood(self, snap_dir: Path):
        ip = _make_ipython()
        install(
            ip,
            out_dir=str(snap_dir),
            max_captures_per_minute=2,
        )

        wrapper = ip.showtraceback

        for _ in range(5):
            try:
                raise ValueError("repeated error")
            except ValueError:
                with patch("llmdebug.jupyter_plugin._display_snapshot_banner"):
                    wrapper()

        from llmdebug._snapshot_loader import find_all_snapshots

        snaps = find_all_snapshots(str(snap_dir))
        assert len(snaps) <= 3  # 2 allowed + possible timing slack

    def test_capture_failure_does_not_propagate(self):
        ip = _make_ipython()
        original_called = []

        def tracking_showtraceback(*args: Any, **kwargs: Any) -> None:
            original_called.append(True)

        ip.showtraceback = tracking_showtraceback
        install(ip, out_dir="/tmp/test_llmdebug_fail")

        wrapper = ip.showtraceback

        try:
            raise ValueError("test")
        except ValueError:
            with patch(
                "llmdebug.jupyter_plugin._capture_safely",
                side_effect=Exception("capture broke"),
            ):
                # Should NOT raise
                wrapper()

        assert len(original_called) == 1


# ---------------------------------------------------------------------------
# Display tests
# ---------------------------------------------------------------------------


class TestDisplay:
    def test_banner_html_contains_exception(self):
        html = _format_snapshot_banner_html(SAMPLE_SNAPSHOT)
        assert "ValueError" in html
        assert "llmdebug" in html
        assert "cell.py:3" in html

    def test_banner_text_contains_exception(self):
        text = _format_snapshot_banner_text(SAMPLE_SNAPSHOT)
        assert "ValueError" in text
        assert "cell.py:3" in text
        assert "%llmdebug" in text

    def test_full_html_contains_details(self):
        html = _format_full_snapshot_html(SAMPLE_SNAPSHOT)
        assert "ValueError" in html
        assert "invalid literal" in html
        assert "cell.py" in html
        assert "value_error" in html
        assert "Cell In[5]" in html

    def test_full_text_contains_details(self):
        text = _format_full_snapshot_text(SAMPLE_SNAPSHOT)
        assert "ValueError" in text
        assert "invalid literal" in text
        assert "cell.py" in text
        assert "value_error" in text
        assert "Cell In[5]" in text

    def test_no_display_when_no_snapshot(self, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        from llmdebug.jupyter_plugin import _display_snapshot_banner

        _display_snapshot_banner(ip, "/nonexistent/dir")
        captured = capsys.readouterr()
        assert "llmdebug" not in captured.out

    def test_html_escapes_special_chars(self):
        snapshot: dict[str, Any] = {
            "exception": {
                "type": "ValueError",
                "message": '<script>alert("xss")</script>',
            },
            "frames": [{"file": "test.py", "line": 1, "function": "f", "code": ""}],
        }
        html = _format_snapshot_banner_html(snapshot)
        assert "<script>" not in html
        assert "&lt;script&gt;" in html

    def test_banner_truncates_long_messages(self):
        snapshot: dict[str, Any] = {
            "exception": {
                "type": "ValueError",
                "message": "x" * 200,
            },
            "frames": [{}],
        }
        html = _format_snapshot_banner_html(snapshot)
        assert "..." in html

    def test_full_html_shows_array_locals(self):
        snapshot: dict[str, Any] = {
            "exception": {"type": "ValueError", "message": "shape mismatch"},
            "frames": [
                {
                    "file": "train.py",
                    "line": 10,
                    "function": "step",
                    "code": "x + y",
                    "locals": {
                        "x": {
                            "__array__": "numpy.ndarray",
                            "shape": [32, 64],
                            "dtype": "float32",
                        },
                    },
                }
            ],
        }
        html = _format_full_snapshot_html(snapshot)
        assert "numpy.ndarray" in html
        assert "shape=[32, 64]" in html

    def test_displays_text_in_terminal(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython(notebook=False)
        from llmdebug.jupyter_plugin import _display_snapshot_banner

        with patch("llmdebug.output.get_latest_snapshot", return_value=SAMPLE_SNAPSHOT):
            _display_snapshot_banner(ip, str(snap_dir))

        captured = capsys.readouterr()
        assert "ValueError" in captured.out
        assert "%llmdebug" in captured.out


# ---------------------------------------------------------------------------
# Magic command tests
# ---------------------------------------------------------------------------


class TestMagicCommands:
    def _get_magic_instance(self, ip: Any, snap_dir: str) -> Any:
        """Install and return an instantiated magic class."""
        install(ip, out_dir=snap_dir)
        assert len(ip._registered_magics) > 0
        magics_cls = ip._registered_magics[0]
        # IPython Magics expects a Configurable shell, pass None to skip validation
        instance = magics_cls(shell=None)
        # Manually set shell for _is_notebook() checks
        instance.shell = ip
        return instance

    def test_magic_show_no_snapshot(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        magic.llmdebug("show")
        captured = capsys.readouterr()
        assert "No snapshots found" in captured.out

    def test_magic_show_displays_snapshot(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        with patch("llmdebug.output.get_latest_snapshot", return_value=SAMPLE_SNAPSHOT):
            magic.llmdebug("show")

        captured = capsys.readouterr()
        assert "ValueError" in captured.out

    def test_magic_list_shows_recent(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        with patch(
            "llmdebug._snapshot_loader.find_all_snapshots",
            return_value=[
                {
                    "path": snap_dir / "snap1.json",
                    "exc_type": "ValueError",
                    "exc_message": "test",
                    "file": "test.py",
                    "line": 1,
                }
            ],
        ):
            magic.llmdebug("list")

        captured = capsys.readouterr()
        assert "ValueError" in captured.out

    def test_magic_list_no_snapshots(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        magic.llmdebug("list")
        captured = capsys.readouterr()
        assert "No snapshots found" in captured.out

    def test_magic_hypothesize_generates_output(
        self, snap_dir: Path, capsys: pytest.CaptureFixture[str]
    ):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        with patch("llmdebug.output.get_latest_snapshot", return_value=SAMPLE_SNAPSHOT):
            magic.llmdebug("hypothesize")

        captured = capsys.readouterr()
        assert len(captured.out) > 0

    def test_magic_diff_no_snapshots(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        magic.llmdebug("diff")
        captured = capsys.readouterr()
        # Should report not enough snapshots
        assert "Need at least 2" in captured.out or "No snapshots found" in captured.out

    def test_magic_install_uninstall(self, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, "/tmp/test_llmdebug_magic")

        magic.llmdebug("uninstall")
        captured = capsys.readouterr()
        assert "uninstalled" in captured.out

    def test_magic_config_shows_settings(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        magic.llmdebug("config")
        captured = capsys.readouterr()
        assert "out_dir" in captured.out
        assert "Active configuration" in captured.out

    def test_magic_unknown_shows_help(self, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, "/tmp/test_llmdebug_help")

        magic.llmdebug("nonexistent_command")
        captured = capsys.readouterr()
        assert "Usage:" in captured.out
        assert "show" in captured.out
        assert "hypothesize" in captured.out

    def test_magic_default_is_show(self, snap_dir: Path, capsys: pytest.CaptureFixture[str]):
        ip = _make_ipython()
        magic = self._get_magic_instance(ip, str(snap_dir))

        magic.llmdebug("")
        captured = capsys.readouterr()
        assert "No snapshots found" in captured.out


# ---------------------------------------------------------------------------
# Extension API tests
# ---------------------------------------------------------------------------


class TestExtensionAPI:
    def test_load_extension_installs_and_registers_magics(self):
        ip = _make_ipython()
        load_ipython_extension(ip)

        assert len(ip._registered_magics) > 0

    def test_unload_extension_uninstalls(self):
        ip = _make_ipython()
        load_ipython_extension(ip)

        unload_ipython_extension(ip)
        from llmdebug.jupyter_plugin import _installed

        assert not _installed


# ---------------------------------------------------------------------------
# Jupyter context extraction
# ---------------------------------------------------------------------------


class TestJupyterContext:
    def test_extracts_execution_count(self):
        ip = _make_ipython()
        ip.execution_count = 7
        ctx = _extract_jupyter_context(ip)
        assert ctx["execution_count"] == 7

    def test_extracts_cell_source(self):
        ip = _make_ipython()
        ip.history_manager = _FakeHistoryManager()
        ctx = _extract_jupyter_context(ip)
        assert ctx["cell_source"] == "x = 1/0"

    def test_truncates_long_cell_source(self):
        ip = _make_ipython()

        class LongHistory:
            def get_tail(self, n: int = 1, raw: bool = True) -> list[tuple[int, int, str]]:
                return [(1, 1, "x" * 600)]

        ip.history_manager = LongHistory()
        ctx = _extract_jupyter_context(ip)
        assert len(ctx["cell_source"]) == 500
        assert ctx["cell_source"].endswith("...")

    def test_handles_missing_history_manager(self):
        ip = _make_ipython()
        ip.history_manager = None
        ctx = _extract_jupyter_context(ip)
        assert "execution_count" in ctx
        assert "cell_source" not in ctx

    def test_handles_missing_execution_count(self):
        class BareShell:
            pass

        ip = BareShell()
        ctx = _extract_jupyter_context(ip)
        assert "execution_count" not in ctx


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


class TestEdgeCases:
    def test_handles_keyboard_interrupt(self, snap_dir: Path):
        ip = _make_ipython()
        install(ip, out_dir=str(snap_dir))

        wrapper = ip.showtraceback

        try:
            raise KeyboardInterrupt()
        except KeyboardInterrupt:
            with patch("llmdebug.jupyter_plugin._display_snapshot_banner") as mock_display:
                wrapper()

            mock_display.assert_not_called()

    def test_handles_system_exit(self, snap_dir: Path):
        ip = _make_ipython()
        install(ip, out_dir=str(snap_dir))

        wrapper = ip.showtraceback

        try:
            raise SystemExit(1)
        except SystemExit:
            with patch("llmdebug.jupyter_plugin._display_snapshot_banner") as mock_display:
                wrapper()

            mock_display.assert_not_called()

    def test_is_notebook_detects_zmq(self):
        ip = _make_ipython(notebook=True)
        assert _is_notebook(ip) is True

    def test_is_notebook_detects_terminal(self):
        ip = _make_ipython(notebook=False)
        assert _is_notebook(ip) is False

    def test_empty_snapshot_frames(self):
        snapshot: dict[str, Any] = {
            "exception": {"type": "RuntimeError", "message": "boom"},
            "frames": [],
        }
        html = _format_snapshot_banner_html(snapshot)
        assert "RuntimeError" in html

        text = _format_full_snapshot_text(snapshot)
        assert "RuntimeError" in text


# ---------------------------------------------------------------------------
# __init__.py integration
# ---------------------------------------------------------------------------


class TestInitIntegration:
    def test_load_jupyter_exported(self):
        import llmdebug

        assert hasattr(llmdebug, "load_jupyter")
        assert "load_jupyter" in llmdebug.__all__

    def test_load_ipython_extension_forwarding(self):
        import llmdebug

        ip = _make_ipython()
        llmdebug.load_ipython_extension(ip)
        assert len(ip._registered_magics) > 0
        uninstall()

    def test_unload_ipython_extension_forwarding(self):
        import llmdebug

        ip = _make_ipython()
        llmdebug.load_ipython_extension(ip)
        llmdebug.unload_ipython_extension(ip)

        from llmdebug.jupyter_plugin import _installed

        assert not _installed
