import base64
import logging
import os
import time
from contextlib import asynccontextmanager
from http.cookies import SimpleCookie
from pathlib import Path
from typing import Optional

from clarabot_crypto.auth.login import (
    create_account_login_request_params,
    create_account_login_private_keyring,
    create_account_login_keyring,
    create_account_login_session,
)
from clarabot_crypto.base.key.kdf import HKDFKeyGeneratorV1, HKDFKeyGeneratorV2
from clarabot_crypto.keyring.private import PRIVATE_KEYRING_CLASSES, AbstractPrivateKeyring

from clarabot_nano.common.auth.two_factor import TwoFactorError, gateway_to_control
from clarabot_nano.common.crypto.keyring import NanoSelfAccountKeyring
from clarabot_nano.utility.serialize import msgpack_loads, msgpack_dumps

logger = logging.getLogger(__name__)


class LoginFailed(Exception):
    """
    This is used to be distinguished from unexpected errors during authentication.
    """


class ServerAuthenticator(object):
    SELF_KR = "self_kr"
    CLIENT_SESSION_KEY = "client_session_key"
    SALT = "salt"

    def __init__(self, session, url_domain, url_path, data_dir: Path):
        self._s = session
        self._d = url_domain
        self._p = url_path
        self._store_dir = data_dir / "gateway"

    def _cookie_storage_path(self):
        self._store_dir.mkdir(parents=True, exist_ok=True)
        return self._store_dir / ".cooks"

    def _local_storage_path(self):
        self._store_dir.mkdir(parents=True, exist_ok=True)
        return self._store_dir / ".lstore"

    def _continue_token_path(self):
        self._store_dir.mkdir(parents=True, exist_ok=True)
        return self._store_dir / ".ctoken"

    def _combine_server_and_client_keys(self, server_key: bytes, client_key: bytes, salt: Optional[bytes]) -> bytes:
        if not isinstance(server_key, bytes) or len(server_key) < 16:
            raise Exception("invalid or length of server key is low")
        if not isinstance(client_key, bytes) or len(client_key) < 16:
            raise Exception("invalid or length of client key is low")
        if salt is None:  # older key derivation with static salt
            return HKDFKeyGeneratorV1.new_key(64, server_key + client_key, HKDFKeyGeneratorV1.SALTS[1], None)
        if not isinstance(salt, bytes) or len(salt) < 16:
            raise Exception("invalid or length of salt is low")
        return HKDFKeyGeneratorV2.new_key(64, server_key + client_key, salt, None)

    def _read_local_storage(self, server_session_key: bytes) -> dict:
        store = self._local_storage_path().read_bytes()
        store = msgpack_loads(store)
        if not isinstance(store, dict):
            raise Exception("invalid value from store")
        combined_session_key = self._combine_server_and_client_keys(
            server_session_key, store[self.CLIENT_SESSION_KEY], store.get(self.SALT)
        )
        store[self.SELF_KR] = NanoSelfAccountKeyring.load_local_storage(store[self.SELF_KR], combined_session_key)
        return store

    def _write_local_storage(self, server_session_key: bytes, self_kr: NanoSelfAccountKeyring):
        combined_session_key = self._combine_server_and_client_keys(
            server_session_key, client_session_key := os.urandom(32), salt := os.urandom(16)
        )
        store = {
            self.SELF_KR: self_kr.dump_local_storage(combined_session_key),
            self.CLIENT_SESSION_KEY: client_session_key,
            self.SALT: salt,
        }
        store = msgpack_dumps(store)
        self._local_storage_path().write_bytes(store)

    async def login(
        self,
        email: str,
        password: str,
        timeout=10.0,
        server_session_key: Optional[bytes] = None,
        otp_salt: Optional[str] = None,
        tfa_code: Optional[str] = None,
    ) -> NanoSelfAccountKeyring:
        if not server_session_key:
            # NOTE: For some 2FA logins this must be provided by the client code to have it be consistent.
            server_session_key = os.urandom(16)
        t = time.perf_counter()
        # 1. acquire salt for PPF
        async with self.post_auto_csrf(
            f"https://{self._d}{self._p}r/login-salt", json={"email": email}, timeout=timeout
        ) as resp:
            if resp.status != 200:
                raise Exception("Failed login (salt)")
            salt = base64.b64decode((await resp.json())["salt"], validate=True)
        # 2. login with best private keyring
        for private_kr_cls, password in AbstractPrivateKeyring.flattened_password_forms_with_cls(
            PRIVATE_KEYRING_CLASSES, password
        ):
            private_kr = create_account_login_private_keyring(private_kr_cls, password, salt)
            request = create_account_login_request_params(private_kr, email, server_session_key, nano=True)
            if otp_salt is not None:
                request["otpSalt"] = otp_salt
            if tfa_code is not None:
                request["2faCode"] = tfa_code
            async with self.post_auto_csrf(
                f"https://{self._d}{self._p}r/login",
                json=request,
                timeout=None if timeout is None else (t + timeout) - time.perf_counter(),
            ) as resp:
                if resp.status != 200:
                    logger.warning("Login response is unsuccessful: %s", resp.status)
                    continue
                login_data = await resp.json()
                if not isinstance(login_data, dict):
                    logger.warning("Login data response is incorrect: %s", login_data)
                    continue
                if two_factor := login_data.get("2FA"):
                    raise TwoFactorError(gateway_to_control(two_factor))
                for k in [
                    "encrypted_master_key",
                    "encrypting_secret_ec",
                    "signing_secret_ec",
                    "encrypted_session_key",
                    "public_key",
                ]:
                    login_data[k] = base64.b64decode(login_data[k], validate=True)
                break
        else:
            raise LoginFailed("Failed login (auth)")
        del password  # no longer needed
        del salt  # no longer needed
        # 3. decrypt the account's keyring and session key
        self_kr = create_account_login_keyring(
            NanoSelfAccountKeyring,
            private_kr,
            login_data["account_id"],
            login_data["encrypted_master_key"],
            login_data["encrypting_secret_ec"],
            login_data["signing_secret_ec"],
        )
        session_key = create_account_login_session(
            self_kr, login_data["encrypted_session_key"], login_data["public_key"]
        )
        del login_data  # no longer needed
        del private_kr  # no longer needed
        # 4.load session key (auth token) into session
        cooks = SimpleCookie({"k": base64.b64encode(session_key).decode("ascii")})
        cooks["k"]["domain"] = self._d
        cooks["k"]["path"] = "/"
        cooks["k"]["secure"] = True
        # other attributes are not applicable really; just mirror browser behaviour
        cooks["k"]["httponly"] = False
        cooks["k"]["samesite"] = "strict"
        self._s.cookie_jar.update_cookies({"k": cooks["k"]})
        # 5. save authentication to local storage with combined session encryption
        self._s.cookie_jar.save(self._cookie_storage_path())
        self._write_local_storage(server_session_key, self_kr)
        # 6. return keys for usage
        return self_kr

    async def logout(self, timeout=10.0) -> bool:
        """
        When logout is requested, the local tokens will be removed even if the server session cannot logout
        for any reason. (network-offline, already destroyed, etc.)
        """
        cookie_storage_path = self._cookie_storage_path()
        if cookie_storage_path.is_file():
            try:
                self._s.cookie_jar.load(cookie_storage_path)
                async with self.post_auto_csrf(f"https://{self._d}{self._p}r/logout", timeout=timeout) as resp:
                    resp.raise_for_status()
            except Exception:
                logger.exception("Failed to perform logout on server, local tokens will still be purged")
        self._s.cookie_jar.clear()
        cookie_storage_path.unlink(missing_ok=True)
        self._continue_token_path().unlink(missing_ok=True)
        self._local_storage_path().unlink(missing_ok=True)
        return True

    async def resume(self, timeout=10.0) -> NanoSelfAccountKeyring:
        """
        Holistic authenticated session resumption. Also acquires new CSRF token if needed.
        """
        cookie_storage_path = self._cookie_storage_path()
        if cookie_storage_path.is_file():
            self._s.cookie_jar.load(cookie_storage_path)
        else:
            self._s.cookie_jar.clear()
        async with self.post_auto_csrf(
            f"https://{self._d}{self._p}h", json={"query": "query sessionEncryptionKey {"}, timeout=timeout
        ) as resp:
            if resp.status != 200:
                raise Exception("Failed resume (sess)")
            server_session_key = base64.b64decode((await resp.json())["data"]["sessionEncryptionKey"], validate=True)
        store = self._read_local_storage(server_session_key)
        if not isinstance(store[self.SELF_KR], NanoSelfAccountKeyring):
            raise Exception("Failed resume (lstore)")
        if store.get(self.SALT) is None:  # upgrade session encryption
            self._write_local_storage(server_session_key, store[self.SELF_KR])
        return store[self.SELF_KR]

    def save_continue_token(self, token: str):
        self._continue_token_path().write_text(token, "utf-8")

    def get_continue_token(self):
        continue_token_path = self._continue_token_path()
        if continue_token_path.is_file():
            return continue_token_path.read_text("utf-8")

    def could_try_resume(self) -> bool:
        return self._cookie_storage_path().is_file() and self._local_storage_path().is_file()

    async def csrf_from_session(self, timeout=None):
        for c in self._s.cookie_jar:
            if c.key == "__Secure-csrf":
                return c.value
        else:
            async with self._s.get(f"https://{self._d}{self._p}r/csrf", timeout=timeout) as r:
                if r.status != 200:
                    raise Exception("Failed to request CSRF token")
        for c in self._s.cookie_jar:
            if c.key == "__Secure-csrf":
                return c.value
        else:
            raise Exception("Failed to acquire CSRF token")

    @asynccontextmanager
    async def post_auto_csrf(self, url, **kwargs):
        t = time.perf_counter()
        timeout = kwargs.get("timeout")
        csrf = await self.csrf_from_session(timeout=timeout)
        kwargs["headers"] = list(kwargs.get("headers", ()))
        kwargs["headers"].append(("pcsrf", csrf))
        if timeout is not None:
            kwargs["timeout"] = (t + timeout) - time.perf_counter()
        async with self._s.post(url, **kwargs) as r:
            yield r
