---

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


def display_loop():
    cv2.namedWindow("Камера робота")
    cv2.namedWindow("Сонар")

    print("""
w — вперёд      s — назад
a — влево       d — вправо
пробел — стоп
1 — расстояние до ближайшей стенки
+ — быстрее вращение сонара    - — медленнее вращение сонара
q — выход
(окно с видео или сонаром должно быть в фокусе)
""")
    current_cmd["sweep_speed"] = 0.02
    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)

        key = cv2.waitKey(30) & 0xFF
        if key == 255:
            continue
        ch = chr(key) if key < 128 else ''

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

        if ch == '1':
            request_distance()
        elif ch == 'q':
            break


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("остановка, выход")