"""Tests for figure auto-capture in _send_context_notification."""

import asyncio
import base64
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from skop_runner.message_handler import MessageHandler, MAX_FIGURE_BASE64_BYTES
from skop_runner.mock_backend import MockKernelBackend


# Minimal 1x1 transparent PNG (valid base64)
TINY_PNG_B64 = (
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4"
    "2mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)


def _make_handler(project_root: Path | None = None) -> MessageHandler:
    backend = MockKernelBackend()
    fm = MagicMock()
    fm.project_root = project_root
    fm.get_project_id.return_value = "proj_test" if project_root else None
    handler = MessageHandler(backend, fm)
    handler.set_notification_sender(AsyncMock())
    return handler


# ── _save_figure_from_outputs ──


@pytest.mark.asyncio
async def test_save_figure_from_display_data():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        outputs = [
            {"output_type": "display_data", "data": {"image/png": TINY_PNG_B64}}
        ]
        path = await handler._save_figure_from_outputs("cell-1", outputs)

        assert path is not None
        assert path.startswith(".skop/figures/")
        assert path.endswith(".png")
        # File actually exists
        assert (root / path).exists()
        # Content is valid decoded PNG
        content = (root / path).read_bytes()
        assert content == base64.b64decode(TINY_PNG_B64)


@pytest.mark.asyncio
async def test_save_figure_from_execute_result():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        outputs = [
            {"output_type": "execute_result", "data": {"image/png": TINY_PNG_B64}}
        ]
        path = await handler._save_figure_from_outputs("cell-2", outputs)

        assert path is not None
        assert (root / path).exists()


@pytest.mark.asyncio
async def test_save_figure_no_image():
    with tempfile.TemporaryDirectory() as tmpdir:
        handler = _make_handler(Path(tmpdir))

        outputs = [{"output_type": "stream", "text": "hello"}]
        path = await handler._save_figure_from_outputs("cell-3", outputs)

        assert path is None


@pytest.mark.asyncio
async def test_save_figure_stream_output_ignored():
    with tempfile.TemporaryDirectory() as tmpdir:
        handler = _make_handler(Path(tmpdir))

        outputs = [
            {"output_type": "stream", "data": {"image/png": TINY_PNG_B64}}
        ]
        path = await handler._save_figure_from_outputs("cell-4", outputs)

        assert path is None


@pytest.mark.asyncio
async def test_save_figure_oversized_image():
    with tempfile.TemporaryDirectory() as tmpdir:
        handler = _make_handler(Path(tmpdir))

        big_b64 = "A" * (MAX_FIGURE_BASE64_BYTES + 1)
        outputs = [
            {"output_type": "display_data", "data": {"image/png": big_b64}}
        ]
        path = await handler._save_figure_from_outputs("cell-big", outputs)

        assert path is None


@pytest.mark.asyncio
async def test_save_figure_creates_directory():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        figures_dir = root / ".skop" / "figures"
        assert not figures_dir.exists()

        outputs = [
            {"output_type": "display_data", "data": {"image/png": TINY_PNG_B64}}
        ]
        path = await handler._save_figure_from_outputs("cell-mkdir", outputs)

        assert path is not None
        assert figures_dir.exists()


@pytest.mark.asyncio
async def test_save_figure_sanitizes_cell_id():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        outputs = [
            {"output_type": "display_data", "data": {"image/png": TINY_PNG_B64}}
        ]
        path = await handler._save_figure_from_outputs(
            "../../../etc/passwd", outputs
        )

        assert path is not None
        # Filename should not contain path separators or traversal sequences
        filename = path.split("/")[-1]
        assert "/" not in filename
        assert ".." not in filename
        # Original path separators should be stripped
        assert "../../" not in path


@pytest.mark.asyncio
async def test_save_figure_no_project_root():
    handler = _make_handler(None)

    outputs = [
        {"output_type": "display_data", "data": {"image/png": TINY_PNG_B64}}
    ]
    path = await handler._save_figure_from_outputs("cell-noroot", outputs)

    assert path is None


# ── context notification integration ──


@pytest.mark.asyncio
async def test_context_notification_includes_figure_path():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        result = {
            "outputs": [
                {"output_type": "display_data", "data": {"image/png": TINY_PNG_B64}}
            ],
            "status": "ok",
        }

        session = MagicMock()
        session.notebook_path = "test.ipynb"

        await handler._send_context_notification("cell-notify", "plt.show()", result, session)

        handler._send_notification.assert_called_once()
        call_args = handler._send_notification.call_args
        params = call_args[0][1]
        assert params["figure_path"] is not None
        assert params["figure_path"].startswith(".skop/figures/")


@pytest.mark.asyncio
async def test_context_notification_null_figure_path():
    with tempfile.TemporaryDirectory() as tmpdir:
        root = Path(tmpdir)
        handler = _make_handler(root)

        result = {
            "outputs": [{"output_type": "stream", "text": "hello"}],
            "status": "ok",
        }

        session = MagicMock()
        session.notebook_path = "test.ipynb"

        await handler._send_context_notification("cell-noimg", "print('hi')", result, session)

        handler._send_notification.assert_called_once()
        call_args = handler._send_notification.call_args
        params = call_args[0][1]
        assert params["figure_path"] is None
