import requests
import typing as t
from uuid import uuid4


HEADERS = (
    ACCOUNT_COUNTRY_HEADER := "X-Apple-ID-Account-Country",
    SESSION_ID_HEADER := "X-Apple-ID-Session-Id",
    SESSION_TOKEN_HEADER := "X-Apple-Session-Token",
    TWOSV_TRUST_TOKEN_HEADER := "X-Apple-TwoSV-Trust-Token",
    SCNT_HEADER := "scnt",
)


class BaseUrlConfig(t.NamedTuple):
    auth: str = "https://idmsa.apple.com/appleauth/auth"
    home: str = "https://www.icloud.com"
    setup: str = "https://setup.icloud.com/setup/ws/1"


class ICloudClient:
    def __init__(
        self, apple_id: str, password: str, base_urls: BaseUrlConfig = BaseUrlConfig()
    ) -> None:
        self.apple_id = apple_id
        self.password = password
        self.base_urs = base_urls
        self.requester = requests.Session()
        self.requester.headers["Origin"] = self.base_urs.home
        self.requester.headers["Referer"] = f"{self.base_urs.home}/"
        self.logger = print
        self.client_id = f"auth-{uuid4()}"
        self._webservices = {}

    def _handle_response(self, response: requests.Response) -> None:
        if not 199 < response.status_code < 300:
            cookie_names = ", ".join([c.name for c in self.requester.cookies])
            self.logger(
                f"Call failed with {response.status_code} code and "
                f"body: {response.text} and cookies: {cookie_names}"
            )
            raise RuntimeError

        for header in HEADERS:
            if header in response.headers:
                self.requester.headers[header] = response.headers[header]

    def _get_auth_headers(self, overrides: t.Optional[t.Mapping[str, str]] = None):
        headers = {
            "Accept": "*/*",
            "Content-Type": "application/json",
            "X-Apple-OAuth-Client-Id": "d39ba9916b7251055b22c7f910e2ea796ee65e98b2ddecea8f5dde8d9d1a815d",
            "X-Apple-OAuth-Client-Type": "firstPartyAuth",
            "X-Apple-OAuth-Redirect-URI": self.base_urs.home,
            "X-Apple-OAuth-Require-Grant-Code": "true",
            "X-Apple-OAuth-Response-Mode": "web_message",
            "X-Apple-OAuth-Response-Type": "code",
            "X-Apple-OAuth-State": self.client_id,
            "X-Apple-Widget-Key": "d39ba9916b7251055b22c7f910e2ea796ee65e98b2ddecea8f5dde8d9d1a815d",
        }
        if overrides:
            headers.update(overrides)

        return headers

    def webservice_url(self, service_name: str) -> str:
        try:
            return self._webservices[service_name]["url"]
        except KeyError:
            raise RuntimeError(f"{service_name} does not exist")

    def sign_in(self, remember_me: bool = True) -> None:
        resp = self.requester.post(
            f"{self.base_urs.auth}/signin",
            json={
                "accountName": self.apple_id,
                "password": self.password,
                "rememberMe": remember_me,
                "trustTokens": [],
            },
            headers=self._get_auth_headers(),
        )

        self._handle_response(resp)

    def account_login(self) -> None:
        resp = self.requester.post(
            f"{self.base_urs.setup}/accountLogin",
            json={
                "accountCountryCode": self.requester.headers.get(
                    ACCOUNT_COUNTRY_HEADER
                ),
                "dsWebAuthToken": self.requester.headers.get(SESSION_TOKEN_HEADER),
                "extended_login": True,
                "trustToken": self.requester.headers.get(TWOSV_TRUST_TOKEN_HEADER),
            },
        )
        self._handle_response(resp)
        self._webservices = resp.json()["webservices"]

    def verify_2fa_code(self, code: str) -> None:
        resp = self.requester.post(
            f"{self.base_urs.auth}/verify/trusteddevice/securitycode",
            json={"securityCode": {"code": code}},
            headers=self._get_auth_headers({"Accept": "application/json"}),
        )
        self._handle_response(resp)

    def trust_device(self) -> None:
        resp = self.requester.get(
            f"{self.base_urs.auth}/2sv/trust", headers=self._get_auth_headers()
        )
        self._handle_response(resp)


class PremiumMailSettings:
    def __init__(self, client: ICloudClient):
        self.client = client
        self.base_url = f"{client.webservice_url('premiummailsettings')}/v1"

    def list_hme(self) -> t.Sequence[t.Mapping[str, str]]:
        resp = self.client.requester.get(f"{self.base_url}/hme/list")
        self.client._handle_response(resp)
        return resp.json()["result"]["hmeEmails"]


if __name__ == "__main__":
    client = ICloudClient(apple_id="foo@bar.com", password="FooBarBuz!")
    client.sign_in()
    client.account_login()
    code = input("Please enter verification code: ")
    client.verify_2fa_code(code)
    client.trust_device()
    client.account_login()
    pms = PremiumMailSettings(client)
    print(pms.list_hme())