from __future__ import annotations

import json
from pathlib import Path

import pytest

from astronote.cli import (
    AnalysisError,
    CliError,
    EntrypointSelectionError,
    ModuleExpansionError,
    analyze_source,
    build_parser,
    choose_entrypoint,
    generate_notebook,
    resolve_cli_args,
)

SOURCE = """
from astronote import notebook_entry

@notebook_entry
def run(alpha: int, beta: str = "x") -> None:
    return None
"""

MULTI_SOURCE = """
from astronote import notebook_entry

@notebook_entry
def run() -> None:
    return None

@notebook_entry
def alternate() -> None:
    return None
"""


def test_cli_help_mentions_parameter_file() -> None:
    help_text = build_parser().format_help()

    assert "--parameter-file" in help_text
    assert "--expand-module" in help_text
    assert "--embed-file" in help_text
    assert "--show-schema" in help_text


def test_resolve_cli_args_reads_pyproject_defaults(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    source = tmp_path / "sample.py"
    source.write_text(SOURCE, encoding="utf-8")
    params = tmp_path / "params.json"
    params.write_text(json.dumps({"alpha": 3}), encoding="utf-8")
    pyproject = tmp_path / "pyproject.toml"
    pyproject.write_text(
        "[tool.astronote]\n"
        'source = "sample.py"\n'
        'parameter_file = "params.json"\n'
        "override = ['beta=\"from_pyproject\"']\n"
        'entrypoint = "run"\n'
        'expand_module = ["sub_mod"]\n'
        'embed_file = ["sub_mod.py"]\n'
        'output = "dist/out.ipynb"\n'
        "show_analysis = true\n"
        "show_schema = true\n",
        encoding="utf-8",
    )
    monkeypatch.chdir(tmp_path)

    args = resolve_cli_args([])

    assert args.source == str(source.resolve())
    assert args.parameter_file == params.resolve()
    assert args.override == ['beta="from_pyproject"']
    assert args.entrypoint == "run"
    assert args.expand_module == ["sub_mod"]
    assert args.embed_file == ["sub_mod.py"]
    assert args.output == (tmp_path / "dist" / "out.ipynb").resolve()
    assert args.show_analysis is True
    assert args.show_schema is True


def test_resolve_cli_args_prefers_cli_values_over_pyproject(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    source = tmp_path / "sample.py"
    source.write_text(SOURCE, encoding="utf-8")
    params = tmp_path / "params.json"
    params.write_text(json.dumps({"alpha": 3}), encoding="utf-8")
    cli_params = tmp_path / "cli.json"
    cli_params.write_text(json.dumps({"alpha": 9}), encoding="utf-8")
    pyproject = tmp_path / "pyproject.toml"
    pyproject.write_text(
        "[tool.astronote]\n"
        'source = "sample.py"\n'
        'parameter_file = "params.json"\n'
        "override = ['beta=\"from_pyproject\"']\n"
        'entrypoint = "run"\n'
        'expand_module = ["sub_mod"]\n'
        'embed_file = ["sub_mod.py"]\n'
        'output = "dist/out.ipynb"\n'
        "show_analysis = true\n"
        "show_schema = true\n",
        encoding="utf-8",
    )
    monkeypatch.chdir(tmp_path)

    args = resolve_cli_args(
        [
            str(source),
            "--parameter-file",
            str(cli_params),
            "--override",
            'beta="from_cli"',
            "--entrypoint",
            "cli_run",
            "--expand-module",
            "other_mod",
            "--embed-file",
            "other_mod.py",
            "--output",
            str(tmp_path / "cli.ipynb"),
            "--show-analysis",
            "--show-schema",
        ]
    )

    assert args.source == str(source)
    assert args.parameter_file == cli_params
    assert args.override == ['beta="from_pyproject"', 'beta="from_cli"']
    assert args.entrypoint == "cli_run"
    assert args.expand_module == ["sub_mod", "other_mod"]
    assert args.embed_file == ["sub_mod.py", "other_mod.py"]
    assert args.output == tmp_path / "cli.ipynb"
    assert args.show_analysis is True
    assert args.show_schema is True


def test_resolve_cli_args_requires_source_when_missing(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.chdir(tmp_path)

    with pytest.raises(CliError, match="determine a source target"):
        resolve_cli_args(["--show-schema"])


def test_resolve_cli_args_reports_missing_source_without_cli_or_pyproject(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.chdir(tmp_path)

    with pytest.raises(CliError, match="determine a source target"):
        resolve_cli_args([])


def test_analyze_source_rejects_non_python_file(tmp_path: Path) -> None:
    target = tmp_path / "sample.txt"
    target.write_text("print('nope')\n", encoding="utf-8")

    with pytest.raises(AnalysisError, match=r"existing \.py source file"):
        analyze_source(str(target))


def test_choose_entrypoint_requires_explicit_name_when_multiple_candidates(
    tmp_path: Path,
) -> None:
    source = tmp_path / "multi.py"
    source.write_text(MULTI_SOURCE, encoding="utf-8")
    analysis = analyze_source(str(source))

    with pytest.raises(EntrypointSelectionError, match="multiple candidates"):
        choose_entrypoint(analysis, None)


def test_generate_notebook_writes_manifest_metadata(tmp_path: Path) -> None:
    source = tmp_path / "sample.py"
    source.write_text(SOURCE, encoding="utf-8")
    params = tmp_path / "params.json"
    params.write_text(json.dumps({"alpha": 7}), encoding="utf-8")
    output = tmp_path / "sample.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(analysis, entrypoint, params, ['beta="override"'], output)

    notebook = json.loads(output.read_text(encoding="utf-8"))
    meta = notebook["metadata"]["astronote"]

    assert meta["manifest"]["source_path"] == str(source.resolve())
    assert meta["manifest"]["entrypoint"] == "run"
    assert meta["manifest"]["parameters"] == {"alpha": 7, "beta": "override"}
    assert meta["manifest"]["parameter_schema"]["fields"][0]["name"] == "alpha"


def test_generate_notebook_embeds_source_definition_without_main_guard(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example.py"
    source.write_text(
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def run() -> str:\n"
        "    return 'ok'\n\n"
        "if __name__ == '__main__':\n"
        "    run()\n",
        encoding="utf-8",
    )
    output = tmp_path / "example.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(analysis, entrypoint, None, [], output)

    notebook = json.loads(output.read_text(encoding="utf-8"))
    roles = [cell["metadata"]["astronote"]["role"] for cell in notebook["cells"]]

    assert roles == ["generated_header", "source_definition", "entrypoint"]
    assert "# Generated by Astronote" in "".join(notebook["cells"][0]["source"])
    assert str(source.resolve()) in "".join(notebook["cells"][0]["source"])
    source_definition = "".join(notebook["cells"][1]["source"])
    assert "notebook_entry" not in source_definition
    assert "if __name__ == '__main__':" not in source_definition
    assert source_definition == "def run() -> str:\n    return 'ok'\n\n"
    assert notebook["cells"][2]["source"] == ["run()\n"]


def test_generate_notebook_strips_entrypoint_alias_imports(
    tmp_path: Path,
) -> None:
    source = tmp_path / "alias_example.py"
    source.write_text(
        "import astronote as an\n\n"
        "@an.notebook_entry()\n"
        "def run() -> str:\n"
        "    return 'ok'\n",
        encoding="utf-8",
    )
    output = tmp_path / "alias_example.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(analysis, entrypoint, None, [], output)

    notebook = json.loads(output.read_text(encoding="utf-8"))
    source_definition = "".join(notebook["cells"][1]["source"])

    assert source_definition == "def run() -> str:\n    return 'ok'\n"


def test_generate_notebook_expands_requested_modules(tmp_path: Path) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.get_hello_world())\n",
        encoding="utf-8",
    )
    module_file = tmp_path / "sub_mod.py"
    module_file.write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello, world!'\n",
        encoding="utf-8",
    )
    output = tmp_path / "example_multiple_file.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        output,
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(output.read_text(encoding="utf-8"))
    roles = [cell["metadata"]["astronote"]["role"] for cell in notebook["cells"]]

    assert roles == [
        "generated_header",
        "source_definition",
        "source_definition",
        "entrypoint",
    ]

    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert f"# Embedded from: {module_file.resolve()}" in module_definition
    assert "def get_hello_world() -> str:" in module_definition
    assert f"# Embedded from: {source.resolve()}" in main_definition
    assert "import sub_mod" not in main_definition
    assert "def main() -> None:" in main_definition
    assert "print(get_hello_world())" in main_definition
    assert "notebook_entry" not in main_definition
    assert notebook["cells"][3]["source"] == ["main()\n"]


def test_generate_notebook_expands_requested_python_file_paths(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.get_hello_world())\n",
        encoding="utf-8",
    )
    module_file = tmp_path / "sub_mod.py"
    module_file.write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello from file path!'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        tmp_path / "example_multiple_file.ipynb",
        embed_files=["sub_mod.py"],
    )

    notebook = json.loads(
        (tmp_path / "example_multiple_file.ipynb").read_text(encoding="utf-8")
    )
    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert f"# Embedded from: {module_file.resolve()}" in module_definition
    assert "Hello from file path!" in module_definition
    assert "import sub_mod" not in main_definition
    assert "print(get_hello_world())" in main_definition


