#!/usr/bin/env python3
"""
xSTEM Scope — осциллограф из школьного датчика «Датчик напряжения
осциллографический» НАУРА (модель B202, USB, FTDI FT232R).

Запуск:  python3 xstem-scope.py            (порт ищется сам)
         python3 xstem-scope.py --port /dev/ttyUSB0
Потом в браузере: http://localhost:8765  (откроется сам)

Нужен только Python 3 из коробки Linux: никаких pip, Wine и драйверов.
В Chromium программа не обязательна: страница с сайта работает с датчиком
напрямую (Web Serial). Программа нужна для браузеров без Web Serial и работы без интернета.
Страница открывается только на этом компьютере: экраны групп учитель смотрит в Veyon.
Протокол датчика описан в THEORY.md и в xstem-scope.js рядом.
"""

import argparse
import fcntl
import glob
import json
import math
import os
import re
import struct
import termios
import threading
import time
import urllib.parse
import urllib.request
import webbrowser
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

HERE = os.path.dirname(os.path.abspath(__file__))

# Скорости, допустимые периоды и USB-коды датчика и генератора — в b202.json.
# Его же читает страница (xstem-scope.js), поэтому числа живут в одном месте.
with open(os.path.join(HERE, 'b202.json'), encoding='utf-8') as _f:
    DEVICE = json.load(_f)

CMD_BAUD = getattr(termios, f'B{DEVICE["baud"]["command"]}')     # скорость команд
DATA_BAUD = getattr(termios, f'B{DEVICE["baud"]["data"]}')       # скорость потока данных
GEN_BAUD = getattr(termios, f'B{DEVICE["baud"]["generator"]}')
# Датчик называет параметр FREQ, но это период в мкс: FREQ=10 → 100 кГц
PERIODS_US = DEVICE['periods_us']
PACKET_MAX = DEVICE['packet_max']
BOOT_S = DEVICE['boot_s']        # сколько датчик включается после того, как его вставили
SENSOR_IDS = {(u['vid'], u['pid']) for u in DEVICE['usb']['sensor']}
GENERATOR_IDS = {(u['vid'], u['pid']) for u in DEVICE['usb']['generator']}


def usb_ports(ids):
    """Последовательные порты, у которых USB VID:PID из списка ids."""
    found = []
    for tty in sorted(glob.glob('/sys/class/tty/ttyUSB*') + glob.glob('/sys/class/tty/ttyACM*')):
        dev = os.path.realpath(os.path.join(tty, 'device'))
        for up in (dev, os.path.dirname(dev), os.path.dirname(os.path.dirname(dev))):
            try:
                vid = open(os.path.join(up, 'idVendor')).read().strip()
                pid = open(os.path.join(up, 'idProduct')).read().strip()
            except OSError:
                continue
            if (vid, pid) in ids:
                found.append('/dev/' + os.path.basename(tty))
            break
    return found


def find_port():
    """Датчик Linux видит как ttyUSB с FTDI FT232R (0403:6001)."""
    ports = usb_ports(SENSOR_IDS)
    return ports[0] if ports else None


