import subprocess
import json
import time
import os

instance_id = "i-0e4eab5fe67cf9abe"

def run_cmd():
    # Sends a bash command to get folder size and the latest git speed log
    bash_script = "du -sh /home/ubuntu/models/Qwen-Image-2512 2>/dev/null; echo ''; tail -n 8 /home/ubuntu/qwen_git_download.log"
    
    send_req = subprocess.run([
        "aws", "ssm", "send-command",
        "--instance-ids", instance_id,
        "--document-name", "AWS-RunShellScript",
        "--parameters", 'commands=["' + bash_script + '"]',
        "--output", "json"
    ], capture_output=True, text=True)
    
    try:
        cmd_id = json.loads(send_req.stdout)["Command"]["CommandId"]
    except Exception as e:
        return f"Failed to send SSM command: {send_req.stderr}"
        
    time.sleep(4) # Wait to allow the command to execute on AWS
    
    res = subprocess.run([
        "aws", "ssm", "get-command-invocation",
        "--command-id", cmd_id,
        "--instance-id", instance_id,
        "--output", "json"
    ], capture_output=True, text=True)
    
    try:
        out = json.loads(res.stdout)
        if out["Status"] in ["Pending", "InProgress"]:
            return "Command still executing. Will retry next loop..."
        return out.get("StandardOutputContent", "") + out.get("StandardErrorContent", "")
    except:
        return "Waiting for AWS to return output..."

if __name__ == "__main__":
    print("Initializing AWS SSM connection...")
    while True:
        output = run_cmd()
        os.system('cls' if os.name == 'nt' else 'clear')
        print("======== LIVE DOWNLOAD MONITOR ========")
        print("Model: Qwen/Qwen-Image-2512 (40.86 GB)")
        print("=======================================\n")
        print(output.strip())
        print("\n=======================================")
        print("Refreshing every 10 seconds... (Press Ctrl+C to stop)")
        time.sleep(10)