import os
import sys
import time
import json
import subprocess
import runpod

if sys.platform.startswith("win"):
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

runpod.api_key = "rpa_5SFMBV8T5MB9XG5SSQ615M43WUUN72F1NOVEJNEU1je8b4"
SSH_KEY = os.path.expanduser(r"~/.ssh/id_ed25519")
SESSION_FILE = os.path.abspath(r"c:\Projects\Historical\runpod_session.json")

print("==========================================================")
print("DEPLOYING RUNPOD TENSORRT 4K PURESCALE ENGINE (EP3 - EP8)")
print("==========================================================")

# 1. Tìm và cấp phát GPU RTX 4090
gpu_candidates = [
    ("NVIDIA GeForce RTX 4090", "COMMUNITY"),
    ("NVIDIA GeForce RTX 4090", "SECURE"),
    ("NVIDIA RTX 4000 Ada Generation", "COMMUNITY"),
    ("NVIDIA GeForce RTX 3090", "COMMUNITY")
]

pod_info = None
for gpu_id, cloud in gpu_candidates:
    try:
        print(f"[*] Thử cấp phát Pod: {gpu_id} ({cloud} Cloud)...")
        pod = runpod.create_pod(
            name="Historical-PureScale-4K",
            image_name="runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04",
            gpu_type_id=gpu_id,
            cloud_type=cloud,
            gpu_count=1,
            volume_in_gb=0,           # KHÔNG TẠO PERSISTENT VOLUME -> PHÍ STORAGE = 0Đ
            container_disk_in_gb=50,  # 50GB ổ đĩa tạm thời
            ports="8888/http,22/tcp",
            start_ssh=True
        )
        if pod and "id" in pod:
            pod_info = pod
            print(f"✓ CẤP PHÁT POD THÀNH CÔNG! Pod ID: {pod['id']} ({gpu_id} - {cloud})")
            break
    except Exception as e:
        print(f"[-] Không thể tạo Pod {gpu_id} ({cloud}): {e}")

if not pod_info:
    print("[X] Không tìm thấy máy ảo GPU khả dụng lúc này trên Runpod.")
    sys.exit(1)

pod_id = pod_info["id"]

# 2. Chờ Pod khởi động và lấy cổng SSH
print(f"\n[2/4] Đang chờ Pod {pod_id} khởi động container...")
ssh_host = None
ssh_port = None

for attempt in range(40):
    time.sleep(5)
    p = runpod.get_pod(pod_id)
    runtime = p.get("runtime")
    if runtime and runtime.get("uptimeInSeconds", 0) > 0:
        ports = runtime.get("ports", [])
        for pt in ports:
            if pt.get("privatePort") == 22:
                ssh_host = pt.get("ip")
                ssh_port = pt.get("publicPort")
                break
        if ssh_host and ssh_port:
            print(f"✓ Pod đã online! SSH: root@{ssh_host}:{ssh_port} (Uptime: {runtime.get('uptimeInSeconds')}s)")
            break

if not ssh_host or not ssh_port:
    print("[-] Không lấy được thông tin SSH từ Pod. Đang hủy Pod để tránh mất tiền...")
    runpod.terminate_pod(pod_id)
    sys.exit(1)

# Lưu thông tin session
session_data = {
    "pod_id": pod_id,
    "host": ssh_host,
    "port": ssh_port,
    "episodes": ["ep3", "ep4", "ep5", "ep6", "ep7", "ep8"],
    "created_at": time.strftime("%Y-%m-%d %H:%M:%S")
}
with open(SESSION_FILE, "w", encoding="utf-8") as f:
    json.dump(session_data, f, indent=2)

# Chờ SSH sẵn sàng
print("[*] Đang kiểm tra kết nối SSH vào Pod...")
for _ in range(12):
    time.sleep(3)
    cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
        "-o", "ConnectTimeout=5",
        "-i", SSH_KEY, "-p", str(ssh_port), f"root@{ssh_host}", "echo ready"
    ]
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if "ready" in res.stdout:
        print("✓ Kết nối SSH thành công!")
        break

# 3. Tải lên mã nguồn và mô hình ONNX
print("\n[3/4] Đang tải lên scripts và mô hình ONNX (purescale_4k_end2end.onnx)...")
scp_base = [
    "scp", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
    "-P", str(ssh_port), "-i", SSH_KEY
]

files_to_upload = [
    ("c:\\Projects\\Historical\\build_engine.py", "/workspace/build_engine.py"),
    ("c:\\Projects\\Historical\\trt_pipeline.py", "/workspace/trt_pipeline.py"),
    ("c:\\Projects\\Historical\\setup_pod.sh", "/workspace/setup_pod.sh"),
    ("c:\\Projects\\Historical\\models\\purescale_4k_end2end.onnx", "/workspace/purescale_4k_end2end.onnx")
]

for local_f, remote_f in files_to_upload:
    print(f"[*] Uploading {os.path.basename(local_f)} ({os.path.getsize(local_f)/(1024*1024):.1f} MB)...", flush=True)
    subprocess.run(scp_base + [local_f, f"root@{ssh_host}:{remote_f}"], check=True)
print("✓ Toàn bộ file đã được tải lên Pod!")

# 4. Kích hoạt setup và build TensorRT engine
print("\n[4/4] Khởi chạy cài đặt và biên dịch TensorRT FP16 Engine trên RTX 4090...")
ssh_cmd = [
    "ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
    "-i", SSH_KEY, "-p", str(ssh_port), f"root@{ssh_host}",
    "chmod +x /workspace/setup_pod.sh && bash /workspace/setup_pod.sh"
]
subprocess.run(ssh_cmd, check=True)

print("\n" + "=" * 60)
print("🎉 SETUP THÀNH CÔNG VÀ TIẾN TRÌNH TRT 4K ĐANG CHẠY TRÊN RUNPOD!")
print(f"Pod ID: {pod_id}")
print("Bây giờ bạn có thể kích hoạt 'python watchdog_runpod.py' để tự động giám sát và hủy Pod.")
print("=" * 60)