class Generator:
    """
    ESP32 с прошивкой-генератором (src/main.cpp): текстовые команды на 115200.
    Линии DTR и RTS у DevKit управляют сбросом платы, поэтому держим их
    снятыми и не опускаем при закрытии порта (без HUPCL) — иначе плата
    перезагрузится или уйдёт в режим прошивки.
    """

    def __init__(self, port):
        self.port = port
        self.lock = threading.Lock()
        self.fd = os.open(port, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
        a = termios.tcgetattr(self.fd)
        a[0] = a[1] = a[3] = 0
        a[2] = termios.CS8 | termios.CREAD | termios.CLOCAL      # без HUPCL
        a[4] = a[5] = GEN_BAUD
        termios.tcsetattr(self.fd, termios.TCSANOW, a)
        fcntl.ioctl(self.fd, termios.TIOCMBIC, struct.pack('I', termios.TIOCM_DTR | termios.TIOCM_RTS))

    def send(self, cmd):
        """Отправить команду, вернуть ответ платы (строки до 0,15 с тишины)."""
        cmd = ''.join(ch for ch in str(cmd) if ch.isalnum() or ch in ' .?')[:40]
        with self.lock:
            termios.tcflush(self.fd, termios.TCIFLUSH)
            os.write(self.fd, (cmd + '\n').encode('ascii'))
            end, quiet, buf = time.time() + 1.0, time.time() + 0.3, b''
            while time.time() < min(end, quiet):
                time.sleep(0.02)
                try:
                    chunk = os.read(self.fd, 1024)
                except BlockingIOError:
                    chunk = b''
                if chunk:
                    buf += chunk
                    quiet = time.time() + 0.15
            return buf.decode('utf-8', 'replace').strip()


def clean(raw):
    """
    Ответ на команду: 0x00, байт длины, текст, добитый пробелами; на установку
    значения — один 0x00; 0x05 — «не знаю такую команду». Байт длины бывает любым:
    0x27 — это апостроф, 0x32 — цифра «2», поэтому отбрасываем его по смыслу, а не
    как «лишний символ». Как clean() в xstem-scope.js; ответы записаны с B202 3.02.
    """
    if len(raw) >= 2 and raw[0] == 0 and raw[1] == len(raw) - 2:
        return raw[2:].decode('cp1251', 'replace').strip()
    if b'\x05' in raw:
        return None
    # Ответ пришёл не целиком или с остатками потока — как раньше: убираем то, что по краям
    return raw.replace(b'\x00', b'').decode('cp1251', 'replace').strip(" '\x01\t\n\r\x0b\x0c")


INFO_RE = re.compile(r'([A-Z]\d{3}[\w-]*),([\d.]+),([\d.]+),(\d+),(\w+),(\d+)', re.ASCII)


def parse_info(raw):
    """
    Ответ на «I»: «B202,3.02,1.12.2011,2,серийный,921600». Перед ним могут
    оказаться остатки потока данных, если прошлое измерение оборвали на ходу, —
    ищем сам шаблон. Как parseInfo() в xstem-scope.js.
    """
    m = INFO_RE.search(raw.decode('cp1251', 'replace'))
    if not m:
        return None
    return {'model': m[1], 'firmware': m[2], 'date': m[3],
            'channels': int(m[4]), 'serial': m[5], 'baud': int(m[6])}


def half_up(x):
    """Округление как Math.round на странице: round() в Python округляет 0,5 к чётному."""
    return math.floor(x + 0.5)


def commands(s):
    """
    Настройки страницы → (период, точек в кадре, команды датчику). Тот же
    пересчёт — commands() в xstem-scope.js; совпадение проверяют тесты (tests/).
    """
    period = s['period_us'] if s['period_us'] in PERIODS_US else 100
    # Точек в кадре: 3072 обычно, меньше — в самописце, чтобы кадры шли чаще
    packet = max(16, min(PACKET_MAX, half_up(s.get('packet') or PACKET_MAX)))
    cmds = [f'PACKET_SIZE={packet}', f'FREQ={period}']
    for ch in (0, 1):
        cmds.append(f'GAIN{ch}={int(s["gain"][ch]) & 3}')
        cmds.append(f'CHANNEL{ch}={"A" if s["ac"][ch] else "D"}')
    trig = s.get('trigger', 'N')
    if trig in ('0', '1'):
        # Уровень страница присылает уже кодом АЦП: калибровка и ноль хранятся в ней
        level = float(s.get('trigger_code', 2048))
        cmds += [f'START_CHANNEL={trig}',
                 f'START_LEVEL={max(1, min(4094, half_up(level)))}',
                 f'DIRECTION={"D" if s.get("falling") else "U"}']
    else:
        cmds.append('START_CHANNEL=N')
    return period, packet, cmds


class Sensor:
    """Обмен с датчиком: текстовые команды на 38400, кадры данных на 921600."""

    def __init__(self, port):
        self.port = port
        self.fd = os.open(port, os.O_RDWR | os.O_NOCTTY)
        self.lock = threading.Lock()
        try:
            self._setup()
        except OSError:
            self.close()      # программа пробует открыть порт снова каждую секунду
            raise

    def _setup(self):
        self._baud(CMD_BAUD)
        # Как в программе производителя: DTR снят, RTS выставлен
        fcntl.ioctl(self.fd, termios.TIOCMBIC, struct.pack('I', termios.TIOCM_DTR))
        fcntl.ioctl(self.fd, termios.TIOCMBIS, struct.pack('I', termios.TIOCM_RTS))
        time.sleep(0.2)
        # Если датчик остался в режиме измерения после прошлого запуска — останавливаем
        self._baud(DATA_BAUD)
        os.write(self.fd, b'T\n')
        time.sleep(0.15)
        self._baud(CMD_BAUD)
        time.sleep(0.1)
        termios.tcflush(self.fd, termios.TCIOFLUSH)

    def _baud(self, baud):
        a = termios.tcgetattr(self.fd)
        a[0] = a[1] = a[3] = 0
        a[2] = termios.CS8 | termios.CREAD | termios.CLOCAL
        a[4] = a[5] = baud
        a[6][termios.VMIN] = 0
        a[6][termios.VTIME] = 1          # read() ждёт не дольше 0,1 с
        termios.tcsetattr(self.fd, termios.TCSANOW, a)

    def _read_reply(self, timeout=0.4):
        end = time.time() + timeout
        buf = b''
        while time.time() < end:
            chunk = os.read(self.fd, 256)
            if chunk:
                buf += chunk
                end = min(end, time.time() + 0.06)   # ответ пришёл — ждём только хвост
        return buf

    def command(self, body):
        """Команда вида «FREQ=100» или «GAIN0=?»; длина тела входит в команду."""
        termios.tcflush(self.fd, termios.TCIOFLUSH)
        os.write(self.fd, f'C {len(body)} {body}\n'.encode('ascii'))
        return clean(self._read_reply())

    def info(self, attempts=3):
        # Сразу после остановки потока датчик может пропустить первый запрос
        for _ in range(attempts):
            termios.tcflush(self.fd, termios.TCIOFLUSH)
            os.write(self.fd, b'I\n')
            info = parse_info(self._read_reply())
            if info:
                return info
            time.sleep(0.3)
        return None

    def configure(self, s):
        period, packet, cmds = commands(s)
        for c in cmds:
            self.command(c)
        return period, packet

    def start(self):
        termios.tcflush(self.fd, termios.TCIOFLUSH)
        os.write(self.fd, b'S\n')
        termios.tcdrain(self.fd)
        time.sleep(0.05)
        self._baud(DATA_BAUD)

    def stop(self):
        # «T» датчик слушает на скорости данных, иначе он продолжит слать кадры
        os.write(self.fd, b'T\n')
        termios.tcdrain(self.fd)
        time.sleep(0.1)
        self._baud(CMD_BAUD)
        time.sleep(0.2)
        termios.tcflush(self.fd, termios.TCIOFLUSH)

    def read_chunk(self):
        return os.read(self.fd, 65536)

    def alive(self):
        """Порт на месте, и это всё ещё наш датчик: вставленный заново получает новый файл порта."""
        try:
            a, b = os.stat(self.port), os.fstat(self.fd)
        except OSError:
            return False
        return (a.st_dev, a.st_ino) == (b.st_dev, b.st_ino)

    def close(self):
        try:
            os.close(self.fd)
        except OSError:
            pass


class FrameParser:
    """
    Данные идут полубайтами (4 бита): «d» + номер канала открывают кадр,
    цифры 0–9 — десятичная запись кода АЦП, «c» закрывает отсчёт,
    «f» — заполнитель. Кадр = packet_size отсчётов канала 0, затем канала 1.
    """

    def __init__(self, packet_size):
        self.n = packet_size
        self.reset()

    def reset(self):
        self.ch = None
        self.expect_ch = False
        self.num = None
        self.frame = {}

    def feed(self, data):
        frames = []
        for byte in data:
            for nib in (byte >> 4, byte & 0xF):
                if self.expect_ch:
                    self.ch = nib
                    self.expect_ch = False
                    if nib == 0:
                        self.frame = {}
                    self.frame[nib] = []
                    self.num = None
                elif nib == 0xD:
                    self.expect_ch = True
                elif nib <= 9:
                    self.num = (self.num or 0) * 10 + nib
                elif nib == 0xC:
                    if self.num is not None and self.ch is not None:
                        vals = self.frame[self.ch]
                        vals.append(self.num)
                        if self.ch == 1 and len(vals) == self.n and 0 in self.frame:
                            frames.append(self.frame)
                            self.frame = {}
                            self.ch = None
                    self.num = None
        return frames


class SensorMissing(Exception):
    """Датчик не подключён: измерять нечем."""


class Scope:
    """
    Фоновый поток измерений и последний кадр для страницы. Датчик выдернули —
    останавливаем всё, закрываем порт и раз в секунду ищем его снова; вставили —
    ждём «Пуск». Программу перезапускать не нужно.
    """

    def __init__(self, sensor, reopen=None, watch_s=1.0, boot_s=BOOT_S):
        self.reopen = reopen or (lambda: None)     # → новый Sensor или None
        self.watch_s = watch_s
        self.boot_s = boot_s
        self.thread = None
        self.stop_flag = threading.Event()
        self.cond = threading.Condition()
        self.frame_id = 0
        self.frame = None
        self.state = 'stopped'          # stopped | waiting | running | nosensor
        self.settings = None
        self.problem = None             # почему датчика нет — показывает страница
        self.sensor = None
        self.info = None
        self.packet = PACKET_MAX
        self._attach(sensor)

    def _attach(self, sensor):
        info = sensor.info()
        if not info:
            sensor.close()
            return False
        sensor.command(f'PACKET_SIZE={PACKET_MAX}')      # самописец мог оставить короткий кадр
        self.packet = int(sensor.command('PACKET_SIZE=?') or PACKET_MAX)
        self.info, self.problem, self.sensor = info, None, sensor
        return True

    def watch(self):
        threading.Thread(target=self._watch, daemon=True).start()

    def _watch(self):
        while True:
            time.sleep(self.watch_s)
            if self.sensor is not None:
                if not self.sensor.alive():
                    self._lost('датчик отключён — вставьте его и нажмите «Пуск»')
                continue
            try:
                sensor = self.reopen()
            except PermissionError as e:
                self.problem = f'нет доступа к {e.filename}'
                continue
            except OSError:
                continue
            if sensor is None:
                continue
            # Только что вставленный датчик первые секунды мерит со своими настройками
            # (b202.json → boot_s): «Пуск» разрешаем, когда он включится
            time.sleep(self.boot_s)
            try:
                ok = self._attach(sensor)
            except (OSError, ValueError):
                sensor.close()
                continue
            if not ok:
                self.problem = f'на {sensor.port} что-то есть, но датчик не ответил'
                continue
            self._set_state('stopped')

    def _lost(self, why):
        sensor, self.sensor = self.sensor, None
        if sensor is None:
            return
        self.stop_flag.set()
        if self.thread and self.thread is not threading.current_thread():
            self.thread.join()
        self.thread = None
        sensor.close()
        self.problem = why
        self._set_state('nosensor')

    def start(self, settings):
        self.stop()
        sensor = self.sensor
        if sensor is None:
            raise SensorMissing(self.problem or 'датчик не подключён')
        try:
            with sensor.lock:
                period, packet = sensor.configure(settings)
        except OSError:
            self._lost('датчик отключён — вставьте его и нажмите «Пуск»')
            raise SensorMissing(self.problem)
        self.settings = dict(settings, period_us=period, packet=packet)
        self.stop_flag.clear()
        self._set_state('waiting')       # до запуска потока: страница уже может ждать кадр
        self.thread = threading.Thread(target=self._run, args=(sensor,), daemon=True)
        self.thread.start()

    def stop(self):
        if self.thread:
            self.stop_flag.set()
            self.thread.join()
            self.thread = None
        self._set_state('stopped' if self.sensor else 'nosensor')

    def _set_state(self, st):
        with self.cond:
            self.state = st
            self.cond.notify_all()

    def _run(self, sensor):
        parser = FrameParser(self.settings['packet'])
        single = self.settings.get('single')
        lost = False
        with sensor.lock:
            try:
                sensor.start()
                while not self.stop_flag.is_set():
                    chunk = sensor.read_chunk()
                    if not chunk:
                        continue
                    for fr in parser.feed(chunk):
                        with self.cond:
                            self.frame_id += 1
                            self.frame = {'id': self.frame_id, 'time': time.time(),
                                          'period_us': self.settings['period_us'],
                                          'gain': self.settings['gain'],
                                          'roll': self.settings.get('roll') or 0,
                                          'ch': [fr[0], fr[1]]}
                            self.state = 'running'
                            self.cond.notify_all()
                        if single:
                            self.stop_flag.set()
            except OSError:
                lost = True
            finally:
                try:
                    sensor.stop()
                except OSError:
                    lost = True
        if lost:
            self._lost('датчик отключён — вставьте его и нажмите «Пуск»')
        elif self.sensor is sensor and self.thread is threading.current_thread():
            self._set_state('stopped')

    def wait_frame(self, after, timeout=20):
        with self.cond:
            self.cond.wait_for(lambda: self.frame_id > after or self.state in ('stopped', 'nosensor'), timeout)
            return self.frame if self.frame_id > after else None


def already_running(url):
    try:
        with urllib.request.urlopen(url + '/api/info', timeout=1) as r:
            return r.status == 200
    except OSError:
        return False


STATIC = {
    '/': ('index.html', 'text/html; charset=utf-8'),
    '/index.html': ('index.html', 'text/html; charset=utf-8'),
    '/xstem-scope.js': ('xstem-scope.js', 'text/javascript; charset=utf-8'),
    '/b202.json': ('b202.json', 'application/json; charset=utf-8'),
    '/experiments.json': ('experiments.json', 'application/json; charset=utf-8'),
}


def make_handler(scope, generator=None):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass

        def _from_page(self):
            # Команды принимаем только от страницы программы, а не от любого сайта,
            # открытого в том же браузере. Тело application/json чужой сайт пошлёт
            # только после запроса OPTIONS, а на него программа не отвечает.
            # Host закрывает подмену DNS: чужое имя, указывающее на 127.0.0.1.
            ctype = (self.headers.get('Content-Type') or '').split(';')[0].strip().lower()
            host = urllib.parse.urlsplit('//' + (self.headers.get('Host') or '')).hostname
            return ctype == 'application/json' and host in ('localhost', '127.0.0.1', '::1')

        def _send(self, code, body=b'', ctype='application/json'):
            self.send_response(code)
            self.send_header('Content-Type', ctype)
            self.send_header('Cache-Control', 'no-store')
            self.send_header('Content-Length', str(len(body)))
            self.end_headers()
            try:
                self.wfile.write(body)
            except (BrokenPipeError, ConnectionResetError):
                pass        # вкладку закрыли или обновили, пока страница ждала кадр

        def _json(self, obj, code=200):
            self._send(code, json.dumps(obj, ensure_ascii=False).encode())

        def _static(self, path):
            if path not in STATIC:
                return False
            name, ctype = STATIC[path]
            with open(os.path.join(HERE, name), 'rb') as f:
                self._send(200, f.read(), ctype)
            return True

        def do_GET(self):
            url = urllib.parse.urlsplit(self.path)
            if self._static(url.path):
                pass
            elif url.path == '/api/info':
                self._json({'info': scope.info, 'port': scope.sensor.port if scope.sensor else None,
                            'packet': scope.packet, 'state': scope.state, 'settings': scope.settings,
                            'problem': scope.problem,
                            'generator': generator.port if generator else None})
            elif url.path == '/api/frame':
                try:
                    after = int(urllib.parse.parse_qs(url.query).get('after', ['0'])[0])
                except ValueError:
                    after = 0
                fr = scope.wait_frame(after)
                self._json({'state': scope.state, 'frame': fr, 'last': scope.frame_id})
            else:
                self._send(404)

        def do_POST(self):
            if not self._from_page():
                return self._send(403, b'{"error": "only from the xSTEM Scope page"}')
            try:
                length = int(self.headers.get('Content-Length') or 0)
                body = json.loads(self.rfile.read(length) or b'{}')
            except ValueError:
                return self._send(400, b'{"error": "bad json"}')
            if self.path == '/api/start':
                try:
                    scope.start(body)
                except SensorMissing as e:
                    return self._json({'error': str(e)}, 503)
                self._json({'ok': True, 'settings': scope.settings})
            elif self.path == '/api/gen':
                if not generator:
                    self._send(404, b'{"error": "no generator"}')
                else:
                    try:
                        self._json({'reply': generator.send(body.get('cmd', '?'))})
                    except OSError:
                        # Плату выдернули или переставили в другой разъём: порт стал другим
                        self._json({'error': 'ESP32 отключена — подключите её и перезапустите программу'}, 503)
            elif self.path == '/api/stop':
                scope.stop()
                self._json({'ok': True})
            else:
                self._send(404)

    return Handler


def main():
    ap = argparse.ArgumentParser(description='xSTEM Scope: осциллограф из датчика НАУРА B202')
    ap.add_argument('--port', help='последовательный порт датчика (по умолчанию ищется сам)')
    ap.add_argument('--http', type=int, default=8765, help='порт веб-страницы')
    ap.add_argument('--no-browser', action='store_true', help='не открывать браузер')
    args = ap.parse_args()

    url = f'http://localhost:{args.http}'
    if already_running(url):
        # Ученик щёлкнул ярлык второй раз: программа уже работает, просто открываем страницу
        print(f'xSTEM Scope уже запущен: {url}')
        if not args.no_browser:
            webbrowser.open(url)
        return

    port = args.port or find_port()
    if not port:
        raise SystemExit('Датчик не найден: подключите его к USB и проверьте, что появился /dev/ttyUSB*')
    try:
        sensor = Sensor(port)
    except PermissionError:
        raise SystemExit(f'Нет доступа к {port}. Выполните: sudo usermod -aG dialout $USER — и перезайдите в систему')
    except OSError as e:
        raise SystemExit(f'Не удалось открыть {port} ({e.strerror}). Переподключите датчик и запустите снова')

    def reopen():
        # Вставленный заново датчик может получить другое имя: ttyUSB1 вместо ttyUSB0
        p = args.port or find_port()
        return Sensor(p) if p and os.path.exists(p) else None

    scope = Scope(sensor, reopen)
    if not scope.info:
        raise SystemExit(f'На {port} что-то есть, но датчик не ответил на запрос «I»')
    scope.watch()
    generator = None
    for gport in usb_ports(GENERATOR_IDS):
        try:
            generator = Generator(gport)
            print(f'Генератор ESP32: {gport}')
            break
        except OSError:
            pass
    print(f'Датчик {scope.info["model"]} № {scope.info["serial"]}, прошивка {scope.info["firmware"]}, порт {port}')

    server = ThreadingHTTPServer(('127.0.0.1', args.http), make_handler(scope, generator))
    print(f'Откройте в браузере: {url}   (остановить: Ctrl+C)')
    if not args.no_browser:
        threading.Timer(0.5, webbrowser.open, (url,)).start()
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        scope.stop()


if __name__ == '__main__':
    main()