from __future__ import annotations

import os
from pathlib import Path
import socket

import nbformat
from nbformat.validator import normalize
from nbclient import NotebookClient
import pytest

from quantum_learning import load_notebook_execution_plan, project_root


def _local_kernel_ports_available() -> bool:
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.bind(("127.0.0.1", 0))
        return True
    except PermissionError:
        return False


def _prepare_runtime_dirs() -> None:
    root = project_root() / ".tmp_test_artifacts" / "notebook_execution"
    ipython_dir = root / "ipython"
    mpl_dir = root / "matplotlib"
    ipython_dir.mkdir(parents=True, exist_ok=True)
    mpl_dir.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("IPYTHONDIR", str(ipython_dir))
    os.environ.setdefault("MPLCONFIGDIR", str(mpl_dir))
    os.environ.setdefault("MPLBACKEND", "Agg")


def _load_notebook(path: Path):
    with path.open() as handle:
        notebook = nbformat.read(handle, as_version=4)
    normalize(notebook)
    return notebook


@pytest.mark.notebook
@pytest.mark.parametrize(
    "target",
    load_notebook_execution_plan().targets,
    ids=lambda target: target.identifier,
)
def test_execution_plan_notebooks_run_without_errors(target):
    if not _local_kernel_ports_available():
        pytest.skip("Notebook execution needs localhost kernel ports, which this sandbox blocks.")

    _prepare_runtime_dirs()
    notebook_path = project_root() / target.path
    notebook = _load_notebook(notebook_path)
    client = NotebookClient(
        notebook,
        kernel_name="python3",
        timeout=target.timeout_seconds,
        resources={"metadata": {"path": str(notebook_path.parent)}},
    )

    executed = client.execute()
    error_outputs = [
        output
        for cell in executed.cells
        for output in cell.get("outputs", [])
        if output.get("output_type") == "error"
    ]

    assert notebook_path.exists()
    assert not error_outputs