import importlib.resources as pkg_resources
import inspect
import logging
import sys
import traceback
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from time import sleep
from typing import List, Optional

import nbclient
import papermill
import typer
import yaml
from papermill.iorw import NoDatesSafeLoader, read_yaml_file
from taegis_magic.core.log import tracing
from taegis_magic.core.normalizer import TaegisResult
from taegis_magic.core.notebook import generate_report
from taegis_magic.core.service import get_service
from taegis_magic.commands.utils.role_checker import has_role
from taegis_sdk_python import GraphQLService
from taegis_sdk_python.services.notebooks.types import Notebook
from gql.transport.exceptions import TransportQueryError
from typing_extensions import Annotated

log = logging.getLogger(__name__)

app = typer.Typer(help="Taegis Notebook Commands.")

remote = typer.Typer(help="Taegis Remote Notebook Commands.")
app.add_typer(remote, name="remote")


@dataclass
class NotebookResult:
    action: str

@dataclass
class RemoteNotebookResult:
    success: bool = False
    message: Optional[str] = None
    remote_instance_data: Optional[any] = None

class LOG_LEVEL(str, Enum):
    NOTSET = "NOTSET"
    DEBUG = "DEBUG"
    INFO = "INFO"
    WARNING = "WARNING"
    ERROR = "ERROR"
    CRITICAL = "CRITICAL"

class NOTEBOOK_STATUS:
    PENDING = "Pending" 
    IN_SERVICE = "InService"
    STOPPING = "Stopping"
    STOPPED = "Stopped"
    FAILED = "Failed"
    DELETING = "Deleting"
    UPDATING = "Updating"
    # Unknown is not a real status that can be returned from API call, used for helping determine if instance exists. 
    UNKNOWN = "Unknown" 

