#!/usr/bin/env python
#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import atexit
import contextlib
import glob
import json
import logging
import os
import sys
import time
from enum import Enum
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, TypeVar, cast
from urllib.parse import urlparse, urlunparse
from uuid import UUID

import msal
import requests
from azure.storage.blob import ContainerClient
from pydantic import BaseModel
from tenacity import Future as tenacity_future
from tenacity import Retrying, retry
from tenacity.retry import retry_if_exception_type
from tenacity.stop import stop_after_attempt
from tenacity.wait import wait_random

_ACCESSTOKENCACHE_UMASK = 0o077

DEFAULT_CONFIG_PATH = os.path.join("~", ".cache", "onefuzz", "config.json")
DEFAULT_TOKEN_PATH = os.path.join("~", ".cache", "onefuzz", "access_token.json")

LOGGER = logging.getLogger("nsv-backend")


@contextlib.contextmanager
def _temporary_umask(new_umask: int) -> Generator[None, None, None]:
    prev_umask = None
    try:
        prev_umask = os.umask(new_umask)
        yield
    finally:
        if prev_umask is not None:
            os.umask(prev_umask)


class Backend:
    def __init__(
        self,
        config: Optional[Dict[str, str]] = None,
        config_path: Optional[str] = None,
        token_path: Optional[str] = None,
    ):
        self.config_path = os.path.expanduser(config_path or DEFAULT_CONFIG_PATH)
        self.token_path = os.path.expanduser(token_path or DEFAULT_TOKEN_PATH)
        self.config = config or {}
        self.token_cache: Optional[msal.SerializableTokenCache] = None
        self.init_cache()
        self.app: Optional[Any] = None
        self.token_expires = 0
        self.load_config()
        self.session = requests.Session()

        atexit.register(self.save_cache)

    def load_config(self) -> None:
        if os.path.exists(self.config_path):
            with open(self.config_path, "r") as handle:
                self.config.update(json.load(handle))

    def save_config(self) -> None:
        with open(self.config_path, "w") as handle:
            json.dump(self.config, handle)

    def init_cache(self) -> None:
        # Ensure the token_path directory exists
        try:
            dir_name = os.path.dirname(self.token_path)
            with _temporary_umask(_ACCESSTOKENCACHE_UMASK):
                os.makedirs(dir_name)
        except FileExistsError:
            pass

        self.token_cache = msal.SerializableTokenCache()
        if os.path.exists(self.token_path):
            with open(self.token_path, "r") as handle:
                self.token_cache.deserialize(handle.read())

    def save_cache(self) -> None:
        if self.token_cache is None:
            return

        if self.token_path is None:
            return

        with _temporary_umask(_ACCESSTOKENCACHE_UMASK):
            with open(self.token_path, "w") as handle:
                handle.write(self.token_cache.serialize())

    def logout(self) -> None:
        self.app = None
        self.token_cache = None
        if os.path.exists(self.token_path):
            os.unlink(self.token_path)

    def headers(self) -> Dict[str, str]:
        value = {}
        if self.config["client_id"] is not None:
            access_token = self.get_access_token()
            value["Authorization"] = "%s %s" % (
                access_token["token_type"],
                access_token["access_token"],
            )
        return value

    def get_access_token(self) -> Any:
        scopes = [self.config["endpoint"] + "/.default"]

        if "client_secret" in self.config:
            return self.client_secret(scopes)
        return self.device_login(scopes)

    def client_secret(self, scopes: List[str]) -> Any:
        if not self.app:
            self.app = msal.ConfidentialClientApplication(
                self.config["client_id"],
                authority=self.config["authority"],
                client_credential=self.config["client_secret"],
                token_cache=self.token_cache,
            )
        result = self.app.acquire_token_for_client(scopes=scopes)
        if "error" in result:
            raise Exception(
                "error: %s\n'%s'"
                % (result.get("error"), result.get("error_description"))
            )
        return result

    def device_login(self, scopes: List[str]) -> Any:
        if not self.app:
            self.app = msal.PublicClientApplication(
                self.config["client_id"],
                authority=self.config["authority"],
                token_cache=self.token_cache,
            )

        accounts = self.app.get_accounts()
        if accounts:
            access_token = self.app.acquire_token_silent(scopes, account=accounts[0])
            if access_token:
                return access_token

        def check_msal_error(value: Dict[str, Any], expected: List[str]) -> None:
            if "error" in value:
                if "error_description" in value:
                    raise Exception(
                        "error: %s\n%s" % (value["error"], value["error_description"])
                    )

                raise Exception("error: %s" % (value["error"]))
            for entry in expected:
                if entry not in value:
                    raise Exception(
                        "interactive login missing value: %s - %s" % (entry, value)
                    )

        LOGGER.info("Attempting interactive device login")
        print("Please login", flush=True)

        flow = self.app.initiate_device_flow(scopes=scopes)
        check_msal_error(flow, ["user_code", "message"])
        print(flow["message"], flush=True)

        access_token = self.app.acquire_token_by_device_flow(flow)
        check_msal_error(access_token, ["access_token"])

        LOGGER.info("Interactive device authentication succeeded")
        print("Login succeeded", flush=True)
        return access_token

    def request(
        self,
        method: str,
        path: str,
        json_data: Optional[Any] = None,
        params: Optional[Any] = None,
        _retry_on_auth_failure: bool = True,
    ) -> Any:
        if not self.config["endpoint"]:
            raise Exception("endpoint not configured")
        url = self.config["endpoint"] + "/api/" + path
        headers = self.headers()
        json_data = serialize(json_data)

        # 401 errors with IDX10501: Signature validation failed occur
        # on rolling new oauth2 client secrets

        # 404 errors happen when new revisions of the functions code are rolling out
        # to the app service environment.
        # TODO: remove this once swapping deployment are in use

        # 502, 503, and 504 errors are often to Azure App issues.
        retry_codes = [401, 404, 429, 502, 503, 504]

        response = None
        for backoff in range(1, 10):
            try:
                LOGGER.debug("request %s %s %s", method, url, repr(json_data))
                response = self.session.request(
                    method,
                    url,
                    headers=headers,
                    json=json_data,
                    params=params,
                    timeout=(30.0, 30.0),
                )

                if response.status_code not in retry_codes:
                    break

                LOGGER.info("request bad status code: %s", response.status_code)
            except requests.exceptions.ConnectionError as err:
                LOGGER.info("request connection error: %s", err)
            except requests.exceptions.ReadTimeout as err:
                LOGGER.info("request timed out: %s", err)

            time.sleep(1.5 ** backoff)

        if response is None:
            raise Exception("request failed: %s %s" % (method, url))

        if response.status_code / 100 != 2:
            error_text = str(
                response.content, encoding="utf-8", errors="backslashreplace"
            )
            raise Exception(
                "request did not succeed: HTTP %s - %s"
                % (response.status_code, error_text)
            )
        return response.json()