def test_generate_notebook_rewrites_expanded_alias_imports(tmp_path: Path) -> None:
    source = tmp_path / "alias_multiple_file.py"
    source.write_text(
        "import sub_mod as helper\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(helper.get_hello_world())\n",
        encoding="utf-8",
    )
    (tmp_path / "sub_mod.py").write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello through alias!'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        tmp_path / "alias_multiple_file.ipynb",
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(
        (tmp_path / "alias_multiple_file.ipynb").read_text(encoding="utf-8")
    )
    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert (
        f"# Embedded from: {(tmp_path / 'sub_mod.py').resolve()}" in module_definition
    )
    assert "def get_hello_world() -> str:" in module_definition
    assert f"# Embedded from: {source.resolve()}" in main_definition
    assert "import sub_mod as helper" not in main_definition
    assert "def main() -> None:" in main_definition
    assert "print(get_hello_world())" in main_definition


def test_generate_notebook_expands_requested_modules_from_imports(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "from sub_mod import func\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(func())\n",
        encoding="utf-8",
    )
    package_dir = tmp_path / "sub_mod"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text(
        "def func() -> str:\n" "    return 'Called sub function'\n",
        encoding="utf-8",
    )
    output = tmp_path / "example_multiple_file.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        output,
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(output.read_text(encoding="utf-8"))
    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert (
        f"# Embedded from: {(package_dir / '__init__.py').resolve()}"
        in module_definition
    )
    assert "def func() -> str:" in module_definition
    assert "from sub_mod import func" not in main_definition
    assert "print(func())" in main_definition


