# Copyright (c) Microsoft. All rights reserved.
import os

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
from pytest import mark, param
from traitlets.config import Config

c = Config()

c.RegexRemovePreprocessor.patterns = ["^!pip .*"]
c.ExecutePreprocessor.exclude_input_prompt = True

# These environment variable names are used to control which samples are run during integration testing.
# This has to do with the setup of the tests and the services they depend on.
COMPLETIONS_CONCEPT_SAMPLE = "COMPLETIONS_CONCEPT_SAMPLE"
MEMORY_CONCEPT_SAMPLE = "MEMORY_CONCEPT_SAMPLE"


def run_notebook(notebook_name: str):
    with open(f"samples/getting_started/{notebook_name}") as f:
        nb = nbformat.read(f, as_version=4)
    ep = ExecutePreprocessor(timeout=600, kernel_name="python3", config=c)
    ep.preprocess(nb, {"metadata": {"path": "samples/getting_started/"}})


notebooks = [
    param(
        "00-getting-started.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "01-basic-loading-the-kernel.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "02-running-prompts-from-file.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "03-prompt-function-inline.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "04-kernel-arguments-chat.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "05-memory-and-embeddings.ipynb",
        marks=mark.skipif(
            True, reason="Issue with missing property. Need to investigate and fix. Skip to unblock CI/CD pipeline."
        ),
    ),
    param(
        "06-hugging-face-for-plugins.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "07-native-function-inline.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "08-groundedness-checking.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "09-multiple-results-per-prompt.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
    param(
        "10-streaming-completions.ipynb",
        marks=mark.skipif(
            os.getenv(COMPLETIONS_CONCEPT_SAMPLE, None) is None, reason="Not running completion samples."
        ),
    ),
]


@mark.parametrize("name", notebooks)
def test_notebooks(name):
    run_notebook(name)