def before_sleep(
    retry_object: Retrying, sleep: float, last_result: tenacity_future
) -> None:
    name = retry_object.fn.__name__ if retry_object.fn else "blob function"

    why = getattr(last_result, "_exception")
    if why:
        LOGGER.warning("%s failed with %s, retrying ...", name, repr(why))
    else:
        LOGGER.warning("%s failed, retrying ...", name)


class ContainerWrapper:
    def __init__(self, container_url: str) -> None:
        self.client = ContainerClient.from_container_url(container_url)

    @retry(
        stop=stop_after_attempt(10),
        wait=wait_random(min=1, max=3),
        retry=retry_if_exception_type(),
        before_sleep=before_sleep,
        reraise=True,
    )
    def upload_file(self, file_path: str, blob_name: str) -> None:
        with open(file_path, "rb") as handle:
            self.client.upload_blob(
                name=blob_name, data=handle, overwrite=True, max_concurrency=10
            )
        return None

    def upload_file_data(self, data: str, blob_name: str) -> None:
        self.client.upload_blob(
            name=blob_name, data=data, overwrite=True, max_concurrency=10
        )

    def upload_dir(self, dir_path: str, recursive: bool = True) -> None:
        for path in glob.glob(os.path.join(dir_path, "**"), recursive=recursive):
            if os.path.isfile(path):
                blob_name = os.path.relpath(path, start=dir_path)
                self.upload_file(path, blob_name)

    @retry(
        stop=stop_after_attempt(10),
        wait=wait_random(min=1, max=3),
        retry=retry_if_exception_type(),
        before_sleep=before_sleep,
        reraise=True,
    )
    def delete_blob(self, blob_name: str) -> None:
        self.client.delete_blob(blob_name)
        return None

    @retry(
        stop=stop_after_attempt(10),
        wait=wait_random(min=1, max=3),
        retry=retry_if_exception_type(),
        before_sleep=before_sleep,
        reraise=True,
    )
    def download_blob(self, blob_name: str) -> bytes:
        return cast(bytes, self.client.download_blob(blob_name).content_as_bytes())

    @retry(
        stop=stop_after_attempt(10),
        wait=wait_random(min=1, max=3),
        retry=retry_if_exception_type(),
        before_sleep=before_sleep,
        reraise=True,
    )
    def list_blobs(self, *, name_starts_with: Optional[str] = None) -> List[str]:
        result = [
            x.name for x in self.client.list_blobs(name_starts_with=name_starts_with)
        ]
        return cast(List[str], result)


def container_file_path(container_url: str, blob_name: str) -> str:
    scheme, netloc, path, params, query, fragment = urlparse(container_url)

    blob_url = urlunparse(
        (scheme, netloc, path + "/" + blob_name, params, query, fragment)
    )

    return blob_url


def serialize(data: Any) -> Any:
    if data is None:
        return data
    if isinstance(data, BaseModel):
        return {serialize(a): serialize(b) for (a, b) in data.dict().items()}
    if isinstance(data, dict):
        return {serialize(a): serialize(b) for (a, b) in data.items()}
    if isinstance(data, list):
        return [serialize(x) for x in data]
    if isinstance(data, tuple):
        return tuple([serialize(x) for x in data])
    if isinstance(data, Enum):
        return data.name
    if isinstance(data, UUID):
        return str(data)
    if isinstance(data, (int, str)):
        return data

    raise Exception("unknown type %s" % type(data))


A = TypeVar("A")


def wait(func: Callable[[], Tuple[bool, str, A]], frequency: float = 1.0) -> A:
    """
    Wait until the provided func returns True

    Provides user feedback via a spinner if stdout is a TTY.
    """

    isatty = sys.stdout.isatty()
    frames = ["-", "\\", "|", "/"]
    waited = False
    last_message = None
    result = None

    try:
        while True:
            result = func()
            if result[0]:
                break
            message = result[1]

            if isatty:
                if last_message:
                    if last_message == message:
                        sys.stdout.write("\b" * (len(last_message) + 2))
                    else:
                        sys.stdout.write("\n")
                sys.stdout.write("%s %s" % (frames[0], message))
                sys.stdout.flush()
            elif last_message != message:
                print(message, flush=True)

            last_message = message
            waited = True
            time.sleep(frequency)
            frames.sort(key=frames[0].__eq__)
    finally:
        if waited and isatty:
            print(flush=True)

    return result[2]
