#!/usr/bin/env python3
import os, sys, time, hashlib, random, string, requests, re, shutil, tkinter as tk
from tkinter import messagebox, simpledialog

# --- Logger ---
class Logger:
    def __init__(self, logfile):
        self.logfile = logfile
        with open(self.logfile, "w") as f:
            f.write(f"Starting OpenAI_Import V2.0 -- {time.ctime()}\n")

    def log(self, msg=""):
        print(msg)
        with open(self.logfile, "a") as f:
            f.write(msg + "\n")


# --- Cloudshare API Client ---
class CloudshareClient:
    def __init__(self, client_id, client_secret, logger):
        self.client_id = client_id
        self.client_secret = client_secret
        self.logger = logger
        self.oauth_url = "https://api.accelerate.cloudshare.com/v4/oauth/token"

    def get_class_id(self, timeout: int = 3):
        """
        Retrieve the class-id from CloudShare metadata.
        Timeout is reduced to a few seconds to avoid long hangs.
        """
        url = "https://metadata.cloudshare.com/api/v3/unauthenticated/metadata"
        try:
            resp = requests.get(url, timeout=timeout)
            resp.raise_for_status()
            class_id = resp.json().get("class-id")
            if not class_id or class_id == "null":
                raise ValueError("No valid class-id found in metadata response.")
            return class_id    
        except requests.exceptions.Timeout:
            self.logger.log(f"Error: Metadata request timed out after {timeout} seconds.")
            return None
        except requests.exceptions.RequestException as e:
            self.logger.log(f"Error retrieving class-id: {e}")
            return None
        except ValueError as ve:
            self.logger.log(f"Error: {ve}")
            return None

    def get_class_name(self, classid):
        api_url = f"https://use.cloudshare.com/api/v3/class/{classid}"
        timestamp = str(int(time.time()))
        token = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
        hmac_input = f"{self.client_secret}{api_url}{timestamp}{token}"
        hmac_val = hashlib.sha1(hmac_input.encode()).hexdigest()
        auth_param = f"userapiid:{self.client_id};timestamp:{timestamp};token:{token};hmac:{hmac_val}"

        resp = requests.get(api_url, headers={
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"cs_sha1 {auth_param}"
        })
        env = resp.json().get("name")
        self.logger.log(f": Class name pulled from API is: {env}")
        return env

    def get_access_token(self):
        post_data = {
            "grant_type": "client_credentials",
            "client_id": self.client_id,
            "client_secret": self.client_secret
        }
        resp = requests.post(self.oauth_url, data=post_data)
        tokdata = resp.json()
        access_token = tokdata.get("access_token")
        token_type = tokdata.get("token_type")
        expires_in = tokdata.get("expires_in")

        if not access_token:
            self.logger.log("Error: Failed to obtain OAuth2 Access Token.")
            sys.exit(1)

        self.logger.log(": Cloudshare API access token obtained successfully!")
        self.logger.log(f": Token Type: {token_type}")
        self.logger.log(f": Expires In: {expires_in} seconds")
        return token_type, access_token

    def get_openai_key(self, env, token_type, access_token):
        api_url = f"https://api.accelerate.cloudshare.com/v4/trainings?name=contains({env})"
        resp = requests.get(api_url, headers={
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"{token_type} {access_token}"
        })
        desc = resp.json()["items"][0]["description"]
        openai = desc.split("=")[1] if "=" in desc else None
        self.logger.log(f": Description field pulled is: {desc}")
        self.logger.log(f": Pulled value to use is: {openai}")
        return openai


# --- File Editor ---
class FileEditor:
    def __init__(self, logger):
        self.logger = logger

    def update_key_in_files(self, files, key):
        edited = False
        for f in files:
            if os.path.isfile(f):
                with open(f, "r+") as fh:
                    content = fh.read()
                    content = re.sub(r'OPENAI_API_KEY=.*', f'OPENAI_API_KEY={key}', content)
                    fh.seek(0); fh.write(content); fh.truncate()
                self.logger.log(f": Edited OpenAI key into file {f}")
                edited = True
            else:
                self.logger.log(f": Did not locate file {f}")
        return edited

    def replace_file_with_key(self, files, key):
        edited = False
        for f in files:
            if os.path.isfile(f):
                with open(f, "w") as fh:
                    fh.write(key)
                self.logger.log(f": Inserted OpenAI key into file {f}")
                edited = True
            else:
                self.logger.log(f": Did not locate file {f}")
        return edited


# --- Tkinter Popup Manager with Timeout ---
class TkPopupManager:
    def __init__(self, enabled, logger, timeout=120):
        self.enabled = enabled
        self.logger = logger
        self.timeout = timeout
        self.root = tk.Tk()
        self.root.withdraw()

    def ask_install(self):
        if not self.enabled:
            return True

        result = [None]

        def default_yes():
            if result[0] is None:
                result[0] = True
                self.root.quit()

        # Schedule timeout
        self.root.after(self.timeout * 1000, default_yes)

        result[0] = messagebox.askyesno("OpenAI Key",
                                        "Do you wish to install the course OpenAI key?\n\n(times out defaults to YES)")
        if result[0] is False:
            self.logger.log("Exited without installing the key")
            return False
        return True

    def notify_success(self):
        if self.enabled:
            messagebox.showinfo("OpenAI Key", "The OpenAI key has been updated")


# --- Main ---
def main():
    SETPOPUPS = "YES" if not any(arg in sys.argv for arg in ["-s", "/s"]) else ""
    # Use current working directory for log file to ensure it's writable in PyInstaller
    log_path = os.path.join(os.getcwd(), "OpenAI_Import.log")
    logger = Logger(log_path)
    client = CloudshareClient(os.getenv("CLIENT_ID", "538GACRZ9FABA07E"),
                              os.getenv("CLIENT_SECRET", "bkAyYSB3dL8GREF87zfqhjkUcW5xNSVyj0HDpQBqheX5vcHpbLE8WXqYjGmGMORt"),
                              logger)
    editor = FileEditor(logger)
    popups = TkPopupManager(SETPOPUPS, logger, timeout=120)

    classid = client.get_class_id()
    if not classid or classid == "null":
        sys.exit(1)

    env = client.get_class_name(classid)
    token_type, access_token = client.get_access_token()
    openai_key = client.get_openai_key(env, token_type, access_token)

    if not openai_key:
        logger.log("No key provided in VM metadata")
        sys.exit(1)

    if not popups.ask_install():
        sys.exit(0)

    edited = False
    edited |= editor.update_key_in_files([
        "/home/student/.config/shell_gpt/.sgptrc",
        "/home/.bashsrc",
        "/home/student/.bashrc"
    ], openai_key)

    edited |= editor.replace_file_with_key([
        "/home/student/keys/key.txt",
        "/home/student/Keys/Key.txt",
        "/home/keys/key.txt"
    ], openai_key)

    if edited:
        popups.notify_success()

    logger.log("\nAll listed files have been edited if found, Exiting")


if __name__ == "__main__":
    main()
