from flytekit import kwtypes, workflow, dynamic
from flytekitplugins.papermill import NotebookTask
from flytekit.types.file import HTMLPage

from flytekit import ImageSpec
from src.images import BaseImageSpec, NotebookImageSpec

task1 = NotebookTask(
    name="notebook_task1",
    notebook_path=str(
        pathlib.Path(__file__)
        .parent.absolute()
        .joinpath("notebooks/task1.ipynb")
    ),
    render_deck=True,
    enable_deck=True,
    inputs=kwtypes(
        some_input=str
    ),
    container_image=NotebookImageSpec,
)


task2 = NotebookTask(
    name="notebook_task2",
    notebook_path=str(
        pathlib.Path(__file__)
        .parent.absolute()
        .joinpath("notebooks/task2.ipynb")
    ),
    render_deck=True,
    enable_deck=True,
    inputs=kwtypes(
        some_other_input=str,
    ),
    container_image=NotebookImageSpec,
)


@dynamic(container_image=BaseImageSpec)
def run_dynamic_workflow() -> HTMLPage:
    _, nb2 = task2(some_other_input="some_input")
    return nb2


@workflow
def run_workflow() -> None:
    _, nb1 = task1(some_input="some_input")
    nb2 = run_dynamic_workflow()