@app.command()
@tracing
def execute(
    input_notebook: Annotated[
        Path,
        typer.Argument(
            exists=True,
            file_okay=True,
            dir_okay=False,
            writable=True,
            readable=True,
            resolve_path=True,
            help="Input notebook to parameterize and execute.",
        ),
    ],
    output_notebook: Annotated[
        Optional[Path],
        typer.Argument(
            exists=False,
            file_okay=True,
            dir_okay=False,
            writable=True,
            readable=True,
            resolve_path=True,
            help="Output notebook with parameters and results.",
        ),
    ] = None,
    parameter: Annotated[
        Optional[List[str]],
        typer.Option(
            "--parameter",
            "-p",
            help="Parameters to pass to the parameters cell. Use the format PARAMETER_NAME=VALUE.",
        ),
    ] = None,
    parameters_file: Annotated[
        Optional[List[Path]],
        typer.Option(
            "--parameters-file",
            "-f",
            exists=True,
            file_okay=True,
            dir_okay=False,
            writable=False,
            readable=True,
            resolve_path=True,
            help="YAML files to read parameters from. Can be used multiple times.",
        ),
    ] = None,
    parameters_yaml: Annotated[
        Optional[List[str]],
        typer.Option(
            "--parameters-yaml",
            "-y",
            help="YAML strings to read parameters from. Can be used multiple times.",
        ),
    ] = None,
    inject_input_path: Annotated[
        bool,
        typer.Option(
            "--inject-input-path",
            help="Insert the path of the input notebook as PAPERMILL_INPUT_PATH as a notebook parameter.",
        ),
    ] = False,
    inject_output_path: Annotated[
        bool,
        typer.Option(
            "--inject-output-path",
            help="Insert the path of the output notebook as PAPERMILL_OUTPUT_PATH as a notebook parameter.",
        ),
    ] = False,
    inject_paths: Annotated[
        bool,
        typer.Option(
            "--inject-paths",
            help="Insert the paths of input/output notebooks as PAPERMILL_INPUT_PATH/PAPERMILL_OUTPUT_PATH"
            " as notebook parameters.",
        ),
    ] = False,
    engine: Annotated[
        Optional[str],
        typer.Option(
            "--engine",
            help="The execution engine name to use in evaluating the notebook.",
        ),
    ] = None,
    request_save_on_cell_execute: Annotated[
        bool,
        typer.Option(
            "--request-save-on-cell-execute",
            help="Request save notebook after each cell execution",
        ),
    ] = True,
    autosave_cell_every: Annotated[
        int,
        typer.Option(
            "--autosave-cell-every",
            help="How often in seconds to autosave the notebook during long cell executions (0 to disable)",
        ),
    ] = 30,
    prepare_only: Annotated[
        bool,
        typer.Option(
            "--prepare-only/--prepare-execute",
            help="Flag for outputting the notebook without execution, but with parameters applied.",
        ),
    ] = False,
    kernel: Annotated[
        Optional[str],
        typer.Option(
            "--kernel",
            "-k",
            help="Name of kernel to run. Ignores kernel name in the notebook document metadata.",
        ),
    ] = None,
    language: Annotated[
        Optional[str],
        typer.Option(
            "--language",
            "-l",
            help="Language for notebook execution. Ignores language in the notebook document metadata.",
        ),
    ] = None,
    cwd: Annotated[
        Optional[str],
        typer.Option(
            "--cwd",
            help="Working directory to run notebook in.",
        ),
    ] = None,
    progress_bar: Annotated[
        bool,
        typer.Option(
            "--progress-bar/--no-progress-bar",
            help="Flag for turning on the progress bar.",
        ),
    ] = True,
    log_output: Annotated[
        bool,
        typer.Option(
            "--log-output/--no-log-output",
            help="Flag for writing notebook output to the configured logger.",
        ),
    ] = False,
    log_level: Annotated[
        LOG_LEVEL, typer.Option("--log-level", help="Log level to use.")
    ] = LOG_LEVEL.INFO,
    stdout_file: Annotated[
        Optional[typer.FileTextWrite],
        typer.Option(
            mode="w",
            encoding="utf-8",
            help="File to write notebook stdout output to.",
        ),
    ] = None,
    stderr_file: Annotated[
        Optional[typer.FileTextWrite],
        typer.Option(
            mode="w",
            encoding="utf-8",
            help="File to write notebook stderr output to.",
        ),
    ] = None,
    start_timeout: Annotated[
        int,
        typer.Option(
            "--start-timeout",
            help="Time in seconds to wait for the kernel to start.",
        ),
    ] = 60,
    execution_timeout: Annotated[
        int,
        typer.Option(
            "--execution-timeout",
            help="Time in seconds to wait for each cell before failing execution (default: forever)",
        ),
    ] = None,
    report_mode: Annotated[
        bool,
        typer.Option("--report-mode/--no-report-mode", help="Flag for hiding input."),
    ] = False,
    title: Annotated[Optional[str], typer.Option(help="Investigation Title.")] = None,
    tenant: Annotated[Optional[str], typer.Option(help="Taegis Tenant ID.")] = None,
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
):
    """Execute a Jupyter notebook."""
    if not prepare_only:
        service = get_service(tenant_id=tenant, environment=region)
        service.access_token

    input_notebook = input_notebook or "-"

    if output_notebook is None:
        stem = input_notebook.stem
        if region:
            stem = f"{stem}_{region}"
        if tenant:
            stem = f"{stem}_{tenant}"

        stem = f"{stem}{input_notebook.suffix}"

        if stem != input_notebook.stem:
            output_notebook = input_notebook.with_name(stem)

    output_notebook = output_notebook or "-"

    if output_notebook == "-":
        # Save notebook to stdout just once
        request_save_on_cell_execute = False

        # Reduce default log level if we pipe to stdout
        if log_level == LOG_LEVEL.INFO:
            log_level = LOG_LEVEL.ERROR

    logging.basicConfig(level=log_level.value, format="%(message)s")

    parameters = {}
    bad_parameters = []
    if inject_input_path or inject_paths:
        parameters["PAPERMILL_INPUT_PATH"] = str(input_notebook)
    if inject_output_path or inject_paths:
        parameters["PAPERMILL_OUTPUT_PATH"] = str(output_notebook)
    for param in parameter or []:
        try:
            name, value = param.split("=", maxsplit=1)
        except ValueError:
            bad_parameters.append(param)
            continue

        parameters[name] = value
    for files in parameters_file or []:
        parameters.update(read_yaml_file(str(files)) or {})
    for params in parameters_yaml or []:
        parameters.update(yaml.load(params, Loader=NoDatesSafeLoader) or {})

    if title:
        parameters["INVESTIGATION_TITLE"] = title
    if tenant:
        parameters["TENANT_ID"] = tenant
    if region:
        parameters["REGION"] = region

    if bad_parameters:
        print(
            f"Please ensure the following parameters use the PARAMETER_NAME=VALUE format: {bad_parameters}"
        )
        raise typer.Exit()

    if "TAEGIS_MAGIC_NOTEBOOK_FILENAME" not in parameters:
        if output_notebook and output_notebook != "-":
            parameters["TAEGIS_MAGIC_NOTEBOOK_FILENAME"] = str(
                output_notebook.resolve()
            )
        else:
            parameters["TAEGIS_MAGIC_NOTEBOOK_FILENAME"] = str(input_notebook.resolve())

    try:
        papermill.execute_notebook(
            input_path=input_notebook,
            output_path=output_notebook,
            parameters=parameters,
            engine_name=engine,
            request_save_on_cell_execute=request_save_on_cell_execute,
            autosave_cell_every=autosave_cell_every,
            prepare_only=prepare_only,
            kernel_name=kernel,
            language=language,
            progress_bar=progress_bar,
            log_output=log_output,
            stdout_file=stdout_file,
            stderr_file=stderr_file,
            start_timeout=start_timeout,
            report_mode=report_mode,
            cwd=cwd,
            execution_timeout=execution_timeout,
        )
    except nbclient.exceptions.DeadKernelError:
        # Exiting with a special exit code for dead kernels
        traceback.print_exc()
        sys.exit(138)

    """Execute the current notebook."""
    return TaegisResult(
        raw_results=NotebookResult(action="execute"),
        service="notebook",
        tenant_id=None,
        region=None,
        arguments=inspect.currentframe().f_locals,
    )


