import os
import logging
from vrchatapi import ApiClient as _ApiClient
from vrchatapi import GroupsApi, Configuration, AuthenticationApi
from http.cookiejar import LWPCookieJar
from vrchatapi.exceptions import UnauthorizedException, ApiException
from vrchatapi.models.two_factor_auth_code import TwoFactorAuthCode
from vrchatapi.models.two_factor_email_code import TwoFactorEmailCode
from totp import get_totp

# Modded ApiClient with cookie support
# Credit: Katsi on the VRC API Community Discord 
class ApiClient(_ApiClient):
    def save_cookies(self):
        # validate env variable and create cookie jar
        filename = os.getenv("COOKIE_FILE")
        if filename is not None:
            cookie_jar = LWPCookieJar(filename=filename)
        else:
            raise ValueError("COOKIE_FILE env variable not set")
        # load cookies from API client into memory
        for cookie in self.rest_client.cookie_jar:
            cookie_jar.set_cookie(cookie)

        # save cookies
        logging.info("Attempting to save cookies")     
        cookie_jar.save()
    
    def load_cookies(self):
        # validate env variable and create cookie jar
        filename = os.getenv("COOKIE_FILE")
        if filename is not None:
            cookie_jar = LWPCookieJar(filename=filename)
        else:
            raise ValueError("COOKIE_FILE env variable not set")
        
        # load existing cookies if possible, else create a new cookie jar
        try:
            logging.info("Attempting to load cookies")
            cookie_jar.load()
        except FileNotFoundError:
            logging.info("Attempting to save cookies")
            cookie_jar.save()
            return
        # transfer cookies from our session
        for cookie in cookie_jar:
            self.rest_client.cookie_jar.set_cookie(cookie)

def get_group_instances():
    group_instances = []

    configuration = Configuration( 
        # Get username/password from env variables
        username = os.getenv("USER_NAME"),
        password = os.getenv("USER_PASSWORD"),
    )

    group_id = os.getenv("GROUP_ID")

    with ApiClient(configuration) as api_client:
        # Set our User-Agent as per VRChat Usage Policy and call the Auth API
        api_client.user_agent = "GroupInstanceLogger/0.1alpha me@williamtpeebles.com"

        # Attempt to load cookies
        api_client.load_cookies()

        auth_api = AuthenticationApi(api_client)

        try:
            # Calling getCurrentUser on Authentication API logs you in if the user isn't already logged in.
            current_user = auth_api.get_current_user()
            api_client.save_cookies() # Save cookies after successful login
        except UnauthorizedException as e:
            handle_auth_exception(auth_api, e)
            current_user = auth_api.get_current_user()
            api_client.save_cookies() # Save cookies after successful login

        logging.info(f"Logged in as: {current_user.display_name}")

        # Call groups APIs
        groups_api_instance = GroupsApi(api_client)

        try:
            api_response = groups_api_instance.get_group_instances(group_id) # Get group instances
            for instance in api_response:
                group_instances.append(instance)
        except ApiException as e:
            logging.error(f"Exception when calling GroupsApi->get_group_instances: {e}")
        except Exception as e:
            logging.error(f"Unexpected Error: {e}")

    return group_instances

def handle_auth_exception(auth_api, e):
    if e.status == 200:
        if "Email 2 Factor Authentication" in e.reason:
            # Calling email verify2fa if the account has 2FA disabled, give warning that automatic runs cannot be done
            logging.warning("IMPORTANT: Your VRChat Account is only enabled for **email** based 2-factor authentication! "
                            "You will not be able to automate runs of GroupInstanceLogger\n"
                            "Set up your account to use TOTP-based 2FA so the application can log you in automatically\n\n")
            auth_api.verify2_fa_email_code(two_factor_email_code=TwoFactorEmailCode(input("Email 2FA Code: ")))
        elif "2 Factor Authentication" in e.reason:
            # Check if TOTP code is defined in env variables
            totp_key = os.getenv("TOTP_KEY")
            if totp_key:
                logging.info("Attempting to authenticate with defined TOTP key")
                auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(get_totp()))
            else:
                logging.info("No TOTP key defined, prompting manual input of 2FA code")
                auth_api.verify2_fa(two_factor_auth_code=TwoFactorAuthCode(input("Enter 2FA Code: ")))
    else:
        logging.error(f"UnauthorizedException: {e}")