---

import socket
import threading
import json
import time

import cv2
import numpy as np

HOST = '0.0.0.0'
CMD_PORT = 5001
SCAN_PORT = 5002
VIDEO_PORT = 5000
SONAR_PORT = 5003

current_cmd = {"linear_x": 0.0, "angular_z": 0.0, "sweep_speed": 0.12}
cmd_lock = threading.Lock()


def command_server():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((HOST, CMD_PORT))
    server_socket.listen(1)
    print(f"жду подключения робота на порту {CMD_PORT}...")

    conn, addr = server_socket.accept()
    print(f"робот подключился: {addr}")

    try:
        while True:
            with cmd_lock:
                data = json.dumps(current_cmd).encode('utf-8')
            conn.sendall(len(data).to_bytes(4, 'big') + data)
            time.sleep(0.1)
    except (ConnectionResetError, BrokenPipeError, OSError):
        print("робот отключился")
    finally:
        conn.close()
        server_socket.close()


def recv_exact(sock, n):
    buf = b''
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


def video_server():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((HOST, VIDEO_PORT))
    server_socket.listen(1)
    print(f"жду подключения робота на порту {VIDEO_PORT} (видео)...")

    conn, addr = server_socket.accept()
    print(f"робот подключился (видео): {addr}")

    try:
        while True:
            size_bytes = recv_exact(conn, 4)
            if size_bytes is None:
                break
            frame_size = int.from_bytes(size_bytes, byteorder='big')

            frame_data = recv_exact(conn, frame_size)
            if frame_data is None:
                break

            nparr = np.frombuffer(frame_data, np.uint8)
            frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            if frame is not None:
                with display_lock:
                    latest_frames["video"] = frame
    except (ConnectionResetError, OSError):
        print("робот отключился (видео)")
    finally:
        conn.close()
        server_socket.close()


scan_conn = None
scan_lock = threading.Lock()

latest_frames = {"video": None, "sonar": None}
display_lock = threading.Lock()


def sonar_server():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((HOST, SONAR_PORT))
    server_socket.listen(1)
    print(f"жду подключения робота на порту {SONAR_PORT} (сонар)...")

    conn, addr = server_socket.accept()
    print(f"робот подключился (сонар): {addr}")

    try:
        while True:
            size_bytes = recv_exact(conn, 4)
            if size_bytes is None:
                break
            frame_size = int.from_bytes(size_bytes, byteorder='big')

            frame_data = recv_exact(conn, frame_size)
            if frame_data is None:
                break

            nparr = np.frombuffer(frame_data, np.uint8)
            frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            if frame is not None:
                with display_lock:
                    latest_frames["sonar"] = frame
    except (ConnectionResetError, OSError):
        print("робот отключился (сонар)")
    finally:
        conn.close()
        server_socket.close()


SPEED = 0.15
TURN = 0.3
SWEEP_STEP = 0.02
KEY_TIMEOUT = 0.25
def display_loop():
    cv2.namedWindow("Камера робота")
    cv2.namedWindow("Сонар")

    print("""
w — вперёд      s — назад
a — влево       d — вправо
пробел — стоп
1 — расстояние до ближайшей стенки
+ — быстрее вращение сонара    - — медленнее вращение сонара
q — выход
(окно с видео или сонаром должно быть в фокусе, клавишу нужно УДЕРЖИВАТЬ)
""")

    current_cmd["sweep_speed"] = 0.02

    # время последнего "подтверждения" каждой зажатой клавиши
    held_keys = {}   # char -> last_seen_time

    def press(ch):
        held_keys[ch] = time.time()

    def is_held(ch):
        t = held_keys.get(ch)
        return t is not None and (time.time() - t) <= KEY_TIMEOUT

    # запоминаем предыдущее состояние команд, чтобы не спамить одинаковыми значениями
    last_sent = {"linear_x": None, "angular_z": None}

    while True:
        with display_lock:
            video_frame = latest_frames["video"]
            sonar_frame = latest_frames["sonar"]
        if video_frame is not None:
            cv2.imshow("Камера робота", video_frame)
        if sonar_frame is not None:
            cv2.imshow("Сонар", sonar_frame)

        # --- 1. Читаем события клавиатуры ---
        # waitKey вызываем несколько раз с малым интервалом, чтобы ловить
        # повторяющиеся события удержания клавиши (key repeat от ОС).
        for _ in range(10):
            k = cv2.waitKey(1) & 0xFF
            if k == 255:
                continue
            ch = chr(k) if k < 128 else ''

            if ch == 'q':
                cv2.destroyAllWindows()
                with cmd_lock:
                    current_cmd["linear_x"] = 0.0
                    current_cmd["angular_z"] = 0.0
                return

            if ch == '1':
                request_distance()
                continue

            if ch in ('w', 'a', 's', 'd', ' '):
                press(ch)
            elif ch == '+':
                with cmd_lock:
                    current_cmd["sweep_speed"] = min(1.0, current_cmd["sweep_speed"] + SWEEP_STEP)
                    print(f"скорость вращения сонара: {current_cmd['sweep_speed']:.2f}")
            elif ch == '-':
                with cmd_lock:
                    current_cmd["sweep_speed"] = max(0.0, current_cmd["sweep_speed"] - SWEEP_STEP)
                    print(f"скорость вращения сонара: {current_cmd['sweep_speed']:.2f}")

        # --- 2. Вычисляем команду из зажатых клавиш ---
        lin = 0.0
        ang = 0.0
        if is_held('w'):
            lin = SPEED
        elif is_held('s'):
            lin = -SPEED
        if is_held('a'):
            ang = TURN
        elif is_held('d'):
            ang = -TURN
        if is_held(' '):
            lin = 0.0
            ang = 0.0

        # если зажато что-то одно из w/a/d/s — едем/поворачиваем; если ничего — стоим
        if lin != last_sent["linear_x"] or ang != last_sent["angular_z"]:
            with cmd_lock:
                current_cmd["linear_x"] = lin
                current_cmd["angular_z"] = ang
            last_sent["linear_x"] = lin
            last_sent["angular_z"] = ang

        # небольшая пауза, чтобы не грузить CPU
        time.sleep(0.01)


def scan_server():
    global scan_conn
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((HOST, SCAN_PORT))
    server_socket.listen(1)
    print(f"жду подключения робота на порту {SCAN_PORT} (лидар)...")

    conn, addr = server_socket.accept()
    print(f"робот подключился (лидар): {addr}")
    with scan_lock:
        scan_conn = conn


def request_distance():
    with scan_lock:
        conn = scan_conn
    if conn is None:
        print("канал лидара ещё не подключен")
        return
    try:
        conn.sendall(b'get')
        size_bytes = recv_exact(conn, 4)
        if size_bytes is None:
            print("робот не ответил")
            return
        msg_len = int.from_bytes(size_bytes, 'big')
        payload = recv_exact(conn, msg_len)
        data = json.loads(payload.decode('utf-8'))
        print(f"расстояние до ближайшей стены: {data['distance']:.2f} м")
    except (ConnectionResetError, OSError):
        print("канал лидара оборван")


if __name__ == '__main__':
    threading.Thread(target=command_server, daemon=True).start()
    threading.Thread(target=scan_server, daemon=True).start()
    threading.Thread(target=video_server, daemon=True).start()
    threading.Thread(target=sonar_server, daemon=True).start()

    try:
        display_loop()
    except KeyboardInterrupt:
        pass
    finally:
        with cmd_lock:
            current_cmd["linear_x"] = 0.0
            current_cmd["angular_z"] = 0.0
        cv2.destroyAllWindows()
        print("остановка, выход")