@app.command(name="generate-report")
@tracing
def generate(input_notebook: Annotated[str, typer.Argument(...)]):
    """Generate a report from the current notebook."""
    report = generate_report(filename=input_notebook)

    return TaegisResult(
        raw_results=NotebookResult(action="generate_report"),
        service="notebook",
        tenant_id=None,
        region=None,
        arguments=inspect.currentframe().f_locals,
    )


@app.command()
@tracing
def create(
    output_notebook: Annotated[
        Path,
        typer.Argument(
            exists=False,
            file_okay=True,
            dir_okay=False,
            writable=True,
            readable=True,
            resolve_path=True,
        ),
    ],
):
    """Create a Jupyter Notebook to run against Taegis."""
    template = pkg_resources.read_text("taegis_magic.templates", "template.ipynb")

    output_notebook.write_text(template)

    return TaegisResult(
        raw_results=NotebookResult(action="create"),
        service="notebook",
        tenant_id=None,
        region=None,
        arguments=inspect.currentframe().f_locals,
    )


# Notebook commands for remote jupyter notebook instances on AWS Sagemaker

@remote.command(name="status")
@tracing
def remote_status(
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="More detailed information.")] = False
):
    """Query the status of remote notebook instance."""
    
    _check_notebook_role(region)
        
    service = get_service(environment=region)
    status = RemoteNotebookResult()
    
    try:
        notebook_info = service.notebooks.query.notebook()   
        status.success = True       
        status.remote_instance_data = notebook_info
            
    except TransportQueryError as e:
        status.message = "A remote notebook instance was not found for the current user."
        if verbose:
            status.message += f" Error: {e}"

    return TaegisResult(
        raw_results=status,
        service="notebook",
        tenant_id=service.tenant_id,
        region=service.environment,
        arguments=inspect.currentframe().f_locals,
    )

@remote.command(name="create")
@tracing
def remote_create(
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
    wait: Annotated[bool, typer.Option("--wait", "-w", help="Wait for instance to become available before command returns.")] = False,
    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="More detailed information.")] = False
):
    """Create a new remote notebook instance."""

    _check_notebook_role(region)

    service = get_service(environment=region)

    result = RemoteNotebookResult()
    
    try:
        result.remote_instance_data = service.notebooks.mutation.create_notebook()
        result.success = True
    except TransportQueryError as e:
        notebook_info = _get_remote_notebook_status(service)
        if not notebook_info.status == NOTEBOOK_STATUS.UNKNOWN:
            result.message = f"The user already has a remote notebook instance with status '{notebook_info.status}'"
            result.remote_instance_data = notebook_info
        else:    
            result.message = f"An error occurred while trying to create a remote notebook instance. {e}"

    if result.success and result.remote_instance_data and wait:
        print("Remote notebook instance is being created...")
        while result.remote_instance_data.status == NOTEBOOK_STATUS.PENDING:
            print(f"Status is currently '{NOTEBOOK_STATUS.PENDING}', waiting for 10s.")
            sleep(10)
            result.remote_instance_data = service.notebooks.query.notebook()
    
    return TaegisResult(
        raw_results=result,
        service="notebook",
        tenant_id=service.tenant_id,
        region=service.environment,
        arguments=inspect.currentframe().f_locals,
    )

