import requests
import webbrowser

from dotenv import load_dotenv
from os import getenv

load_dotenv()

# TODO: Add support for paper trading API endpoint.

class TradeStationClient:

    def __init__(self, refresh_token: str = None) -> None:
        self.id = getenv('TRADESTATION_CLIENT_ID')
        self.redirect_uri = getenv('TRADESTATION_REDIRECT_URI')
        self.scope = getenv('TRADESTATION_CLIENT_SCOPE')
        self.secret = getenv('TRADESTATION_CLIENT_SECRET')
        self.state = getenv('TRADESTATION_CLIENT_STATE')

        # Must be retrieved via authorization flow.
        # TODO: Automate authorization.
        self.access_token = None
        self.id_token = None
        self.refresh_token = None

        # For bypassing normal authorization flow.
        if refresh_token:
            self.refresh_token = refresh_token
            self.refresh_access_token()

    def open_authorization_url(self):
        """
        Open the TradeStation authorization URL in the default web browser.
        """
        url = (
            f"https://signin.tradestation.com/authorize?"
            f"response_type=code&"
            f"client_id={self.id}&"
            f"redirect_uri={self.redirect_uri}&"
            f"audience=https://api.tradestation.com&"
            f"state={self.state}&"
            f"scope={self.scope}"
        )
        webbrowser.open(url)

    def get_tokens(self, authorization_code):
        """
        Exchange the authorization code for access token, ID token, and refresh token.
        """
        url = 'https://signin.tradestation.com/oauth/token'
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        payload = {
            'grant_type': 'authorization_code',
            'client_id': self.id,
            'client_secret': self.secret,
            'code': authorization_code,
            'redirect_uri': self.redirect_uri
        }

        response = requests.post(url, headers=headers, data=payload)

        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Failed to get tokens: {response.text}")

    def refresh_access_token(self):
        """
        Use the refresh token to obtain a new access token.
        """
        if not self.refresh_token:
            raise Exception("Refresh token is not available.")

        url = 'https://signin.tradestation.com/oauth/token'
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        payload = {
            'grant_type': 'refresh_token',
            'client_id': self.id,
            'refresh_token': self.refresh_token,
            'client_secret': self.secret
        }

        response = requests.post(url, headers=headers, data=payload)

        if response.status_code == 200:
            tokens = response.json()
            self.access_token = tokens.get('access_token')
            self.id_token = tokens.get('id_token')
        else:
            raise Exception(f"Failed to refresh token: {response.text}")

if __name__ == '__main__':
    client = TradeStationClient()
    client.open_authorization_url()
    authorization_code = input("Please enter the authorization code you received: ")

    try:
        tokens = client.get_tokens(authorization_code)
        print("Access Token:", tokens['access_token'])
        print("ID Token:", tokens['id_token'])
        print("Refresh Token:", tokens['refresh_token'])
    except Exception as e:
        print(str(e))