import os
import time
import subprocess
import sys
from google import genai
from google.genai import types, errors

# --- 1. CONFIGURATION ---
API_KEYS = [os.environ.get("AIzaSyB2lVccgq6Vbk0B__jcAqPeDDTEPzc9FqY"), os.environ.get("AIzaSyBavbqf2UxtUzHQ1zyuYnfl9Hyy1CIdMpQ")]
MODEL_ID = "gemini-2.0-flash" 

current_key_index = 0

def get_client():
    key = API_KEYS[current_key_index]
    if not key:
        print("\033[91m[Error] No API Key found. Set Z_KEY_1 in your environment variables.\033[0m")
        sys.exit(1)
    return genai.Client(api_key=key)

client = get_client()

# --- 2. UI UTILITIES ---
def typewriter(text: str, delay: float = 0.04):
    for char in text:
        sys.stdout.write(char)
        sys.stdout.flush()
        time.sleep(delay)
    print()

def run_startup_animation():
    """Starts instantly without asking for input."""
    print("\033[94m--- Zassistant - easiest Assistant in CMD ---\033[0m")
    time.sleep(0.5)
    
    steps = [
        ("Loading Modules", "OK"),
        ("Establishing Connection", "OK"),
        ("Securing Shell Access", "READY")
    ]

    for label, status in steps:
        for i in range(0, 101, 10):
            chars = ["|", "/", "-", "\\"]
            for char in chars:
                sys.stdout.write(f"\r{label}... {char} {i}%")
                sys.stdout.flush()
                time.sleep(0.05)
        sys.stdout.write(f"\r{label}... {status}               \n")
    print("\033[92m[System Online - How can I help?]\033[0m\n")

# --- 3. THE SYSTEM TOOL ---
def sys_action(cmd: str) -> str:
    try:
        res = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=15)
        return (res.stdout + res.stderr).strip() or "Done."
    except Exception as e:
        return f"Error: {str(e)}"

# --- 4. MAIN PROGRAM ---
def main():
    global current_key_index, client
    
    # 1. RUN ANIMATION IMMEDIATELY
    run_startup_animation()

    # 2. CONFIGURE AI
    config = types.GenerateContentConfig(
        tools=[sys_action],
        automatic_function_calling=True,
        system_instruction="You are Zassistant. Use sys_action for all commands. Be brief."
    )
    chat = client.chats.create(model=MODEL_ID, config=config)

    # 3. START LISTENING TO USER
    while True:
        try:
            user_input = input("\033[92mUser > \033[0m").strip()
            
            if not user_input: continue
            if user_input.lower() in ["exit", "shutdown", "end", "quit"]:
                print("\033[94m[Shutting Down]\033[0m")
                break

            success = False
            while not success:
                try:
                    response = chat.send_message(user_input)
                    if response.text:
                        print("\n\033[1mZassistant:\033[0m ", end="")
                        typewriter(response.text)
                    success = True
                except errors.ClientError as e:
                    if "429" in str(e) and current_key_index + 1 < len(API_KEYS) and API_KEYS[current_key_index + 1]:
                        current_key_index += 1
                        print(f"\n\033[93m[Quota Full] Switching Key...\033[0m")
                        client = get_client()
                        chat = client.chats.create(model=MODEL_ID, config=config)
                    else:
                        print(f"\n\033[91mAPI Error: {e}\033[0m")
                        success = True
        except KeyboardInterrupt:
            break

if __name__ == "__main__":
    main()