import os
import subprocess
import threading
import time
import requests
import secrets
import string
from jupyter_server.auth import passwd

# ====================
# CONFIG
# =====================
PORT = int(os.environ.get("PORT", 7860))
KEEP_ALIVE_URL = "https://srinetwork-srinet1.hf.space"
INTERVAL = 30 * 60

# 🔑 HF Variable / Secret (PLAIN PASSWORD)
ENV_PASSWORD = os.environ.get("JUPYTER_PASSWORD")
# =====================


def run_cmd(cmd):
    print(f"📦 {cmd}")
    subprocess.run(cmd, shell=True, check=False)


def install_packages():
    print("🔧 Installing system packages...")
    run_cmd("apt-get update")
    run_cmd("apt-get install -y sudo curl wget openjdk-21-jdk")
    run_cmd("java -version")


def keep_alive_loop():
    while True:
        try:
            r = requests.get(KEEP_ALIVE_URL, timeout=10)
            print(f"🔁 Keep-alive OK | {r.status_code}")
        except Exception as e:
            print(f"⚠️ Keep-alive error: {e}")
        time.sleep(INTERVAL)


def generate_password(length=14):
    chars = string.ascii_letters + string.digits
    return ''.join(secrets.choice(chars) for _ in range(length))


def start_jupyterlab():
    # 🔐 Decide password source
    if ENV_PASSWORD and ENV_PASSWORD.strip():
        plain_password = ENV_PASSWORD.strip()
        source = "HF VARIABLE"
    else:
        plain_password = generate_password()
        source = "RANDOM (fallback)"

    password_hash = passwd(plain_password)

    print("\n" + "=" * 60)
    print("🔐 JUPYTERLAB LOGIN INFO")
    print(f"🔑 Password Source : {source}")
    print(f"👉 PASSWORD        : {plain_password}")
    print("⚠️ Save this password (restart = may change)")
    print("=" * 60 + "\n")

    cmd = [
        "jupyter",
        "lab",
        "--ip=0.0.0.0",
        f"--port={PORT}",
        "--no-browser",
        "--allow-root",

        # ✅ Correct auth flags
        "--ServerApp.token=''",
        f"--ServerApp.password='{password_hash}'",
        "--ServerApp.disable_check_xsrf=True",
        "--ServerApp.allow_origin='*'"
    ]

    print("🚀 Starting JupyterLab...")
    subprocess.run(" ".join(cmd), shell=True)


# =====================
# MAIN
# =====================
if __name__ == "__main__":
    install_packages()

    threading.Thread(
        target=keep_alive_loop,
        daemon=True
    ).start()

    start_jupyterlab()