import runpod
import base64
import os
from rich import print

def create(name: str, region: str):
    with open(os.path.expanduser('~/.ssh/id_rsa.pub'), 'r', encoding='utf-8') as f:
        public_key = f.read().strip()
    setup_cmd = (
        # Setting up SSH here
        'prefix_cmd() '
        '{ if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }; '
        '$(prefix_cmd) apt update;'
        'export DEBIAN_FRONTEND=noninteractive;'
        '$(prefix_cmd) apt install openssh-server rsync curl patch -y;'
        '$(prefix_cmd) mkdir -p /var/run/sshd; '
        '$(prefix_cmd) '
        'sed -i "s/PermitRootLogin prohibit-password/PermitRootLogin yes/" '
        '/etc/ssh/sshd_config; '
        '$(prefix_cmd) sed '
        '"s@session\\s*required\\s*pam_loginuid.so@session optional '
        'pam_loginuid.so@g" -i /etc/pam.d/sshd; '
        'cd /etc/ssh/ && $(prefix_cmd) ssh-keygen -A; '
        '$(prefix_cmd) mkdir -p ~/.ssh; '
        '$(prefix_cmd) chown -R $(whoami) ~/.ssh;'
        '$(prefix_cmd) chmod 700 ~/.ssh; '
        f'$(prefix_cmd) echo "{public_key}" >> ~/.ssh/authorized_keys; '
        '$(prefix_cmd) chmod 644 ~/.ssh/authorized_keys; '
        '$(prefix_cmd) service ssh restart; '
        '[ $(id -u) -eq 0 ] && echo alias sudo="" >> ~/.bashrc;'
        # Starting a test HTTP server
        'python3 -m http.server 9000'
    )
    encoded = base64.b64encode(setup_cmd.encode('utf-8')).decode('utf-8')
    pod = runpod.create_pod(
        name=name,
        image_name="runpod/base:0.0.2",
        gpu_type_id="NVIDIA RTX A4000",
        country_code=region,
        ports="22/tcp,9000/tcp",
        support_public_ip=True,
        docker_args=f'bash -c \'echo {encoded} | base64 --decode > init.sh; bash init.sh\''
    )
    return pod['id']

rp1_id = create("rp1", "CA")
rp2_id = create("rp2", "CA")

print(f"rp1_id = '{rp1_id}'")
print(f"rp2_id = '{rp2_id}'")

def get_cmd(pod_id: str):
    pod_stat = runpod.get_pod(pod_id)
    runtime = pod_stat.get('runtime') or {}
    ports_info = runtime.get('ports', [])
    if not ports_info:
        raise ValueError(f"Pod {pod_id} is not ready.")
    ssh_cmd = None
    curl_cmd = None
    for p in ports_info:
        if p['isIpPublic']:
            if p['privatePort'] == 22:
                ssh_cmd = f'ssh -i ~/.ssh/id_rsa -p {p["publicPort"]} root@{p["ip"]}'
            if p['privatePort'] == 9000:
                curl_cmd = f'curl http://{p["ip"]}:{p["publicPort"]}'
    assert ssh_cmd is not None and curl_cmd is not None, f"Pod {pod_id} is not ready."
    return ssh_cmd, curl_cmd

# Fill in the pod id retrieved from previous script
rp1_id = 'qi5a6pnu01x2zl'
rp2_id = '3k3hy87mtr2old'

rp1_ssh, rp1_curl = get_cmd(rp1_id)
rp2_ssh, rp2_curl = get_cmd(rp2_id)

print(rp1_curl)
print(rp2_curl)

print(f'{rp1_ssh} {rp2_curl}')
print(f'{rp2_ssh} {rp1_curl}')