import json
import os
import tornado
import requests
import datetime
import jwt
from multiprocessing.pool import ThreadPool

from aext_shared.handler import BackendHandler, create_rules
from typing import Dict, Type, Optional
from jupyter_server.base.handlers import APIHandler
from tornado.web import HTTPError
from tornado.routing import _RuleList

# TODO: This can be removed if/when we add anaconda-cloud-auth as a required dependency
try:
    from anaconda_cloud_auth.token import TokenInfo
    from anaconda_cloud_auth.config import AuthConfig
    from anaconda_cloud_auth.exceptions import TokenNotFoundError
    from anaconda_cloud_auth import login, logout
    from anaconda_cloud_auth.handlers import shutdown_all_servers
except ImportError:  # pragma: no cover
    TokenInfo = None
    AuthConfig = None
    TokenNotFoundError = None
    login = None
    logout = None

_cached_api_key = None


def _load_api_key() -> Optional[str]:
    if TokenInfo is None:
        return None
    global _cached_api_key
    if _cached_api_key:
        return _cached_api_key
    config = AuthConfig()
    try:
        _cached_api_key = TokenInfo.load(domain=config.domain).api_key
        return _cached_api_key
    except TokenNotFoundError:
        return None


def _clear_api_key():
    global _cached_api_key
    _cached_api_key = None


def raise_if_error(response):
    if not response.ok:
        raise HTTPError(response.status_code, reason=response.reason)


async def getData(get_URL, **kwargs):
    response = requests.get(get_URL, **kwargs)
    return {"remote_status_code": response.status_code, "remote_data": response.json()}


def get_expires_at(token):
    if not token:
        print("[Assistant] No token found")
        return 0
        # raise HTTPError(403, reason="missing nucleus_token")
    info = jwt.decode(token, algorithms=["RS256"], options={"verify_signature": False})
    expires = datetime.datetime.fromtimestamp(info["exp"])
    # Return milliseconds
    return int(expires.timestamp() * 1000)


class AE5NucleusTokenRouteHandler(APIHandler):
    async def get(self, matched_part=None, *args, **kwargs):
        # To test this, run the following in Python the terminal:
        #     from anaconda_cloud_auth import login
        #     login()
        #     from anaconda_cloud_auth.token import TokenInfo
        #     TokenInfo.load('id.anaconda.cloud').api_key
        # Then place the token at the end of the next line:

        nucleus_token = os.environ.get(
            "ANACONDA_AE5_CLOUD_TOKEN"
        )  # or "your-token-here"

        if not nucleus_token:
            raise HTTPError(403, reason="missing nucleus_token")

        expires_at = get_expires_at(nucleus_token)

        self.finish(
            json.dumps(
                {
                    "access_token": nucleus_token,
                    # 12 months from now, expressed in multiples of 60 seconds
                    "expires_at": expires_at,
                }
            )
        )


class NucleusTokenRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self):
        api_key = _load_api_key()

        if not api_key:
            raise HTTPError(403, reason="missing nucleus_token")

        expires_at = get_expires_at(api_key)
        self.finish(
            {
                "access_token": api_key,
                "expires_at": expires_at,
            }
        )


class LoginRouteHandler(APIHandler):
    pool = None

    @tornado.web.authenticated
    async def get(self):
        """
        This is a workaround for the fact that the login() function in anaconda-cloud-auth
        is not async, and it's not possible to run it in a thread because it starts a tornado
        server. So we run it in a thread pool.

        The next time the user makes a request, if the previous login was still in progress,
        we close the thread pool and start a new one, ensuring that the login will be run again
        without waiting for the previous one and without blocking the main thread.
        """
        try:
            api_key = _load_api_key()
            if api_key:
                self.finish()
                return
        except Exception as e:
            # Ignore errors here, we'll try to login again
            print(f"Error occurred: {e}")

        if LoginRouteHandler.pool:
            LoginRouteHandler.pool.close()

        def login_and_load():
            shutdown_all_servers()
            login()
            _load_api_key()

        with ThreadPool(processes=1) as _pool:
            LoginRouteHandler.pool = _pool
            r = LoginRouteHandler.pool.apply_async(login_and_load)
            while not r.ready():
                await tornado.gen.sleep(0.25)
            r.get()

        self.finish()


class LogoutRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self):
        _clear_api_key()
        logout()
        filepath = prepare_assistant_settings()
        with open(filepath, "w", os.O_EXCL) as f:
            f.write("{}")
        self.finish()


class NucleusUserRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self, matched_part=None, *args, **kwargs):
        nucleus_token = self.request.headers.get("X-Nucleus-Token")
        if not nucleus_token:
            nucleus_token = _load_api_key()
        if not nucleus_token:
            raise HTTPError(403, reason="missing nucleus_token")

        try:
            cloud_api_url = os.environ.get(
                "ANACONDA_CLOUD_API_URL", "https://anaconda.cloud/api"
            )
            self.finish(
                json.dumps(
                    await getData(
                        f"{cloud_api_url}/account",
                        headers={"Authorization": f"Bearer {nucleus_token}"},
                    )
                )
            )
        except Exception as e:
            print(f"Error occurred: {e}")
            raise HTTPError(500, reason=f"Server Error: {e}")


def prepare_assistant_settings():
    """
    Creates the ~/.anaconda/assistant.json file if it doesn't exist.
    TODO: Use Jupyterlab's settings system instead of this file
    """
    directory = os.path.join(os.path.expanduser("~"), ".anaconda")
    if not os.path.exists(directory):
        os.makedirs(directory)
    filepath = os.path.join(directory, "assistant.json")
    if not os.path.exists(filepath):
        with open(filepath, "w", os.O_EXCL) as f:
            f.write("{}")
    return filepath


class GetDiskStateRouteHandler(APIHandler):
    async def get(self, matched_part=None, *args, **kwargs):
        """
        Reads the application state from ~/.anaconda/assistant.json
        We never parse the file, we just return it as a string.
        """
        try:
            filepath = prepare_assistant_settings()
            with open(filepath) as f:
                self.finish(f.read())
        except FileNotFoundError:
            self.finish("{}")


class SyncDiskStateRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def post(self, matched_part=None, *args, **kwargs):
        """
        Writes the application state to ~/.anaconda/assistant.json
        We never parse the file, we just write the string as-is.
        We use "x" mode to ensure that the file doesn't exist.
        """
        try:
            filepath = prepare_assistant_settings()
            with open(filepath, "w", os.O_EXCL) as f:
                f.write(self.request.body.decode("utf-8"))
            self.finish()
        except Exception as e:
            print(f"Error occurred: {e}")
            raise HTTPError(500, reason=f"Server Error: {e}")


def get_routes(base_url: str) -> _RuleList:
    handlers: Dict[str, Type[BackendHandler]] = {
        "nucleus_user": NucleusUserRouteHandler,
        "nucleus_token": NucleusTokenRouteHandler,
        "ae5_nucleus_token": AE5NucleusTokenRouteHandler,
        "login": LoginRouteHandler,
        "logout": LogoutRouteHandler,
        "get_disk_state": GetDiskStateRouteHandler,
        "sync_disk_state": SyncDiskStateRouteHandler,
    }
    return create_rules(base_url, "aext_assistant_server", handlers)
