#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Daneshyar — Media Server Manager
================================
Register and manage the TURN/SFU servers that carry online-class audio/video.
New classes are placed on the least-loaded active server, so adding a row here
is all that's needed to spread load across more machines.

Requires:  pip install pymysql paramiko

Usage:
    python scripts/media_servers.py list
    python scripts/media_servers.py migrate
    python scripts/media_servers.py add --name "سرور دوم" --host 1.2.3.4 --secret <hmac> [--capacity 50]
    python scripts/media_servers.py health
    python scripts/media_servers.py install --host 1.2.3.4 --ssh-user root [--ssh-port 22]
    python scripts/media_servers.py drain 2          # stop new classes landing on it
    python scripts/media_servers.py enable 2
    python scripts/media_servers.py remove 2

`install` SSHes to a fresh Ubuntu box, installs coturn, generates an HMAC
secret, opens the firewall, and prints the exact `add` command to run next.
"""
from __future__ import annotations

import argparse
import secrets
import socket
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

if hasattr(sys.stdout, "reconfigure"):
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

ROOT = Path(__file__).resolve().parents[1]

G = "\033[92m"; R = "\033[91m"; Y = "\033[93m"; B = "\033[94m"; DIM = "\033[2m"; RST = "\033[0m"


def read_env() -> dict[str, str]:
    values: dict[str, str] = {}
    env_path = ROOT / ".env"
    if not env_path.exists():
        print(f"{R}.env not found at {env_path}{RST}")
        sys.exit(1)
    for line in env_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        values[k] = v.strip().strip('"').strip("'")
    return values


def connect():
    try:
        import pymysql
    except ImportError:
        print(f"{R}pip install pymysql{RST}")
        sys.exit(1)

    env = read_env()
    return pymysql.connect(
        host=env["DB_HOST"],
        port=int(env.get("DB_PORT", "3306")),
        user=env["DB_USERNAME"],
        password=env["DB_PASSWORD"],
        database=env["DB_DATABASE"],
        charset="utf8mb4",
        autocommit=True,
        cursorclass=__import__("pymysql").cursors.DictCursor,
    )


SCHEMA = """
CREATE TABLE IF NOT EXISTS media_servers (
  id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
  name VARCHAR(120) NOT NULL,
  kind ENUM('turn','sfu') NOT NULL DEFAULT 'turn',
  host VARCHAR(190) NOT NULL,
  port INT UNSIGNED NOT NULL DEFAULT 3478,
  scheme VARCHAR(10) NOT NULL DEFAULT 'turn',
  secret VARCHAR(255) NULL,
  region VARCHAR(60) NULL,
  capacity INT UNSIGNED NOT NULL DEFAULT 50,
  current_load INT UNSIGNED NOT NULL DEFAULT 0,
  priority INT UNSIGNED NOT NULL DEFAULT 100,
  status ENUM('active','draining','offline') NOT NULL DEFAULT 'active',
  manually_disabled BOOLEAN NOT NULL DEFAULT FALSE,
  last_health_at TIMESTAMP NULL,
  last_health_ok BOOLEAN NOT NULL DEFAULT FALSE,
  last_health_note VARCHAR(255) NULL,
  created_at TIMESTAMP NULL,
  updated_at TIMESTAMP NULL,
  UNIQUE KEY uniq_media_server_host_port (host, port, kind),
  INDEX idx_media_server_pick (status, kind, current_load, priority)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
"""


def cmd_migrate(_args) -> None:
    """Create the table and the online_classes.media_server_id column."""
    conn = connect(); cur = conn.cursor()

    cur.execute(SCHEMA)
    print(f"{G}media_servers table ready{RST}")

    cur.execute("SHOW COLUMNS FROM online_classes LIKE 'media_server_id'")
    if cur.fetchone():
        print(f"{DIM}online_classes.media_server_id already exists{RST}")
    else:
        cur.execute("ALTER TABLE online_classes ADD COLUMN media_server_id BIGINT UNSIGNED NULL AFTER schedule_id")
        try:
            cur.execute(
                "ALTER TABLE online_classes ADD CONSTRAINT fk_online_classes_media_server "
                "FOREIGN KEY (media_server_id) REFERENCES media_servers(id) ON DELETE SET NULL"
            )
        except Exception as exc:
            print(f"{Y}column added; FK skipped: {exc}{RST}")
        print(f"{G}added online_classes.media_server_id{RST}")

    # Seed the primary coturn from .env if the table is empty.
    cur.execute("SELECT COUNT(*) c FROM media_servers")
    if cur.fetchone()["c"] == 0:
        env = read_env()
        host = env.get("TURN_HOST", "")
        if host:
            cur.execute(
                "INSERT INTO media_servers (name, kind, host, port, scheme, region, capacity, priority, status, created_at, updated_at) "
                "VALUES (%s,'turn',%s,%s,'turn',%s,%s,%s,'active',NOW(),NOW())",
                ("سرور اصلی (coturn)", host, int(env.get("TURN_PORT", 3478)), "de-fsn1", 50, 10),
            )
            print(f"{G}seeded primary server from .env: {host}{RST}")
        else:
            print(f"{Y}TURN_HOST not in .env — add a server manually{RST}")
    conn.close()


def cmd_list(_args) -> None:
    conn = connect(); cur = conn.cursor()
    cur.execute("SHOW TABLES LIKE 'media_servers'")
    if not cur.fetchone():
        print(f"{Y}table missing — run: python scripts/media_servers.py migrate{RST}")
        return

    cur.execute("SELECT * FROM media_servers ORDER BY priority, id")
    rows = cur.fetchall()
    if not rows:
        print("no media servers registered")
        return

    print(f"\n{'ID':<4}{'NAME':<26}{'KIND':<6}{'HOST':<24}{'PORT':<7}{'LOAD':<10}{'STATUS':<10}HEALTH")
    print("─" * 100)
    for r in rows:
        pct = (r["current_load"] / r["capacity"] * 100) if r["capacity"] else 0
        load = f"{r['current_load']}/{r['capacity']}"
        colour = G if r["status"] == "active" else (Y if r["status"] == "draining" else R)
        health = f"{G}ok{RST}" if r["last_health_ok"] else (f"{R}fail{RST}" if r["last_health_at"] else f"{DIM}never{RST}")
        secret = "env" if not r["secret"] else "db"
        print(f"{r['id']:<4}{r['name'][:25]:<26}{r['kind']:<6}{r['host'][:23]:<24}{r['port']:<7}"
              f"{load:<10}{colour}{r['status']:<10}{RST}{health}  {DIM}secret:{secret} {pct:.0f}%{RST}")
    print()
    conn.close()


def cmd_add(args) -> None:
    conn = connect(); cur = conn.cursor()
    cur.execute("SHOW COLUMNS FROM media_servers LIKE 'manually_disabled'")
    has_manual_disable = bool(cur.fetchone())
    manual_update = ", manually_disabled=FALSE" if has_manual_disable else ""
    cur.execute(
        "INSERT INTO media_servers (name, kind, host, port, scheme, secret, region, capacity, priority, status, created_at, updated_at) "
        "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,'active',NOW(),NOW()) "
        "ON DUPLICATE KEY UPDATE name=VALUES(name), secret=VALUES(secret), capacity=VALUES(capacity), "
        f"priority=VALUES(priority), status='active'{manual_update}, updated_at=NOW()",
        (args.name, args.kind, args.host, args.port, args.scheme, args.secret,
         args.region, args.capacity, args.priority),
    )
    print(f"{G}registered {args.host}:{args.port} ({args.kind}){RST}")
    conn.close()
    cmd_list(args)


def _probe(host: str, port: int, timeout: float = 2.5) -> tuple[bool, str]:
    """TCP-connect to the TURN port. coturn listens on TCP as well as UDP."""
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True, "tcp reachable"
    except Exception as exc:
        return False, f"{type(exc).__name__}: {exc}"


def cmd_health(_args) -> None:
    conn = connect(); cur = conn.cursor()
    cur.execute("SELECT * FROM media_servers ORDER BY id")
    servers = cur.fetchall()
    if not servers:
        print("no servers to check")
        return

    cur.execute("SHOW COLUMNS FROM media_servers LIKE 'manually_disabled'")
    has_manual_disable = bool(cur.fetchone())
    # Network waits happen in parallel so one unreachable host does not delay
    # every other server in the five-second monitoring interval.
    worker_count = min(8, len(servers))
    with ThreadPoolExecutor(max_workers=worker_count) as pool:
        futures = {
            server["id"]: pool.submit(_probe, server["host"], server["port"])
            for server in servers
        }
        probe_results = {server_id: future.result() for server_id, future in futures.items()}

    for s in servers:
        ok, note = probe_results[s["id"]]
        if has_manual_disable:
            cur.execute(
                "UPDATE media_servers SET last_health_at=NOW(), last_health_ok=%s, last_health_note=%s, "
                "status=CASE WHEN manually_disabled=1 THEN 'offline' "
                "            WHEN %s=0 AND status='active' THEN 'offline' "
                "            WHEN %s=1 AND status='offline' THEN 'active' ELSE status END, "
                "updated_at=NOW() WHERE id=%s",
                (ok, note[:250], int(ok), int(ok), s["id"]),
            )
        else:
            cur.execute(
                "UPDATE media_servers SET last_health_at=NOW(), last_health_ok=%s, last_health_note=%s, "
                "status=CASE WHEN %s=0 AND status='active' THEN 'offline' "
                "            WHEN %s=1 AND status='offline' THEN 'active' ELSE status END, "
                "updated_at=NOW() WHERE id=%s",
                (ok, note[:250], int(ok), int(ok), s["id"]),
            )
        mark = f"{G}OK  {RST}" if ok else f"{R}FAIL{RST}"
        print(f"  [{mark}] {s['name']} — {s['host']}:{s['port']}  {DIM}{note}{RST}")

    # Refresh current_load from live participants.
    cur.execute("SHOW COLUMNS FROM online_classes LIKE 'media_server_id'")
    if cur.fetchone():
        cur.execute("UPDATE media_servers SET current_load = 0")
        cur.execute("""
            UPDATE media_servers ms
            SET ms.current_load = (
              SELECT COUNT(*) FROM online_classes oc
              JOIN online_class_participants p ON p.online_class_id = oc.id AND p.kicked = 0
              WHERE oc.media_server_id = ms.id AND oc.status IN ('scheduled','live')
                AND p.last_seen_at >= DATE_SUB(NOW(), INTERVAL 30 SECOND)
            ), ms.updated_at = NOW()
        """)
        print(f"{DIM}load recomputed from live classes{RST}")
    conn.close()


def cmd_install(args) -> None:
    """Install coturn on a fresh Ubuntu host and print the `add` command."""
    try:
        import paramiko
    except ImportError:
        print(f"{R}pip install paramiko{RST}")
        sys.exit(1)

    import getpass
    password = args.ssh_password or getpass.getpass(f"SSH password for {args.ssh_user}@{args.host}: ")
    secret = secrets.token_hex(32)

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    print(f"{B}connecting to {args.host}:{args.ssh_port}…{RST}")
    client.connect(args.host, port=args.ssh_port, username=args.ssh_user, password=password, timeout=30)

    def run(cmd: str, quiet: bool = False) -> str:
        _, out, err = client.exec_command(cmd, timeout=300)
        text = (out.read() + err.read()).decode("utf-8", "replace")
        if not quiet and text.strip():
            print(DIM + text.strip()[:400] + RST)
        return text

    print(f"{B}installing coturn…{RST}")
    run("DEBIAN_FRONTEND=noninteractive apt-get update -qq && "
        "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq coturn")

    conf = f"""# Generated by scripts/media_servers.py
listening-port={args.port}
fingerprint
use-auth-secret
static-auth-secret={secret}
realm={args.realm or args.host}
total-quota=100
stale-nonce=600
min-port=49152
max-port=65535
no-multicast-peers
no-cli
# Relay only — never let a client reach internal networks
denied-peer-ip=10.0.0.0-10.255.255.255
denied-peer-ip=172.16.0.0-172.31.255.255
denied-peer-ip=192.168.0.0-192.168.255.255
denied-peer-ip=127.0.0.0-127.255.255.255
"""
    sftp = client.open_sftp()
    with sftp.open("/etc/turnserver.conf", "w") as f:
        f.write(conf.encode("utf-8"))
    sftp.close()

    run("sed -i 's/^#TURNSERVER_ENABLED=1/TURNSERVER_ENABLED=1/' /etc/default/coturn 2>/dev/null; "
        "grep -q TURNSERVER_ENABLED /etc/default/coturn || echo 'TURNSERVER_ENABLED=1' >> /etc/default/coturn")
    run(f"ufw allow {args.port}/udp 2>/dev/null; ufw allow {args.port}/tcp 2>/dev/null; "
        f"ufw allow 49152:65535/udp 2>/dev/null; true", quiet=True)
    run("systemctl enable coturn >/dev/null 2>&1; systemctl restart coturn; sleep 1; systemctl is-active coturn")

    state = run("systemctl is-active coturn", quiet=True).strip()
    client.close()

    if state != "active":
        print(f"{R}coturn did not start (state={state}). Check `journalctl -u coturn` on the box.{RST}")
        sys.exit(1)

    print(f"\n{G}coturn is running on {args.host}:{args.port}{RST}")
    print(f"\n{B}Register it:{RST}")
    print(f'  python scripts/media_servers.py add --name "{args.name}" \\\n'
          f'      --host {args.host} --port {args.port} --secret {secret} --capacity {args.capacity}')
    print(f"\n{Y}The secret is only shown once — copy it now.{RST}\n")


def cmd_status_change(args, status: str) -> None:
    conn = connect(); cur = conn.cursor()
    cur.execute("SHOW COLUMNS FROM media_servers LIKE 'manually_disabled'")
    if cur.fetchone():
        cur.execute(
            "UPDATE media_servers SET status=%s, manually_disabled=%s, updated_at=NOW() WHERE id=%s",
            (status, status == "offline", args.id),
        )
    else:
        cur.execute("UPDATE media_servers SET status=%s, updated_at=NOW() WHERE id=%s", (status, args.id))
    if cur.rowcount == 0:
        print(f"{R}no server with id {args.id}{RST}")
    else:
        print(f"{G}server {args.id} -> {status}{RST}")
    conn.close()
    cmd_list(args)


def cmd_remove(args) -> None:
    conn = connect(); cur = conn.cursor()
    cur.execute("SELECT COUNT(*) c FROM online_classes WHERE media_server_id=%s AND status='live'", (args.id,))
    live = cur.fetchone()["c"]
    if live:
        print(f"{R}{live} live class(es) still on this server. Drain it first:{RST}")
        print(f"  python scripts/media_servers.py drain {args.id}")
        conn.close()
        return
    cur.execute("DELETE FROM media_servers WHERE id=%s", (args.id,))
    print(f"{G}removed server {args.id}{RST}" if cur.rowcount else f"{R}no server with id {args.id}{RST}")
    conn.close()


def main() -> None:
    p = argparse.ArgumentParser(description="Manage Daneshyar media (TURN/SFU) servers")
    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("list", help="show registered servers").set_defaults(fn=cmd_list)
    sub.add_parser("migrate", help="create table + column, seed primary from .env").set_defaults(fn=cmd_migrate)
    sub.add_parser("health", help="probe every server and recompute load").set_defaults(fn=cmd_health)

    a = sub.add_parser("add", help="register a server")
    a.add_argument("--name", required=True)
    a.add_argument("--host", required=True)
    a.add_argument("--port", type=int, default=3478)
    a.add_argument("--kind", choices=["turn", "sfu"], default="turn")
    a.add_argument("--scheme", default="turn")
    a.add_argument("--secret", default=None, help="HMAC secret; omit to use TURN_SECRET from .env")
    a.add_argument("--region", default=None)
    a.add_argument("--capacity", type=int, default=50)
    a.add_argument("--priority", type=int, default=100)
    a.set_defaults(fn=cmd_add)

    i = sub.add_parser("install", help="install coturn on a fresh Ubuntu host over SSH")
    i.add_argument("--host", required=True)
    i.add_argument("--name", default="سرور رسانه")
    i.add_argument("--port", type=int, default=3478)
    i.add_argument("--realm", default=None)
    i.add_argument("--capacity", type=int, default=50)
    i.add_argument("--ssh-user", default="root")
    i.add_argument("--ssh-port", type=int, default=22)
    i.add_argument("--ssh-password", default=None)
    i.set_defaults(fn=cmd_install)

    for name, status in (("drain", "draining"), ("enable", "active"), ("disable", "offline")):
        s = sub.add_parser(name, help=f"set a server to {status}")
        s.add_argument("id", type=int)
        s.set_defaults(fn=lambda args, _s=status: cmd_status_change(args, _s))

    r = sub.add_parser("remove", help="delete a server")
    r.add_argument("id", type=int)
    r.set_defaults(fn=cmd_remove)

    args = p.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()