"""Tests for context.executionCompleted notification from kernel.execute."""

import json
import pytest
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock

from skop_runner.message_handler import MessageHandler
from skop_runner.file_manager import FileManager

# Reuse MockKernelBackend from test_message_handler
from tests.test_message_handler import MockKernelBackend


@pytest.fixture
def project_with_skop(tmp_path):
    """Create a temp project dir with .skop/project.json."""
    skop_dir = tmp_path / ".skop"
    skop_dir.mkdir()
    (skop_dir / "project.json").write_text(
        json.dumps({"projectId": "proj_test_ctx"})
    )
    return tmp_path


@pytest.fixture
def handler_with_notification(project_with_skop):
    """MessageHandler with send_notification wired up."""
    backend = MockKernelBackend()
    fm = FileManager(str(project_with_skop))
    notifications = []

    async def capture_notification(method, params):
        notifications.append({"method": method, "params": params})

    handler = MessageHandler(backend, fm)
    handler.set_notification_sender(capture_notification)
    return handler, notifications


class TestContextExecutionCompleted:
    """Tests for context.executionCompleted notification."""

    @pytest.mark.asyncio
    async def test_notification_sent_with_cell_id(self, handler_with_notification):
        """kernel.execute with cell_id and project_id → notification sent."""
        handler, notifications = handler_with_notification

        # Start kernel
        resp = await handler.handle_message({
            "id": "1",
            "method": "kernel.start",
            "params": {},
        })
        kernel_id = resp["result"]["kernel_id"]

        try:
            # Execute with cell_id
            resp = await handler.handle_message({
                "id": "2",
                "method": "kernel.execute",
                "params": {
                    "kernel_id": kernel_id,
                    "code": "print('hello')",
                    "cell_id": "cell_abc",
                },
            })

            assert "result" in resp

            # Find the context.executionCompleted notification
            ctx_notifs = [
                n for n in notifications
                if n["method"] == "context.executionCompleted"
            ]
            assert len(ctx_notifs) == 1
            params = ctx_notifs[0]["params"]
            assert params["project_id"] == "proj_test_ctx"
            assert params["cell_id"] == "cell_abc"
            assert params["code"] == "print('hello')"
            assert params["source"] == "user"
            assert isinstance(params["outputs"], list)
            assert "executed_at" in params
            # Validate ISO 8601 format
            from datetime import datetime, timezone
            dt = datetime.fromisoformat(params["executed_at"])
            assert dt.tzinfo is not None  # must be timezone-aware
        finally:
            await handler.handle_message({
                "id": "99",
                "method": "kernel.stop",
                "params": {"kernel_id": kernel_id},
            })

    @pytest.mark.asyncio
    async def test_no_notification_without_cell_id(self, handler_with_notification):
        """kernel.execute without cell_id → no context notification."""
        handler, notifications = handler_with_notification

        resp = await handler.handle_message({
            "id": "1",
            "method": "kernel.start",
            "params": {},
        })
        kernel_id = resp["result"]["kernel_id"]

        try:
            resp = await handler.handle_message({
                "id": "2",
                "method": "kernel.execute",
                "params": {
                    "kernel_id": kernel_id,
                    "code": "1+1",
                    # no cell_id
                },
            })

            assert "result" in resp
            ctx_notifs = [
                n for n in notifications
                if n["method"] == "context.executionCompleted"
            ]
            assert len(ctx_notifs) == 0
        finally:
            await handler.handle_message({
                "id": "99",
                "method": "kernel.stop",
                "params": {"kernel_id": kernel_id},
            })

    @pytest.mark.asyncio
    async def test_no_notification_without_project_id(self, tmp_path):
        """kernel.execute with cell_id but no .skop/project.json → no notification."""
        backend = MockKernelBackend()
        fm = FileManager(str(tmp_path))
        notifications = []

        async def capture(method, params):
            notifications.append({"method": method, "params": params})

        handler = MessageHandler(backend, fm)
        handler.set_notification_sender(capture)

        resp = await handler.handle_message({
            "id": "1",
            "method": "kernel.start",
            "params": {},
        })
        kernel_id = resp["result"]["kernel_id"]

        try:
            resp = await handler.handle_message({
                "id": "2",
                "method": "kernel.execute",
                "params": {
                    "kernel_id": kernel_id,
                    "code": "print('hi')",
                    "cell_id": "cell_xyz",
                },
            })

            assert "result" in resp
            ctx_notifs = [
                n for n in notifications
                if n["method"] == "context.executionCompleted"
            ]
            assert len(ctx_notifs) == 0
        finally:
            await handler.handle_message({
                "id": "99",
                "method": "kernel.stop",
                "params": {"kernel_id": kernel_id},
            })

    @pytest.mark.asyncio
    async def test_batch_sends_context_notifications(self, handler_with_notification):
        """kernel.executeBatch sends context.executionCompleted for each cell."""
        handler, notifications = handler_with_notification

        resp = await handler.handle_message({
            "id": "1",
            "method": "kernel.start",
            "params": {},
        })
        kernel_id = resp["result"]["kernel_id"]

        try:
            resp = await handler.handle_message({
                "id": "2",
                "method": "kernel.executeBatch",
                "params": {
                    "kernel_id": kernel_id,
                    "cells": [
                        {"cell_id": "cell_1", "code": "print('one')"},
                        {"cell_id": "cell_2", "code": "print('two')"},
                    ],
                },
            })

            assert "result" in resp
            ctx_notifs = [
                n for n in notifications
                if n["method"] == "context.executionCompleted"
            ]
            assert len(ctx_notifs) == 2
            assert ctx_notifs[0]["params"]["cell_id"] == "cell_1"
            assert ctx_notifs[1]["params"]["cell_id"] == "cell_2"
        finally:
            await handler.handle_message({
                "id": "99",
                "method": "kernel.stop",
                "params": {"kernel_id": kernel_id},
            })

    @pytest.mark.asyncio
    async def test_code_truncated_in_notification(self, handler_with_notification):
        """Large code is truncated to 5000 chars in notification."""
        handler, notifications = handler_with_notification

        resp = await handler.handle_message({
            "id": "1",
            "method": "kernel.start",
            "params": {},
        })
        kernel_id = resp["result"]["kernel_id"]

        try:
            long_code = "x = 1\n" * 2000  # ~12000 chars
            resp = await handler.handle_message({
                "id": "2",
                "method": "kernel.execute",
                "params": {
                    "kernel_id": kernel_id,
                    "code": long_code,
                    "cell_id": "cell_long",
                },
            })

            assert "result" in resp
            ctx_notifs = [
                n for n in notifications
                if n["method"] == "context.executionCompleted"
            ]
            assert len(ctx_notifs) == 1
            assert len(ctx_notifs[0]["params"]["code"]) <= 5000
        finally:
            await handler.handle_message({
                "id": "99",
                "method": "kernel.stop",
                "params": {"kernel_id": kernel_id},
            })
