import datetime
import json
import os
from multiprocessing.pool import ThreadPool
from typing import Dict, Optional, Type

import jwt
import requests
import tornado
from jupyter_server.base.handlers import APIHandler
from tornado.routing import _RuleList
from tornado.web import HTTPError

from aext_shared.handler import BackendHandler, create_rules

# TODO: This can be removed if/when we add anaconda-cloud-auth as a required dependency
try:
    from anaconda_cloud_auth import login, logout
    from anaconda_cloud_auth.config import AuthConfig
    from anaconda_cloud_auth.exceptions import TokenNotFoundError
    from anaconda_cloud_auth.handlers import shutdown_all_servers
    from anaconda_cloud_auth.token import TokenInfo
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, headers={}):
    # Update headers
    cloud_flare_client = os.getenv("NUCLEUS_CLOUDFLARE_CLIENT_ID")
    cloud_flare_secret = os.getenv("NUCLEUS_CLOUDFLARE_CLIENT_SECRET")
    if cloud_flare_client and cloud_flare_secret:
        headers["CF-Access-Client-Id"] = cloud_flare_client
        headers["CF-Access-Client-Secret"] = cloud_flare_secret

    response = requests.get(get_URL, headers=headers)
    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")
    try:
        info = jwt.decode(token, algorithms=["RS256"], options={"verify_signature": False})
        expires = datetime.datetime.fromtimestamp(info["exp"])
        # Return milliseconds
        return int(expires.timestamp() * 1000)
    except Exception as e:
        print(f"Error occurred: {e}")
        return 0


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}")


# In order to summarize files for the user, we keep track of when they have been updated
files_last_modified = {}


# Time to wait for a file to change
def get_wait_time():
    return int(os.environ.get("ASSISTANT_MONITOR_FILES_WAIT_TIME", 60))


class MonitorFileChangesRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def get(self, matched_part=None, *args, **kwargs):
        """
        Lists all the files `.ipynb` files in the current directory,
        then waits until one of the files change,
        then returns an object with the following properties:
        - path: the path to the file
        - last_modified: the last time the file was modified
        """
        try:
            # Get the current directory, then list all the .ipynb files
            path = os.path.abspath(os.path.join(os.getcwd(), "."))
            files = [f for f in os.listdir(path) if f.endswith(".ipynb")]

            # If there are no files, return an empty object
            if not files:
                self.finish(json.dumps({"path": None, "last_modified": 0}))
                return

            # Wait (up to 1 minute) for a file to change
            current_time = datetime.datetime.now()
            while (datetime.datetime.now() - current_time).seconds < get_wait_time():
                for f in files:
                    last_modified = os.path.getmtime(f)
                    # If the file has not been seen before, save the last modified time and continue to the next file
                    if f not in files_last_modified:
                        files_last_modified[f] = last_modified
                        continue
                    # If the file has been modified, save the last modified time and return the file
                    if last_modified > files_last_modified.get(f, 0):
                        files_last_modified[f] = last_modified
                        self.finish(json.dumps({"path": f, "last_modified": last_modified}))
                        return
                    # In any case, update the last modified time, then continue to the next file.
                    files_last_modified[f] = last_modified
                    continue
                await tornado.gen.sleep(1)
            self.finish(json.dumps({"path": None, "last_modified": 0}))
        except Exception as e:
            print(f"Error occurred: {e}")
            raise HTTPError(500, reason=f"Server Error: {e}")


class SummarizeFileRouteHandler(APIHandler):
    @tornado.web.authenticated
    async def post(self, matched_part=None, *args, **kwargs):
        """
        Receives a file path, then sends it to the Assistant API
        to get the summary of the file by making a request to
        ANACONDA_ASSISTANT_API_URL/v1/pro/summaries
        As the following curl script indicates:
        curl --location 'http://localhost:8000/v1/pro/summaries' --header 'x-client-version: 0.4.0' --header 'x-client-source: anaconda-local-dev' --header 'Content-Type: application/json' --header 'Authorization: Bearer {token}' --data '{
            "source": {
                "name": "chat_123",
                "data": "data",
                "type": "ipynb"
            },
            "session_id": "session_123",
            "response_message_id": "resp_msg_123"
        }'
        Where the token is obtained from _load_api_key()
        """
        try:
            # Load the file from the request body
            json_body = json.loads(self.request.body)
            file_path = json_body.get("file_path")

            # Read the file contents
            with open(file_path) as f:
                file_contents = f.read()

            # Prepare the request to the Assistant API
            assistant_api_request_headers = {
                "x-client-version": "0.4.0",
                "x-client-source": json_body.get("x-client-source", "anaconda-local-dev"),
                "Content-Type": "application/json",
                "Authorization": f"Bearer {_load_api_key()}",
            }
            assistant_api_request_body = {
                "source": {
                    "name": file_path,
                    "data": file_contents,
                    "type": "ipynb",
                },
                "session_id": json_body.get("session_id", "session_123"),
                "response_message_id": json_body.get("response_message_id", "response_123"),
                "skip_logging": json_body.get("skip_logging", False),
            }

            # Send the request to the Assistant API
            response = requests.post(
                os.environ.get("ANACONDA_ASSISTANT_API_URL", "http://localhost:8000") + "/v1/pro/summaries",
                headers=assistant_api_request_headers,
                json=assistant_api_request_body,
            )
            raise_if_error(response)
            self.finish(response.json())
        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,
        "monitor_file_changes": MonitorFileChangesRouteHandler,
        "summarize_file": SummarizeFileRouteHandler,
    }
    return create_rules(base_url, "aext_assistant_server", handlers)
