import base64
import hashlib
import hmac
import time
import urllib.parse

from airflow.exceptions import AirflowException
from airflow.providers.dingding.hooks.dingding import DingdingHook


class DingdingWithSignHook(DingdingHook):
    """
    DingdingHook with sign supported
    """

    def _get_secret(self) -> str:
        """Get DingTalk Custom Robot Secret"""
        conn = self.get_connection(self.http_conn_id)
        secret = conn.extra_dejson.get('secret', '')
        if not secret:
            raise AirflowException(
                "Dingding secret is requests but get nothing, check you conn_id configuration."
            )
        return secret

    def send(self) -> None:
        """
        Copy from DingdingHook.send, with sign supported
        """
        support_type = ["text", "link", "markdown", "actionCard", "feedCard"]
        if self.message_type not in support_type:
            raise ValueError(
                f"DingdingWebhookHook only support {support_type} so far, but receive {self.message_type}"
            )

        timestamp = str(round(time.time()) * 1000)
        secret = self._get_secret()
        secret_enc = secret.encode('utf-8')
        string_to_sign = f'{timestamp}\n{secret}'
        string_to_sign_enc = string_to_sign.encode('utf-8')
        hmac_code = hmac.new(secret_enc, string_to_sign_enc, digestmod=hashlib.sha256).digest()
        sign = urllib.parse.quote_plus(base64.b64encode(hmac_code))

        endpoint = self._get_endpoint()
        endpoint = f'{endpoint}&timestamp={timestamp}&sign={sign}'

        data = self._build_message()
        self.log.info("Sending Dingding type %s message %s", self.message_type, data)
        resp = self.run(
            endpoint=endpoint, data=data, headers={"Content-Type": "application/json"}
        )

        # Success send message will return errcode = 0
        if int(resp.json().get("errcode")) != 0:
            raise AirflowException(f"Send Dingding message failed, receive error message {resp.text}")
        self.log.info("Success Send Dingding message")
