import logging
import os
from base64 import b64encode, urlsafe_b64encode
from datetime import datetime, timedelta
from hashlib import sha256
from json import dumps, loads
from random import getrandbits
from re import search
from socket import AF_INET, SHUT_WR, SOCK_STREAM, socket
from typing import Optional
from webbrowser import open as browse

import requests
from appdirs import user_cache_dir
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from requests_oauthlib import OAuth2Session

CLIENT_ID = "amcat4py"

FELINE_RESPONSE = """HTTP/1.1 200 OK
Content-Type: text/html; charset=utf-8

<html><body><pre>Authorization complete, you can close this tab and return to Python.
 ,_     _
 |\\\\_,-~/
 / _  _ |    ,--.
(  @  @ )   / ,-'
 \\  _T_/-._( (
 /         `. \\
|         _  \\ |
 \\ \\ ,  /      |
  || |-_\\__   /
 ((_/`(____,-'
</pre></body></html>
"""


def get_middlecat_token(host, callback_port=65432, refresh="static") -> dict:
    """
    Authenticate to an AmCAT instance using a middlecat instance (which is automatically retrieved from the AmCAT instance)
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api")
    :param callback_port: Port used to receive the token. The only reason to change this is if the port is already in use.
    :param refresh: Either "refresh" or "static" to en-/disable token rotation.
    """
    # we open a socket and browser for for the interactive authentication and wait for the code from middlecat
    # We do this first so we can change the port if needed
    s = socket(AF_INET, SOCK_STREAM)
    while True:
        try:
            s.bind(("127.0.0.1", callback_port))
        except OSError:
            logging.info(f"Port {callback_port} already in use, trying {callback_port - 1}")
            callback_port -= 1
        else:
            break
    s.listen()

    middlecat = requests.get(f"{host}/config").json()["middlecat_url"]
    auth_url = f"{middlecat}/authorize"
    token_url = f"{middlecat}/api/token"
    pkce = pkce_challange()

    auth_params = {
        "resource": host,
        "refresh_mode": refresh,
        "session_type": "api_key",
        "code_challenge_method": pkce["method"],
        "code_challenge": pkce["challenge"],
    }

    oauth = OAuth2Session(client_id=CLIENT_ID, redirect_uri=f"http://localhost:{callback_port}/")

    authorization_url, state = oauth.authorization_url(auth_url, **auth_params)
    browse(authorization_url)
    print("Waiting for authorization in browser...")
    conn, addr = s.accept()

    data = conn.recv(4096).decode()
    match = search(r"code=([^&\s]+)", data)
    if match is None:
        raise ValueError(f"No authorization code in callback: {data!r}")
    code = match.group(1)
    conn.sendall(FELINE_RESPONSE.encode("ascii"))
    conn.shutdown(SHUT_WR)
    conn.close()

    # using the received code, make a request to get the actual token
    headers = {"Accept": "application/json", "Content-Type": "application/json"}
    params = {"grant_type": "authorization_code", "code": code, "code_verifier": pkce["verifier"], "state": state}
    return _request_token(params, headers, host, token_url)


def token_refresh(token, host) -> dict:
    """
    Usually called by _check_token if token has expired
    :param token: old token
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api").
    """
    middlecat = requests.get(f"{host}/middlecat").json()["middlecat_url"]
    token_url = f"{middlecat}/api/token"
    params = {
        "resource": host,
        "grant_type": "refresh_token",
        "refresh_mode": token["refresh_rotate"],
        "session_type": "api_key",
        "refresh_token": token["refresh_token"],
        "client_id": CLIENT_ID,
    }
    headers = {"Accept": "application/json", "Content-Type": "application/json"}
    return _request_token(params, headers, host, token_url)


def _request_token(auth_params, headers, host, token_url) -> dict:
    r = requests.post(token_url, headers=headers, data=dumps(auth_params))
    r.raise_for_status()
    token = r.json()
    expires_at = timedelta(seconds=token["expires_in"]) + datetime.now()
    token["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%S")
    del token["expires_in"]
    cache_token(token, host)
    return token


def _get_token(host, force_refresh=False, login_if_needed=True) -> Optional[dict]:
    """
    Returns refreshed token if old token has expired
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api").
    :param username,password: optionally provide a username and password to trigger get_password_token()
        instead of get_middlecat_token()
    :param force_refresh: when True, overwrites the cached token and creates a new one
    """
    file_path = user_cache_dir(CLIENT_ID) + "/" + sha256(host.encode()).hexdigest()
    if os.path.exists(file_path) and not force_refresh:
        token = secret_read(file_path, host)
        try:
            return _check_token(token, host)
        except Exception:
            logging.warning("Cached token is invalid or refresh failed; clearing cache.")
            os.remove(file_path)
            if not login_if_needed:
                return None
            token = get_middlecat_token(host)
    elif login_if_needed:
        token = get_middlecat_token(host)
    else:
        return None
    return _check_token(token, host)


def _check_token(token, host) -> dict:
    """
    Returns refreshed token if old token has expired
    :param token: old token
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api").
    """
    if "expires_at" in token:
        if datetime.now() + timedelta(seconds=10) > datetime.strptime(token["expires_at"], "%Y-%m-%dT%H:%M:%S"):
            token = token_refresh(token, host)
    return token


def cache_token(token, host) -> None:
    """
    Caches encrypted token on disk
    :param token: old token
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api").
    """
    file_path = user_cache_dir(CLIENT_ID) + "/" + sha256(host.encode()).hexdigest()
    dir_path = os.path.dirname(file_path)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    fernet = Fernet(make_key(host))
    data = fernet.encrypt(dumps(token).encode())
    with open(file_path, "wb") as f:
        f.write(data)


def secret_read(path, host) -> dict:
    """
    Reads encrypted token from disk
    :param path: path to file, usually in user_cache_dir(CLIENT_ID)
    :param host: The URL to the AmCAT instance (e.g. "http://localhost/api").
    """
    with open(path, "rb") as f:
        token_enc = f.read()
    fernet = Fernet(make_key(host))
    return loads(fernet.decrypt(token_enc).decode())


def make_key(key) -> bytes:
    """
    Helper function to make key for encryption of tokens
    :param key: string that is turned into key (the host URL).
    """
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=key.encode(),
        iterations=600_000,
    )
    return urlsafe_b64encode(kdf.derive(key.encode()))


def base64_url_encode(x) -> str:
    """
    Custom base64 encode for pkce challange nicked from httr2
    https://github.com/r-lib/httr2/blob/main/R/utils.R
    :param x: string to be encoded.
    """
    x = b64encode(x).decode("utf-8")
    # Replace some characters to align output with the javascript version
    x = x.rstrip("=")
    x = x.replace("+", "-").replace("/", "_")
    return x


def pkce_challange() -> dict:
    """
    Generates PKCE code challange for middlecat requests
    :param x: string to be encoded.
    """
    # Generate random 32-octet sequence
    verifier_bytes = getrandbits(256).to_bytes(32, byteorder="big")
    verifier = base64_url_encode(verifier_bytes)
    challenge_bytes = sha256(verifier.encode("utf-8")).digest()
    challenge = base64_url_encode(challenge_bytes)
    return {"verifier": verifier, "method": "S256", "challenge": challenge}
