import typing
from base64 import b64encode
from datetime import datetime
from enum import StrEnum
from typing import List

import niquests
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
from niquests import PreparedRequest, Response

TRANSFERWISE_BASE_URL = "https://api.transferwise.com"


class StatementType(StrEnum):
    COMPACT = "COMPACT"  # a single statement line per transaction
    FLAT = "FLAT"  # accounting statements where transaction fees are on a separate line


class TransferWise2FASession(niquests.Session):
    """
    TransferWise 2FA session implementation.
    """

    def __init__(self, *args, tfa_private_key: typing.Optional[PrivateKeyTypes] = None, **kwargs):
        super().__init__(*args, **kwargs)

        self.tfa_private_key = tfa_private_key

    def send(self, request: PreparedRequest, **kwargs: typing.Any) -> Response:
        resp = super().send(request, **kwargs)

        needs_2fa = (
            resp.status_code == 403
            and "x-2fa-approval" in resp.headers
            and resp.headers.get("x-2fa-approval-result") == "REJECTED"
        )

        if needs_2fa:
            if not self.tfa_private_key:
                raise RuntimeError("2FA required by no private key provided.")

            tfa_flow_id = resp.headers["x-2fa-approval"]

            tfa_signature = self.tfa_private_key.sign(
                tfa_flow_id.encode("utf-8"),
                padding.PKCS1v15(),
                hashes.SHA256(),
            )

            request.headers.update(
                {
                    "x-2fa-approval": tfa_flow_id,
                    "X-Signature": b64encode(tfa_signature).decode("utf-8"),
                }
            )

            return super().send(request, **kwargs)

        return resp


class TransferwiseClient:
    def __init__(self, access_token: str, private_key: PrivateKeyTypes):
        self.session = TransferWise2FASession(tfa_private_key=private_key)
        self.session.headers.update({"Authorization": f"Bearer {access_token}"})

    def get_balance_by_profile_and_balance_id(self, profile_id: str, balance_id: str) -> dict:
        resp = self.session.get(
            TRANSFERWISE_BASE_URL + f"/v4/profiles/{profile_id}/balances/{balance_id}",
        )

        resp.raise_for_status()

        return resp.json()

    def get_statement_by_profile_and_balance_id(
        self,
        profile_id: str,
        balance_id: str,
        currency: str,
        from_: datetime,
        to: datetime,
        type: StatementType = StatementType.FLAT,
        locale: str = "en",
    ) -> List[dict]:
        resp = self.session.get(
            TRANSFERWISE_BASE_URL + f"/v1/profiles/{profile_id}/balance-statements/{balance_id}/statement.json",
            params={
                "currency": currency,
                "intervalStart": from_.isoformat(),
                "intervalEnd": to.isoformat(),
                "type": type,
                "statementLocale": locale,
            },
        )

        resp.raise_for_status()

        return resp.json().get("transactions", [])


class TransferWiseProfileClient(TransferwiseClient):
    def __init__(self, *args, profile_id: str, **kwargs):
        super().__init__(*args, **kwargs)

        self.profile_id = profile_id


class TransferWiseBalanceClient(TransferWiseProfileClient):
    def __init__(self, *args, balance_id: str, **kwargs):
        super().__init__(*args, **kwargs)

        self.balance_id = balance_id

    def get_balance(self):
        return self.get_balance_by_profile_and_balance_id(self.profile_id, self.balance_id)

    def get_statement(
        self,
        currency: str,
        from_: datetime,
        to: datetime,
        type: StatementType = StatementType.FLAT,
        locale: str = "en",
    ):
        return self.get_statement_by_profile_and_balance_id(
            profile_id=self.profile_id,
            balance_id=self.balance_id,
            currency=currency,
            from_=from_,
            to=to,
            type=type,
            locale=locale,
        )