"""
Мост CodeWithMe: запускает код из уроков codewithme.app на этом компьютере.

Запуск:   python cwm_bridge.py        (Windows: py cwm_bridge.py)
Стоп:     Ctrl+C или закрыть окно терминала.

Как это защищено:
  * мост слушает только 127.0.0.1 — с других компьютеров к нему не подключиться;
  * запросы принимаются только со страниц https://codewithme.app (заголовок Origin
    браузер подделать не даёт) и только с адресом 127.0.0.1/localhost (защита от DNS rebinding);
  * запускать код можно только после сопряжения: мост печатает здесь код из 6 цифр,
    его вводят на сайте, и сайт получает секретный токен. После перезапуска моста — новый код;
  * сайт может только передать программу и входные данные: ни путей, ни команд, ни аргументов;
  * каждый запуск виден в этом окне. Мост не песочница: программа работает с вашими правами,
    как если бы вы запустили её сами. Держите мост включённым, только пока занимаетесь.

Нужен только Python 3.8+ — сторонние библиотеки мосту не нужны.
"""

import argparse
import importlib.util
import json
import os
import platform
import re
import secrets
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

VERSION = "1.0.0"
PROTOCOL = 1
PORTS = range(47615, 47620)
SITE = "https://codewithme.app"
DEV_ORIGINS = ("http://localhost:3000", "http://127.0.0.1:3000")

MAX_BODY = 20 << 20  # программа + входные данные (большие тесты — до 8 МБ)
MAX_OUTPUT = 16 << 20
MAX_TIME_MS = 10 * 60 * 1000  # «Запустить у себя» с окном игры — до 10 минут
STARTUP_GRACE_S = 3.0  # запуск интерпретатора и импорт библиотек не входят во время программы
MAX_PARALLEL = 2
PAIR_ATTEMPTS = 5
PAIR_TTL_S = 10 * 60
MODULE_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$")
RUN_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$")

# Выполняется в отдельном процессе Python. Повторяет правила браузерной проверки сайта:
# sys.exit(0) — нормальное завершение, ненулевой код и исключения — ошибка выполнения,
# трейсбек — только по строкам программы ученика.
DRIVER = r'''
import json, sys, time, traceback
src_path, check_path, result_path, name = sys.argv[1:5]
sys.stdout.reconfigure(newline="\n")
sys.stderr.reconfigure(newline="\n")
src = open(src_path, encoding="utf-8").read()
check = open(check_path, encoding="utf-8").read() if check_path != "-" else None
g = {"__name__": name, "__builtins__": __builtins__}
status, error = "ok", None
started = time.perf_counter()
try:
    exec(compile(src, "main.py", "exec"), g)
    if check:
        try:
            exec(compile(check, "<check>", "exec"), g)
        except Exception as e:
            status, error = "re", "Проверка задачи не смогла прочитать результат: %s: %s" % (type(e).__name__, e)
except SystemExit as e:
    if e.code not in (None, 0):
        status, error = "re", "Программа завершилась с кодом %s (sys.exit)" % (e.code,)
except MemoryError:
    status, error = "re", "MemoryError: программе не хватило памяти"
except SyntaxError as e:
    status, error = "re", "".join(traceback.format_exception_only(type(e), e)).rstrip()
except BaseException as e:
    frames = [f for f in traceback.extract_tb(e.__traceback__) if f.filename == "main.py"]
    text = "Traceback (most recent call last):\n" + "".join(traceback.format_list(frames))
    status, error = "re", (text + "".join(traceback.format_exception_only(type(e), e))).rstrip()
elapsed = time.perf_counter() - started
try:
    sys.stdout.flush()
    sys.stderr.flush()
except BaseException:
    pass
with open(result_path, "w", encoding="utf-8") as f:
    json.dump({"status": status, "error": error, "timeMs": round(elapsed * 1000)}, f)
'''


def log(msg):
    print(time.strftime("%H:%M:%S"), msg, flush=True)


