#!/usr/bin/env python3
import os
import time
import random
import string
import hashlib
import requests
from pathlib import Path
import tkinter as tk
from tkinter import messagebox

LOGFILE = Path.cwd() / "OpenAI_Import.log"

def log(msg: str):
    print(msg)
    with open(LOGFILE, "a") as f:
        f.write(msg + "\n")

def get_metadata(timeout=3):
    url = "https://metadata.cloudshare.com/api/v3/unauthenticated/metadata"
    try:
        resp = requests.get(url, timeout=timeout)
        resp.raise_for_status()
        return resp.json()
    except requests.exceptions.Timeout:
        log(f"Error: metadata request timed out after {timeout} seconds.")
        return None
    except requests.exceptions.RequestException as e:
        log(f"Error retrieving metadata: {e}")
        return None

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

    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"cs_sha1 {auth_param}"
    }
    resp = requests.get(api_url, headers=headers, timeout=5)
    resp.raise_for_status()
    return resp.json().get("name")

def get_oauth_token(client_id, client_secret):
    url = "https://api.accelerate.cloudshare.com/v4/oauth/token"
    data = {
        "grant_type": "client_credentials",
        "client_id": client_id,
        "client_secret": client_secret
    }
    resp = requests.post(url, data=data, timeout=5)
    resp.raise_for_status()
    j = resp.json()
    return j.get("access_token"), j.get("token_type"), j.get("expires_in")

def get_openai_key(env_name, token, token_type):
    url = f"https://api.accelerate.cloudshare.com/v4/trainings?name=contains({env_name})"
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"{token_type} {token}"
    }
    resp = requests.get(url, headers=headers, timeout=5)
    resp.raise_for_status()
    description = resp.json()["items"][0]["description"]
    return description.split("=")[1]

def update_files(openai_key):
    files = [
        Path("/home/student/.config/shell_gpt/.sgptrc"),
        Path("/home/.bashsrc"),
        Path("/home/student/.bashrc"),
        Path("/home/student/keys/key.txt"),
        Path("/home/student/Keys/Key.txt"),
        Path("/home/keys/key.txt")
    ]
    edited = False
    for f in files:
        if f.exists():
            if f.suffix in [".rc", ".sgptrc"]:
                lines = f.read_text().splitlines()
                new_lines = []
                for line in lines:
                    if line.startswith("OPENAI_API_KEY="):
                        new_lines.append(f"OPENAI_API_KEY={openai_key}")
                    else:
                        new_lines.append(line)
                f.write_text("\n".join(new_lines))
            else:
                f.write_text(openai_key)
            log(f"Edited OpenAI key into file {f}")
            edited = True
        else:
            log(f"Did not locate file {f}")
    return edited

def popup_question():
    root = tk.Tk()
    root.withdraw()
    result = messagebox.askyesno("OpenAI Key", "Do you wish to install the course OpenAI key?\n\n(times out to YES after 120 seconds)")
    root.destroy()
    return result

def popup_info(msg):
    root = tk.Tk()
    root.withdraw()
    messagebox.showinfo("OpenAI Key", msg)
    root.destroy()

def main(show_popups=True):
    log(f"Starting OpenAI_Import V2.0 -- {time.ctime()}")
    metadata = get_metadata()
    if not metadata:
        return

    class_id = metadata.get("class-id")
    if not class_id or class_id == "null":
        log("Error: No class-id found in metadata, exiting.")
        return
    log(f"Class-id pulled: {class_id}")

    openai_key = metadata.get("custom-properties", {}).get("openAIKey")
    if not openai_key or openai_key == "null":
        client_id = "538GACRZ9FABA07E"
        client_secret = "bkAyYSB3dL8GREF87zfqhjkUcW5xNSVyj0HDpQBqheX5vcHpbLE8WXqYjGmGMORt"
        env_name = hmac_auth_request(class_id, client_id, client_secret)
        log(f"Class name pulled: {env_name}")

        token, token_type, expires_in = get_oauth_token(client_id, client_secret)
        if not token:
            log("Error: Failed to obtain OAuth2 token.")
            return
        log(f"Token obtained, type={token_type}, expires_in={expires_in}")

        openai_key = get_openai_key(env_name, token, token_type)
        log(f"Pulled OpenAI key: {openai_key}")
    else:
        log(f"OpenAI key pulled from metadata custom-properties: {openai_key}")

    if not openai_key:
        log("No key provided in VM metadata")
        return

    if show_popups:
        if not popup_question():
            log("Exited without installing the key")
            return

    edited = update_files(openai_key)
    if show_popups and edited:
        popup_info("The OpenAI key has been updated")

    log("All listed files have been edited if found, exiting.")

if __name__ == "__main__":
    main(show_popups=True)