@remote.command(name="start")
@tracing
def remote_start(
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
    wait: Annotated[bool, typer.Option("--wait", "-w", help="Wait for instance to become available before command returns.")] = False,
    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="More detailed information.")] = False
):
    """Start the remote notebook instance."""
    
    _check_notebook_role(region)
        
    service = get_service(environment=region)
    result = RemoteNotebookResult()
    try:
        result.remote_instance_data = service.notebooks.mutation.start_notebook()
        result.success = True
    except TransportQueryError as e:
        notebook_info = _get_remote_notebook_status(service)
        if not notebook_info.status == NOTEBOOK_STATUS.UNKNOWN:
            result.message = f"Command can only be executed if status is in ['{NOTEBOOK_STATUS.STOPPED}', '{NOTEBOOK_STATUS.FAILED}']. Current status is '{notebook_info.status}'."
            result.remote_instance_data = notebook_info
        else:
            result.message = "A remote notebook instance was not found for the current user."
        if verbose:
            result.message += f" Error: {e}"

    if result.success and result.remote_instance_data and wait:
        print("Remote notebook instance is starting...")
        while result.remote_instance_data.status == NOTEBOOK_STATUS.PENDING:
            print(f"Status is currently '{NOTEBOOK_STATUS.PENDING}', waiting for 10s.")
            sleep(10)
            result.remote_instance_data = service.notebooks.query.notebook()

    return TaegisResult(
        raw_results=result,
        service="notebook",
        tenant_id=service.tenant_id,
        region=service.environment,
        arguments=inspect.currentframe().f_locals,
    )

@remote.command(name="stop")
@tracing
def remote_stop(
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
    wait: Annotated[bool, typer.Option("--wait", "-w", help="Wait for instance to stop before command returns.")] = False,
    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="More detailed error information.")] = False
):
    """Stop remote notebook instance."""

    _check_notebook_role(region)

    service = get_service(environment=region)
    result = RemoteNotebookResult()
    try:
        result.remote_instance_data = service.notebooks.mutation.shutdown_notebook()
        result.success = True
    except TransportQueryError as e:
        notebook_info = _get_remote_notebook_status(service)
        if not notebook_info.status == NOTEBOOK_STATUS.UNKNOWN:
            result.message = f"Status must be '{NOTEBOOK_STATUS.IN_SERVICE}' to execute command. Current status is: '{notebook_info.status}'"
            result.remote_instance_data = notebook_info
        else:
            result.message = "A remote notebook instance was not found for the current user."
        if verbose:
                result.message += f" Error: {e}"
    
    if result.success and result.remote_instance_data and wait:
        print("Remote notebook instance is being stopped...")
        while result.remote_instance_data.status == NOTEBOOK_STATUS.STOPPING:
            print(f"Status is currently '{NOTEBOOK_STATUS.STOPPING}', waiting for 10s.")
            sleep(10)
            result.remote_instance_data = service.notebooks.query.notebook()
    
    return TaegisResult(
        raw_results=result,
        service="notebook",
        tenant_id=service.tenant_id,
        region=service.environment,
        arguments=inspect.currentframe().f_locals,
    )

@remote.command(name="delete")
@tracing
def remote_delete(
    region: Annotated[Optional[str], typer.Option(help="Taegis Region.")] = None,
    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="More detailed information.")] = False
):
    """Delete remote notebook instance."""

    _check_notebook_role(region)

    service = get_service(environment=region)

    result = RemoteNotebookResult()
    try:
        result.remote_instance_data = service.notebooks.mutation.delete_notebook()
        result.success = True
    except TransportQueryError as e:
        notebook_info = _get_remote_notebook_status(service)
        if not notebook_info.status == NOTEBOOK_STATUS.UNKNOWN:
            result.message = f"Status must be in ['{NOTEBOOK_STATUS.STOPPED}', '{NOTEBOOK_STATUS.FAILED}'] to execute command. Current status is: '{notebook_info.status}'"    
            result.remote_instance_data = notebook_info
        else:
            result.message = "A remote notebook instance was not found for the current user."
        if verbose:
                result.message += f" Error: {e}"
    
    return TaegisResult(
        raw_results=result,
        service="notebook",
        tenant_id=service.tenant_id,
        region=service.environment,
        arguments=inspect.currentframe().f_locals,
    )

def _check_notebook_role(region: str):
    if not has_role("notebooks", region):
        print("Current user does not have the notebook role required to access remote notebooks.")
        raise typer.Exit()
    
def _get_remote_notebook_status(service: GraphQLService) -> Notebook:
    notebook = Notebook(status=NOTEBOOK_STATUS.UNKNOWN)
    try:
        notebook = service.notebooks.query.notebook()
    except Exception:
        None
    return notebook