"""Пускач PLUS Studio: сервер + публичная ссылка.

Одна строка в Colab или Kaggle:

    !python3 plus/launch_studio.py

Что делает:
  1. ставит зависимости (fastapi, uvicorn, python-multipart, soundfile, pretty_midi, ...);
  2. если рядом нет plus_app.py — скачивает plus.zip с Hugging Face и распаковывает;
  3. поднимает plus_app.py на локальном порту и ждёт /health;
  4. перебирает туннели по очереди (localhost.run, serveo, pinggy) и берёт первый,
     который выдал адрес и реально ответил;
  5. печатает ссылку, кладёт её в plus_links.txt и держит всё живым,
     перезапуская сервер или туннель, если они умерли.

Флаги:
  --port 8900          локальный порт
  --device auto        auto | cuda | cpu
  --weights PATH       свои веса; иначе модель скачается с HF при первом запросе
  --no-tunnel          только сервер, без ссылки (для локального запуска)
  --tunnel NAME        начать с конкретного туннеля: localhostrun | serveo | pinggy
  --keep 0             сколько секунд держать (0 = бесконечно)
"""

from __future__ import annotations

import argparse
import collections
import os
import re
import shutil
import subprocess
import sys
import threading
import time
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = "Nokia123543/muscriptor-large-fork"
ZIP_URL = "https://huggingface.co/%s/resolve/main/plus.zip" % REPO
PACKAGES = ("fastapi", "uvicorn", "python-multipart", "soundfile", "pretty_midi",
            "librosa", "huggingface_hub", "numpy", "scipy")

LINKS_FILE = "plus_links.txt"


def step(index: int, total: int, text: str) -> None:
    print("[%d/%d] %s" % (index, total, text), flush=True)


def ensure_deps() -> None:
    missing = []
    for name in PACKAGES:
        import importlib.util
        if importlib.util.find_spec(name.replace("-", "_")) is None:
            missing.append(name)
    if not missing:
        return
    print("ставлю: " + ", ".join(missing), flush=True)
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", *missing], check=False)


def ensure_pack() -> str:
    target = os.path.join(HERE, "plus_app.py")
    if os.path.isfile(target):
        return HERE
    print("скачиваю пакет с Hugging Face", flush=True)
    import zipfile
    archive = os.path.join(HERE, "plus.zip")
    urllib.request.urlretrieve(ZIP_URL, archive)
    with zipfile.ZipFile(archive) as bundle:
        bundle.extractall(HERE)
    for candidate in (os.path.join(HERE, "plus"), HERE):
        if os.path.isfile(os.path.join(candidate, "plus_app.py")):
            return candidate
    raise SystemExit("в архиве нет plus_app.py")


def healthy(port: int) -> bool:
    try:
        with urllib.request.urlopen("http://127.0.0.1:%d/health" % port, timeout=3) as answer:
            return answer.status == 200
    except Exception:
        return False


def start_server(pack: str, port: int, device: str, weights: str | None) -> subprocess.Popen:
    command = [sys.executable, os.path.join(pack, "plus_app.py"),
               "--port", str(port), "--device", device, "--host", "0.0.0.0"]
    if weights:
        command += ["--weights", weights]
    log = open("plus_studio.log", "a")
    log.write("\n=== %s ===\n" % time.strftime("%Y-%m-%d %H:%M:%S"))
    process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT,
                               stdin=subprocess.DEVNULL)
    for _ in range(90):
        if healthy(port):
            return process
        if process.poll() is not None:
            raise SystemExit("сервер умер при старте, смотри plus_studio.log")
        time.sleep(2)
    raise SystemExit("сервер не ответил за 3 минуты, смотри plus_studio.log")


class Tunnel:
    """Туннель по ssh -R или по http-клиенту; нужен только адрес в stdout."""

    def __init__(self, name: str, argv: list[str], pattern: str):
        self.name, self.pattern, self.tail = name, re.compile(pattern), collections.deque(maxlen=60)
        self.url: str | None = None
        self.proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                     stdin=subprocess.DEVNULL, text=True, bufsize=1)
        threading.Thread(target=self._drain, daemon=True).start()

    def _drain(self) -> None:
        for line in self.proc.stdout:
            self.tail.append(line.rstrip())
            if self.url is None:
                found = self.pattern.search(line)
                if found:
                    self.url = found.group(0)

    def alive(self) -> bool:
        return self.proc.poll() is None

    def kill(self) -> None:
        try:
            self.proc.kill()
        except Exception:
            pass


SSH_OPTS = ["-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
            "-o", "ServerAliveInterval=15", "-o", "ServerAliveCountMax=3",
            "-o", "TCPKeepAlive=yes", "-o", "ExitOnForwardFailure=yes"]


def tunnel_specs(port: int) -> list[tuple[str, list[str], str]]:
    target = "127.0.0.1:%d" % port
    specs = [("localhostrun",
              ["ssh", "-T", *SSH_OPTS, "-R", "80:" + target, "nokey@localhost.run"],
              r"https://[a-z0-9]+\.lhr\.life"),
             ("serveo",
              ["ssh", "-T", *SSH_OPTS, "-R", "80:" + target, "serveo.net"],
              r"https://[a-z0-9-]+(?:\.[a-z0-9-]+)*\.(?:serveousercontent|serveo)\.[a-z]+")]
    if shutil.which("npx"):
        specs.append(("localtunnel",
                      ["npx", "-y", "localtunnel", "--port", str(port)],
                      r"https://[a-z0-9-]+\.loca\.lt"))
    return specs