class Pairing:
    """Код сопряжения из 6 цифр и выданные токены. Живут, пока работает мост."""

    def __init__(self):
        self.lock = threading.Lock()
        self.tokens = set()
        self.locked = False
        self.new_code()

    def new_code(self):
        self.code = "".join(secrets.choice("0123456789") for _ in range(6))
        self.code_until = time.time() + PAIR_TTL_S
        self.attempts = 0
        print()
        print("  Код для сайта CodeWithMe:  %s %s" % (self.code[:3], self.code[3:]))
        print("  Введите его на странице урока. Код действует 10 минут.")
        print(flush=True)

    def pair(self, code):
        with self.lock:
            if self.locked:
                return None, "locked"
            if time.time() > self.code_until:
                self.new_code()
                return None, "expired"
            if not secrets.compare_digest(str(code).replace(" ", ""), self.code):
                self.attempts += 1
                if self.attempts >= PAIR_ATTEMPTS:
                    self.locked = True
                    log("Слишком много неверных кодов. Сопряжение заблокировано — перезапустите мост.")
                    return None, "locked"
                return None, "wrong"
            token = secrets.token_urlsafe(32)
            self.tokens.add(token)
            log("Сайт подключён.")
            self.new_code()  # код одноразовый
            return token, None

    def valid(self, token):
        with self.lock:
            return any(secrets.compare_digest(token, t) for t in self.tokens)


class Runs:
    def __init__(self):
        self.lock = threading.Lock()
        self.procs = {}
        self.cancelled = set()
        self.slots = threading.Semaphore(MAX_PARALLEL)

    def cancel(self, run_id):
        with self.lock:
            proc = self.procs.get(run_id)
            if proc is None:
                return False
            self.cancelled.add(run_id)
        if proc.poll() is None:
            proc.kill()
        return True


def read_limited(stream, chunks, state, limit, proc):
    """Читает поток процесса; при превышении общего лимита вывода завершает процесс (OLE)."""
    while True:
        data = stream.read(65536)
        if not data:
            return
        with state["lock"]:
            state["total"] += len(data)
            if state["total"] > limit:
                state["over"] = True
                try:
                    proc.kill()
                except OSError:
                    pass
                return
            chunks.append(data)


