import io
from typing import TYPE_CHECKING, Any, Dict, List
from prefect.client import Secret


class Webhook(Storage):
    """
    Args:
        - build_kwargs (dict): Dictionary of keyword arguments to the
            function from ``requests`` used to store the flow. 
        - build_http_method (str): HTTP method identifying the type of request
            to execute when storing the flow. For example, ``"POST"`` for
            ``requests.post()``.
        - get_flow_kwargs (dict): Dictionary of keyword arguments to the
            function from ``requests`` used to retrive the flow.
        - get_flow_http_method (str): HTTP method identifying the type of
            request to execute when storing the flow. For example, ``"GET"``
            for ``requests.post()``.
        - secret_config (dict, optional): A dictionary describing how to set
            request headers from environment variables or Prefect Cloud
            secrets. See example for details on specifying this

    Passing sensitive data in headers
    ---------------------------------

    For services which require authentication, use `secret_config` to pass
    sensitive data like API keys without storing their values in this Storage object.

    This should be a dictionary whose keys are headers, and whose
    values indicate whether to retrieve real values from environment
    variables (``"type": "environment"``) or
    Prefect secrets (``"type": "secret"``).

    So, for example, to get an API key from an environment variable you
    can do the following

    .. code-block:: python

        storage = Webhoook(
            build_kwargs={
                "url": "some-random-service.place.thing",
                "headers" = {
                    "Content-Type" = "application/octet-stream"
                }
            },
            build_http_method="POST",
            ...
            ...
            secret_config={
                "X-Api-Key": {
                    "value": "MY_COOL_ENV_VARIABLE",
                    "type": "environment"
                }
            }
        )
    """

    def __init__(
        self,
        build_kwargs: dict,
        build_http_method: str,
        get_flow_kwargs: dict,
        get_flow_http_method: str,
        secret_config: dict = {},
         **kwargs: Any
    ) -> None:
        self.flows = dict()  # type: Dict[str, str]
        self._flows = dict()  # type: Dict[str, "Flow"]

        self.build_kwargs = build_kwargs
        self.build_http_method = build_http_method

        self.get_flow_kwargs = get_flow_kwargs
        self.get_flow_http_method = build_http_method

        self.secret_config = {}

        self._method_to_function = {
            "GET": requests.get,
            "POST": requests.post,
            "PUT": requests.put
        }

        super().__init__(**kwargs)

    @property
    def default_labels(self) -> List[str]:
        return ["webhook-flow-storage"]

    def _render_headers(self, headers: dict) -> dict:
        out_headers = headers.copy()
        for header, _ in headers.items():
            if header in self.secret_config.keys():
                name = self.secret_config[header]
                if self.secret_config["type"] == "environment":
                    out_headers[header] = os.environ[name]
                elif self.secret_config["type"] == "secret":
                    out_headers[header] = Secret(name).get()
        return out_header

    def get_flow(self, flow_name: str) -> "Flow":
        req_function = self._method_to_function[self.get_flow_http_method]

        get_flow_kwargs = self.get_flow_kwargs.copy()
        get_flow_kwargs["headers"] = self._render_headers(
            get_flow_kwargs["headers"]
        )

        response = req_function(**get_flow_kwargs)

        return cloudpickle.loads(response.content)

    def build(self) -> "Storage":
        self.run_basic_healthchecks()

        for flow_name, flow in self._flows.items():
            # Pickle Flow
            data = cloudpickle.dumps(flow)

            # Write pickled Flow to stream
            try:
                stream = io.BytesIO(data)
            except TypeError:
                stream = io.BytesIO(data.encode())

            # write flow to the service
            req_function = self._method_to_function[self.build_method]

            build_kwargs = self.build_kwargs.copy()
            build_kwargs["headers"] = self._render_headers(
                build_kwargs["headers"]
            )
            response = req_function(**build_kwargs)

        return self