def test_generate_notebook_recursively_expands_nested_package_modules(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "from sub_mod import func\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(func())\n",
        encoding="utf-8",
    )
    package_dir = tmp_path / "sub_mod"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text(
        "from sub_mod.sub import sub_func\n\n"
        "def func() -> str:\n"
        "    return sub_func() + ' root'\n",
        encoding="utf-8",
    )
    (package_dir / "sub.py").write_text(
        "def sub_func() -> str:\n" "    return 'nested'\n",
        encoding="utf-8",
    )
    output = tmp_path / "example_multiple_file.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        output,
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(output.read_text(encoding="utf-8"))
    roles = [cell["metadata"]["astronote"]["role"] for cell in notebook["cells"]]
    nested_definition = "".join(notebook["cells"][1]["source"])
    package_definition = "".join(notebook["cells"][2]["source"])
    main_definition = "".join(notebook["cells"][3]["source"])

    assert roles == [
        "generated_header",
        "source_definition",
        "source_definition",
        "source_definition",
        "entrypoint",
    ]
    assert f"# Embedded from: {(package_dir / 'sub.py').resolve()}" in nested_definition
    assert "def sub_func() -> str:" in nested_definition
    assert (
        f"# Embedded from: {(package_dir / '__init__.py').resolve()}"
        in package_definition
    )
    assert "from sub_mod.sub import sub_func" not in package_definition
    assert "return sub_func() + ' root'" in package_definition
    assert "from sub_mod import func" not in main_definition
    assert "print(func())" in main_definition


