#!/usr/bin/env python
# by Dominik Stanisław Suchora <hexderm@gmail.com>
# License: GNU GPLv3


import json
from typing import List
import io

import treerequests
from reliq import RQ
import requests

from .exceptions import RequestError, AuthenticationError, FailedError

reliq = RQ(cached="True")


def data_decode_val(format, data):
    if isinstance(format, dict):
        return data_decode_dict(format, data)
    elif isinstance(format, list):
        return data_decode_list(format, data)
    else:
        return format


def data_decode_dict(format, data):
    r = {}
    for key, value in format.items():
        r[key] = data_decode_val(data[value], data)

    return r


def data_decode_list(format, data):
    r = []
    for i in format:
        r.append(data_decode_val(data[i], data))

    return r


def data_decode(data):
    assert data["type"] == "data"
    ret = []

    for i in data["nodes"]:
        if i["type"] != "data":
            continue

        if len(data) == 0:
            continue

        r = data_decode_val(i["data"][0], i["data"])
        ret.append(r)

    return ret


class Api:
    def __init__(self, password: str = "", **kwargs):
        """
        kwarg( password: str = "" ) if not empty used for login at initialization.

        All other kwargs are passed to treerequests session
        """

        self.ses = treerequests.Session(
            requests,
            requests.Session,
            lambda x, y: treerequests.reliq(x, y, obj=reliq),
            requesterror=RequestError,
            **kwargs,
        )

        self.password = password
        if len(password) > 0:
            self.login(password)

        self.token = ""

    @staticmethod
    def _valid_password(password: str):
        assert len(password) == 16
        assert password.isdigit()

    def _force_login(self):
        if len(self.token) != 68:
            raise AuthenticationError()
        if self.token != self.ses.cookies["accessToken"]:
            raise AuthenticationError()

    def login(self, password: str):
        """
        Logs user in with arg( password )
        """

        self._valid_password(password)

        try:
            r = self.ses.post_json(
                "https://mullvad.net/en/account/login",
                headers={
                    "Referer": "https://mullvad.net/en/account/login",
                    "content-type": "application/x-www-form-urlencoded",
                    "Origin": "https://mullvad.net",
                },
                data={"account_number": password},
            )
        except Exception:
            raise AuthenticationError()

        if r != {"type": "redirect", "status": 302, "location": "/en/account"}:
            raise AuthenticationError()

        self.token = self.ses.cookies["accessToken"]

        self._force_login()
        self.password = password

    def logout(self):
        """
        Logs user out
        """

        self.token = ""
        self.password = ""
        self.ses.cookies.pop("accessToken")

    def receipt(self, source: str, token: str) -> dict:
        """
        Gets receipt from using arg( source ) payment and arg( token ) of transaction.

        arg( source ) can be "swish", "paypal", "card", "other".

        login required
        """

        self._force_login()

        assert source in ("swish", "paypal", "card", "other")
        assert len(token) > 5

        r = self.ses.post_json(
            "https://mullvad.net/api/receipt",
            headers={
                "Referer": "https://mullvad.net/en/account/request-receipt/" + source,
                "Origin": "https://mullvad.net",
            },
            json={"token": token},
        )
        return r

    def status_raw(self) -> dict:
        """
        Get status of user.

        login required
        """

        self._force_login()

        r = self.ses.get_json(
            "https://mullvad.net/en/account/__data.json?x-sveltekit-invalidated=0100"
        )
        return r

    def status(self) -> dict:
        """
        Get status of user in format easier to use.

        login required
        """

        r = self.status_raw()
        return data_decode(r)[0]

    def devices_raw(self) -> dict:
        """
        List devices of user.

        login required
        """

        self._force_login()

        r = self.ses.get_json(
            "https://mullvad.net/en/account/devices/__data.json?x-sveltekit-invalidated=0101",
            headers={"Referer": "https://mullvad.net/en/account"},
        )
        return r

    def devices(self) -> List[dict]:
        """
        List devices of user in format easier to use.

        login required
        """

        r = self.devices_raw()
        r = data_decode(r)
        return r[1]["devices"]

    def create_device(self, key: str) -> dict:
        """
        Create a new device with arg( key ) wireguard public key.

        On error raises FailedError(), otherwise format similar to method( devices ) is returned about the newly created device.

        login required
        """

        self._force_login()
        assert len(key) == 44

        r = self.ses.post_json(
            "https://mullvad.net/en/account/devices?/upload-key",
            headers={
                "Referer": "https://mullvad.net/en/account/devices",
                "content-type": "application/x-www-form-urlencoded",
                "Origin": "https://mullvad.net",
            },
            data={"key": key},
        )

        if r["type"] != "success":
            raise FailedError()

        r = json.loads(r["data"])
        r = data_decode_val(r[0], r)
        return r["device"]

    def revoke_device(self, uuid: str) -> bool:
        """
        Remove device with arg( uuid ).

        Returns True on success and False on failure.

        login required
        """

        r = self.ses.post_json(
            "https://mullvad.net/en/account/devices?/revoke-device",
            headers={
                "content-type": "application/x-www-form-urlencoded",
                "Origin": "https://mullvad.net",
            },
            data={"id": uuid},
        )

        return r["type"] == "success"

    def version(self) -> str:
        """
        Get version string
        """

        r = self.ses.get_json("https://mullvad.net/_app/version.json")
        return r["version"]

    def create_account(self) -> str:
        """
        Create new account and return password.

        FailedError() is raised on failure
        """

        r = self.ses.post_json(
            "https://mullvad.net/en/account/create",
            headers={
                "Referer": "https://mullvad.net/en/account/create",
                "content-type": "application/x-www-form-urlencoded",
                "Origin": "https://mullvad.net",
            },
        )

        if r["type"] != "success":
            raise FailedError()

        return json.loads(r["data"])[1]

    def create_config(self, instance: dict, device: dict) -> str:
        """
        Creates wireguard configuration from arg( instance ) and arg( device ) derived from method( instances ) and method( devices )
        """

        with io.StringIO() as s:
            s.write("[Interface]\n# Device: ")
            s.write(device["name"])
            s.write("\nPrivateKey = ")
            s.write(device["pubkey"])
            s.write("\nAddress = ")
            s.write(device["ipv4_address"] + "," + device["ipv6_address"])
            s.write("\nDNS = 10.64.0.1\n\n[Peer]\nPublicKey = ")
            s.write(instance["pubkey"])
            s.write("\nAllowedIPs = 0.0.0.0/0,::0/0\nEndpoint = ")
            s.write(instance["ipv4_addr_in"] + ":" + str(instance["multihop_port"]))
            s.write("\n")

            return s.getvalue()

    def instances(self) -> List[dict]:
        """
        Gets full list of available server instances
        """

        rq = self.ses.get_html("https://mullvad.net/en/servers")
        r = rq.search(r"""
            [0] script nonce | "%i" sed "
                /^[ \t]*data: \[/{
                    s/^[ \t]*data: \[.*\{relays:\[/[/
                    s/\},uses:\{\}\},null\],$//
                    s/([{\[,])([a-zA-Z_][A-Z_a-z0-9]*):/\1\"\2\":/g
                    p
                    q
                }
            " "nE"
        """)

        return json.loads(r)