def run_program(req, runs):
    run_id = req["runId"]
    tmp = tempfile.mkdtemp(prefix="cwm-")
    try:
        main_path = os.path.join(tmp, "main.py")
        with open(main_path, "w", encoding="utf-8") as f:
            f.write(req["code"])
        check_path = "-"
        if req.get("check"):
            check_path = os.path.join(tmp, "_cwm_check.py")
            with open(check_path, "w", encoding="utf-8") as f:
                f.write(req["check"])
        driver_path = os.path.join(tmp, "_cwm_driver.py")
        with open(driver_path, "w", encoding="utf-8") as f:
            f.write(DRIVER)
        result_path = os.path.join(tmp, "_cwm_result.json")

        env = dict(os.environ)
        for k in list(env):
            if k.startswith("PYTHON"):
                del env[k]
        env["PYTHONIOENCODING"] = "utf-8"
        env["PYGAME_HIDE_SUPPORT_PROMPT"] = "1"  # иначе pygame печатает приветствие в вывод программы
        if not req.get("window"):
            # Проверка без окон и звука
            env["SDL_VIDEODRIVER"] = "dummy"
            env["SDL_AUDIODRIVER"] = "dummy"
            env["MPLBACKEND"] = "Agg"
        # При проверке с check игровой цикл под if __name__ == "__main__" не запускается
        name = "__cwm_check__" if req.get("check") else "__main__"

        cmd = [sys.executable, "-E", "-X", "utf8", driver_path, main_path, check_path, result_path, name]
        flags = subprocess.CREATE_NO_WINDOW if os.name == "nt" and not req.get("window") else 0
        limit_s = req["timeLimitMs"] / 1000
        with runs.slots:
            proc = subprocess.Popen(
                cmd, cwd=tmp, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, creationflags=flags
            )
            with runs.lock:
                runs.procs[run_id] = proc
            out, err = [], []
            state = {"lock": threading.Lock(), "total": 0, "over": False}
            readers = [
                threading.Thread(target=read_limited, args=(proc.stdout, out, state, req["outputLimitBytes"], proc), daemon=True),
                threading.Thread(target=read_limited, args=(proc.stderr, err, state, req["outputLimitBytes"], proc), daemon=True),
            ]
            for t in readers:
                t.start()
            def feed():
                try:
                    proc.stdin.write(req["stdin"].encode("utf-8"))
                    proc.stdin.close()
                except OSError:
                    pass  # программа завершилась, не дочитав ввод

            # Ввод пишем в отдельном потоке: если программа его не читает, запись в pipe блокируется,
            # а лимит времени (proc.wait ниже) должен действовать всё равно
            threading.Thread(target=feed, daemon=True).start()
            timed_out = False
            try:
                proc.wait(timeout=limit_s + STARTUP_GRACE_S)
            except subprocess.TimeoutExpired:
                timed_out = True
                proc.kill()
                proc.wait()
            for t in readers:
                t.join(5)
            with runs.lock:
                runs.procs.pop(run_id, None)
                cancelled = run_id in runs.cancelled
                runs.cancelled.discard(run_id)

        stdout = b"".join(out).decode("utf-8", "replace")
        stderr = b"".join(err).decode("utf-8", "replace")
        if cancelled:
            return {"status": "cancelled", "stdout": "", "stderr": "", "timeMs": 0}
        if state["over"]:
            return {"status": "ole", "stdout": "", "stderr": "", "timeMs": 0}
        if timed_out:
            return {"status": "tle", "stdout": stdout, "stderr": stderr, "timeMs": req["timeLimitMs"]}
        try:
            with open(result_path, encoding="utf-8") as f:
                res = json.load(f)
        except (OSError, ValueError):
            return {
                "status": "re",
                "stdout": stdout,
                "stderr": stderr,
                "timeMs": 0,
                "error": "Программа аварийно завершилась (код %s)" % proc.returncode,
            }
        status = res["status"]
        if res["timeMs"] > req["timeLimitMs"]:
            status = "tle"
        out_res = {"status": status, "stdout": stdout, "stderr": stderr, "timeMs": res["timeMs"]}
        if res.get("error"):
            out_res["error"] = res["error"]
        return out_res
    finally:
        with runs.lock:
            runs.procs.pop(run_id, None)
            runs.cancelled.discard(run_id)
        shutil.rmtree(tmp, ignore_errors=True)


