from requests_oauthlib import OAuth2Session
import requests


class ProsodyRestSession(OAuth2Session):
    def __init__(
        self, base_url, client_name, client_uri, redirect_uri, *args, **kwargs
    ):
        self.base_url = base_url
        discovery_url = base_url + "/.well-known/oauth-authorization-server"

        meta = requests.get(discovery_url).json()
        reg = requests.post(
            meta["registration_endpoint"],
            json={
                "client_name": client_name,
                "client_uri": client_uri,
                "redirect_uris": [redirect_uri],
            },
        ).json()

        super().__init__(client_id=reg["client_id"], *args, **kwargs)

        self.meta = meta
        self.client_secret = reg["client_secret"]
        self.client_id = reg["client_id"]

    def authorization_url(self, *args, **kwargs):
        return super().authorization_url(
            self.meta["authorization_endpoint"], *args, **kwargs
        )

    def fetch_token(self, *args, **kwargs):
        return super().fetch_token(
            token_url=self.meta["token_endpoint"],
            client_secret=self.client_secret,
            *args,
            **kwargs
        )

    def xmpp(self, json=None, *args, **kwargs):
        return self.post(self.base_url + "/rest", json=json, *args, **kwargs)


if __name__ == "__main__":
    # Example usage

    # from prosody_oauth import ProsodyRestSession
    from getpass import getpass

    p = ProsodyRestSession(
        input("Base URL: "),
        "Prosody mod_rest OAuth 2 example",
        "https://modules.prosody.im/mod_rest",
        "urn:ietf:wg:oauth:2.0:oob",
    )

    print("Open the following URL in a browser and login:")
    print(p.authorization_url()[0])

    p.fetch_token(code=getpass("Paste Authorization code: "))

    print(p.xmpp(json={"disco": True, "to": "jabber.org"}).json())