from contextlib import contextmanager
from functools import wraps
import arcpy
from version import __version__

@contextmanager
def ArcEnvState():
    """Context manager to reset arcpy environment variables"""
    orig_ws = arcpy.env.workspace
    orig_scratch = arcpy.env.scratchWorkspace
    orig_overwrite = arcpy.env.overwriteOutput
    orig_extent = arcpy.env.extent
    orig_snap = arcpy.env.snapRaster
    orig_cellsize = arcpy.env.cellSize
    orig_parallel = arcpy.env.parallelProcessingFactor
    yield None
    arcpy.env.workspace = orig_ws
    arcpy.env.scratchWorkspace = orig_scratch
    arcpy.env.overwriteOutput = orig_overwrite
    arcpy.env.extent = orig_extent
    arcpy.env.snapRaster = orig_snap
    arcpy.env.cellSize = orig_cellsize
    arcpy.env.parallelProcessingFactor = orig_parallel


def isolate_env(func):
    """Decorator function to reset any altered arcpy environment variables"""

    # Using functools.wraps passes through the name and docstring of the decorated function
    @wraps(func)
    def wrapper(*args, **kwargs):
        with ArcEnvState():
            output = func(*args, **kwargs)
        return output

    return wrapper


@contextmanager
def UseExtension(name):
    """Context manager to use and check back in ArcGIS extensions"""
    if arcpy.CheckExtension(name) == "Available":
        status = arcpy.CheckOutExtension(name)
        yield status
    else:
        raise RuntimeError("{} license not available".format(name))
    arcpy.CheckInExtension(name)


def print_header(tool):
    """Prints header string with the given name and a timestamp"""
    line = 80 * "-"
    message = f"""
WBD QC v{__version__}

Tool: {tool.label}

Description: {tool.description}
"""
    arcpy.AddMessage("\n\n".join([line, message, line]))