def make_handler(port, origins, pairing, runs):
    hosts = {"127.0.0.1:%d" % port, "localhost:%d" % port}

    class Handler(BaseHTTPRequestHandler):
        server_version = "cwm-bridge/" + VERSION
        sys_version = ""

        def log_message(self, *args):
            pass  # своё короткое журналирование вместо стандартного

        def allowed(self):
            origin = self.headers.get("Origin", "")
            if origin not in origins or self.headers.get("Host", "") not in hosts:
                self.send_json(403, {"error": "forbidden"}, cors=False)
                return None
            return origin

        def send_json(self, code, data, cors=True):
            body = json.dumps(data, ensure_ascii=False).encode("utf-8")
            self.send_response(code)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            if cors:
                self.send_header("Access-Control-Allow-Origin", self.headers.get("Origin", ""))
                self.send_header("Vary", "Origin")
            self.end_headers()
            self.wfile.write(body)

        def read_json(self):
            try:
                length = int(self.headers.get("Content-Length", "0"))
            except ValueError:
                return None
            if length <= 0 or length > MAX_BODY:
                return None
            try:
                data = json.loads(self.rfile.read(length).decode("utf-8"))
            except (ValueError, UnicodeDecodeError):
                return None
            return data if isinstance(data, dict) else None

        def authorized(self):
            token = self.headers.get("X-CWM-Token", "")
            if not token or not pairing.valid(token):
                self.send_json(401, {"error": "not-paired"})
                return False
            return True

        def do_OPTIONS(self):
            origin = self.allowed()
            if not origin:
                return
            self.send_response(204)
            self.send_header("Access-Control-Allow-Origin", origin)
            self.send_header("Access-Control-Allow-Methods", "GET, POST")
            self.send_header("Access-Control-Allow-Headers", "Content-Type, X-CWM-Token")
            self.send_header("Access-Control-Allow-Private-Network", "true")
            self.send_header("Access-Control-Max-Age", "600")
            self.send_header("Vary", "Origin")
            self.send_header("Content-Length", "0")
            self.end_headers()

        def do_GET(self):
            if not self.allowed():
                return
            path, _, query = self.path.partition("?")
            if path != "/status":
                self.send_json(404, {"error": "not-found"})
                return
            modules = {}
            # parse_qs раскодирует %2C: сайт присылает modules=pygame%2Cnumpy
            for value in urllib.parse.parse_qs(query).get("modules", []):
                for m in value.split(",")[:20]:
                    if MODULE_RE.match(m):
                        try:
                            modules[m] = importlib.util.find_spec(m) is not None
                        except (ImportError, ValueError):
                            modules[m] = False
            token = self.headers.get("X-CWM-Token", "")
            self.send_json(
                200,
                {
                    "bridge": PROTOCOL,
                    "version": VERSION,
                    "python": platform.python_version(),
                    "os": platform.system(),
                    "paired": bool(token) and pairing.valid(token),
                    "modules": modules,
                },
            )

        def do_POST(self):
            if not self.allowed():
                return
            if self.path == "/pair":
                data = self.read_json()
                if not data or not isinstance(data.get("code"), str) or len(data["code"]) > 20:
                    self.send_json(400, {"error": "bad-request"})
                    return
                token, error = pairing.pair(data["code"])
                if token:
                    self.send_json(200, {"token": token})
                else:
                    self.send_json(403, {"error": error})
                return
            if self.path not in ("/run", "/cancel"):
                self.send_json(404, {"error": "not-found"})
                return
            if not self.authorized():
                return
            data = self.read_json()
            if not data or not isinstance(data.get("runId"), str) or not RUN_ID_RE.match(data["runId"]):
                self.send_json(400, {"error": "bad-request"})
                return
            if self.path == "/cancel":
                self.send_json(200, {"cancelled": runs.cancel(data["runId"])})
                return
            ok = (
                isinstance(data.get("code"), str)
                and isinstance(data.get("stdin"), str)
                and (data.get("check") is None or isinstance(data.get("check"), str))
                and isinstance(data.get("timeLimitMs"), int)
                and 1 <= data["timeLimitMs"] <= MAX_TIME_MS
                and isinstance(data.get("outputLimitBytes"), int)
                and 1 <= data["outputLimitBytes"] <= MAX_OUTPUT
                and isinstance(data.get("window", False), bool)
            )
            if not ok:
                self.send_json(400, {"error": "bad-request"})
                return
            label = str(data.get("label", ""))[:80]
            mode = "окно" if data.get("window") else "проверка"
            log("Запуск: %s (%s)" % (label or "программа", mode))
            result = run_program(data, runs)
            log("  → %s, %s мс" % (result["status"].upper(), result.get("timeMs", 0)))
            self.send_json(200, result)

    return Handler


def main():
    parser = argparse.ArgumentParser(description="Мост CodeWithMe: запуск кода из уроков на этом компьютере")
    parser.add_argument("--dev", action="store_true", help="разрешить также локальную разработку сайта (localhost:3000)")
    args = parser.parse_args()
    origins = (SITE,) + (DEV_ORIGINS if args.dev else ())

    server = None
    for port in PORTS:
        try:
            server = ThreadingHTTPServer(("127.0.0.1", port), BaseHTTPRequestHandler)
            break
        except OSError:
            continue
    if server is None:
        print("Не удалось занять порт %d–%d: возможно, мост уже запущен в другом окне." % (PORTS[0], PORTS[-1]))
        sys.exit(1)

    print("Мост CodeWithMe %s · Python %s · %s" % (VERSION, platform.python_version(), sys.executable))
    print("Работает на 127.0.0.1:%d. Чтобы выключить — Ctrl+C или закройте это окно." % port)
    if args.dev:
        print("Режим разработки: принимаются запросы и с localhost:3000.")
    pairing = Pairing()
    server.RequestHandlerClass = make_handler(port, origins, pairing, Runs())
    server.daemon_threads = True
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nМост выключен.")
    finally:
        server.server_close()


if __name__ == "__main__":
    main()