def test_generate_notebook_recursively_expands_relative_nested_package_modules(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "from sub_mod import func\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(func())\n",
        encoding="utf-8",
    )
    package_dir = tmp_path / "sub_mod"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text(
        "from .sub import sub_func\n\n"
        "def func() -> str:\n"
        "    return sub_func()\n",
        encoding="utf-8",
    )
    (package_dir / "sub.py").write_text(
        "def sub_func() -> str:\n" "    return 'nested'\n",
        encoding="utf-8",
    )
    output = tmp_path / "example_multiple_file.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        output,
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(output.read_text(encoding="utf-8"))
    package_definition = "".join(notebook["cells"][2]["source"])

    assert "from .sub import sub_func" not in package_definition
    assert "return sub_func()" in package_definition


def test_generate_notebook_rewrites_expanded_from_import_aliases(
    tmp_path: Path,
) -> None:
    source = tmp_path / "alias_multiple_file.py"
    source.write_text(
        "from sub_mod import func as helper\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(helper())\n",
        encoding="utf-8",
    )
    package_dir = tmp_path / "sub_mod"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text(
        "def func() -> str:\n" "    return 'Called through alias'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        tmp_path / "alias_multiple_file.ipynb",
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(
        (tmp_path / "alias_multiple_file.ipynb").read_text(encoding="utf-8")
    )
    main_definition = "".join(notebook["cells"][2]["source"])

    assert "from sub_mod import func as helper" not in main_definition
    assert "print(helper())" not in main_definition
    assert "print(func())" in main_definition


def test_generate_notebook_expands_requested_package_submodules_from_imports(
    tmp_path: Path,
) -> None:
    package_dir = tmp_path / "pkg"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text("", encoding="utf-8")

    source = tmp_path / "example.py"
    source.write_text(
        "from pkg import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.value())\n",
        encoding="utf-8",
    )
    (package_dir / "sub_mod.py").write_text(
        "def value() -> str:\n" "    return 'Nested module'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        tmp_path / "example.ipynb",
        expand_modules=["pkg.sub_mod"],
    )

    notebook = json.loads((tmp_path / "example.ipynb").read_text(encoding="utf-8"))
    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert (
        f"# Embedded from: {(package_dir / 'sub_mod.py').resolve()}"
        in module_definition
    )
    assert "from pkg import sub_mod" not in main_definition
    assert "print(sub_mod.value())" not in main_definition
    assert "print(value())" in main_definition


def test_generate_notebook_suggests_closest_expand_module_name(tmp_path: Path) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.get_hello_world())\n",
        encoding="utf-8",
    )
    (tmp_path / "sub_mod.py").write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello, world!'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(
        ModuleExpansionError,
        match=r"Detected import targets: sub_mod\. Closest match: sub_mod\. Next step: pass --expand-module sub_mod\.",
    ):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example_multiple_file.ipynb",
            expand_modules=["sub_module"],
        )


def test_generate_notebook_rejects_unimported_python_file_paths(
    tmp_path: Path,
) -> None:
    source = tmp_path / "example_multiple_file.py"
    source.write_text(
        "import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.get_hello_world())\n",
        encoding="utf-8",
    )
    (tmp_path / "sub_mod.py").write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello, world!'\n",
        encoding="utf-8",
    )
    unused_file = tmp_path / "other.py"
    unused_file.write_text(
        "def value() -> str:\n    return 'unused'\n", encoding="utf-8"
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(
        ModuleExpansionError,
        match=r"did not match any imported local module",
    ):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example_multiple_file.ipynb",
            embed_files=[str(unused_file)],
        )