def reachable(url: str, attempts: int = 5) -> bool:
    for _ in range(attempts):
        try:
            with urllib.request.urlopen(url + "/health", timeout=8) as answer:
                if answer.status == 200:
                    return True
        except Exception:
            pass
        time.sleep(3)
    return False


def open_tunnel(port: int, order: list[str]) -> Tunnel:
    specs = {name: (argv, pattern) for name, argv, pattern in tunnel_specs(port)}
    tried = []
    for name in order:
        if name not in specs:
            continue
        argv, pattern = specs[name]
        if not shutil.which(argv[0]):
            print("  [%s] нет команды %s, пропускаю" % (name, argv[0]), flush=True)
            continue
        print("  [%s] открываю" % name, flush=True)
        tunnel = Tunnel(name, argv, pattern)
        deadline = time.time() + 60
        started = time.time()
        while time.time() < deadline and tunnel.url is None and tunnel.alive():
            time.sleep(1)
            waited = int(time.time() - started)
            if waited and waited % 10 == 0:
                print("      [%s] ждём адрес, %d с, жив: %s"
                      % (name, waited, tunnel.alive()), flush=True)
        if tunnel.url and reachable(tunnel.url):
            print("  [%s] работает: %s" % (name, tunnel.url), flush=True)
            return tunnel
        tried.append((name, tunnel.url, list(tunnel.tail)[-3:]))
        tunnel.kill()
    for name, url, tail in tried:
        print("  [%s] не вышло: %s | %s" % (name, url or "нет адреса", " / ".join(tail)), flush=True)
    raise SystemExit("ни один туннель не поднялся: ssh или npx недоступны из этой среды")


def write_links(lines: list[str]) -> None:
    text = "\n".join(lines) + "\n"
    with open(LINKS_FILE, "w", encoding="utf-8") as handle:
        handle.write(text)
    print("\n" + text, flush=True)


def main(argv: list[str] | None = None) -> int:
    print("логи: %s (сервер) и %s (ссылки) в %s"
          % ("plus_studio.log", LINKS_FILE, os.getcwd()), flush=True)
    parser = argparse.ArgumentParser(description="PLUS Studio: сервер и ссылка")
    parser.add_argument("--port", type=int, default=8900)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--weights", default=os.environ.get("MUSCRIPTOR_WEIGHTS"))
    parser.add_argument("--no-tunnel", action="store_true")
    parser.add_argument("--tunnel", default=None)
    parser.add_argument("--keep", type=int, default=0, help="секунд держать, 0 = бесконечно")
    args = parser.parse_args(argv)

    step(1, 5, "зависимости")
    ensure_deps()
    step(2, 5, "пакет PLUS")
    pack = ensure_pack()
    step(3, 5, "сервер на порту %d" % args.port)
    server = start_server(pack, args.port, args.device, args.weights)
    print("      сервер отвечает: http://127.0.0.1:%d" % args.port, flush=True)

    order = ["localhostrun", "serveo", "local" "tunnel"]
    if args.tunnel:
        order = [args.tunnel] + [name for name in order if name != args.tunnel]
    tunnel = None
    if args.no_tunnel:
        write_links(["PLUS Studio " + time.strftime("%H:%M:%S"),
                     "локально: http://127.0.0.1:%d" % args.port])
        print(">>> LOCAL LINK: http://127.0.0.1:%d" % args.port, flush=True)
    else:
        step(4, 5, "публичная ссылка")
        tunnel = open_tunnel(args.port, order)
        write_links(["PLUS Studio " + time.strftime("%H:%M:%S"),
                     "ссылка: " + tunnel.url,
                     "локально: http://127.0.0.1:%d" % args.port])
        if tunnel.url:
            print("\n>>> OPEN THIS LINK: " + tunnel.url, flush=True)

    step(5, 5, "держу живым, Stop останавливает")
    started = time.time()
    fails = 0
    beat = time.time()
    try:
        while True:
            time.sleep(10)
            if time.time() - beat > 60:
                beat = time.time()
                print("      жив: %.1f мин | сервер %s | ссылка %s"
                      % ((time.time() - started) / 60,
                         "ок" if server.poll() is None else "упал",
                         tunnel.url if tunnel else "нет (локально)"), flush=True)
            if args.keep and time.time() - started > args.keep:
                break
            if server.poll() is not None:
                print("сервер упал, поднимаю заново", flush=True)
                server = start_server(pack, args.port, args.device, args.weights)
            if tunnel is None:
                continue
            if not tunnel.alive():
                print("туннель умер, открываю новый", flush=True)
                tunnel.kill()
                tunnel = open_tunnel(args.port, order)
                write_links(["PLUS Studio " + time.strftime("%H:%M:%S"),
                             "новая ссылка: " + tunnel.url])
                if tunnel.url:
                    print(">>> OPEN THIS LINK: " + tunnel.url, flush=True)
                fails = 0
                continue
            if not reachable(tunnel.url, attempts=1):
                fails += 1
                if fails >= 3:
                    print("ссылка перестала отвечать, переоткрываю", flush=True)
                    tunnel.kill()
                    fails = 0
    except KeyboardInterrupt:
        pass
    if tunnel is not None:
        tunnel.kill()
    if server.poll() is None:
        server.terminate()
    print("остановлено", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())