def test_generate_notebook_expands_requested_modules_in_flat_layout(
    tmp_path: Path,
) -> None:
    project_dir = tmp_path / "project"
    project_dir.mkdir()

    source = project_dir / "example_multiple_file.py"
    source.write_text(
        "import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def main() -> None:\n"
        "    print(sub_mod.get_hello_world())\n",
        encoding="utf-8",
    )
    module_file = project_dir / "sub_mod.py"
    module_file.write_text(
        "def get_hello_world() -> str:\n" "    return 'Hello from flat layout!'\n",
        encoding="utf-8",
    )
    output = project_dir / "example_multiple_file.ipynb"

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)
    generate_notebook(
        analysis,
        entrypoint,
        None,
        [],
        output,
        expand_modules=["sub_mod"],
    )

    notebook = json.loads(output.read_text(encoding="utf-8"))
    module_definition = "".join(notebook["cells"][1]["source"])
    main_definition = "".join(notebook["cells"][2]["source"])

    assert f"# Embedded from: {module_file.resolve()}" in module_definition
    assert "Hello from flat layout!" in module_definition
    assert f"# Embedded from: {source.resolve()}" in main_definition
    assert "print(get_hello_world())" in main_definition


def test_generate_notebook_requires_exact_expand_module_match(tmp_path: Path) -> None:
    package_dir = tmp_path / "pkg"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text("", encoding="utf-8")

    source = package_dir / "example.py"
    source.write_text(
        "from . import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def run() -> None:\n"
        "    print(sub_mod.value())\n",
        encoding="utf-8",
    )
    (package_dir / "sub_mod.py").write_text(
        "def value() -> str:\n" "    return 'ok'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(
        ModuleExpansionError,
        match=r"Detected import targets: \.sub_mod",
    ):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example.ipynb",
            expand_modules=["sub_mod"],
        )


def test_generate_notebook_rejects_relative_expand_module_targets(
    tmp_path: Path,
) -> None:
    package_dir = tmp_path / "pkg"
    package_dir.mkdir()
    (package_dir / "__init__.py").write_text("", encoding="utf-8")

    source = package_dir / "example.py"
    source.write_text(
        "from . import sub_mod\n\n"
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def run() -> None:\n"
        "    print(sub_mod.value())\n",
        encoding="utf-8",
    )
    (package_dir / "sub_mod.py").write_text(
        "def value() -> str:\n" "    return 'ok'\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(
        ModuleExpansionError,
        match="relative import targets are not supported for notebook expansion",
    ):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example.ipynb",
            expand_modules=[".sub_mod"],
        )


def test_generate_notebook_rejects_unknown_expand_modules(tmp_path: Path) -> None:
    source = tmp_path / "example.py"
    source.write_text(
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def run() -> None:\n"
        "    return None\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(
        ModuleExpansionError,
        match=r"did not exactly match any imported module string.*Detected import targets: none",
    ):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example.ipynb",
            expand_modules=["missing_module"],
        )


def test_generate_notebook_rejects_path_style_expand_modules(tmp_path: Path) -> None:
    source = tmp_path / "example.py"
    source.write_text(
        "from astronote import notebook_entry\n\n"
        "@notebook_entry\n"
        "def run() -> None:\n"
        "    return None\n",
        encoding="utf-8",
    )

    analysis = analyze_source(str(source))
    entrypoint = choose_entrypoint(analysis, None)

    with pytest.raises(ModuleExpansionError, match="module paths are unsupported"):
        generate_notebook(
            analysis,
            entrypoint,
            None,
            [],
            tmp_path / "example.ipynb",
            expand_modules=["sub_mod=missing.py"],
        )
