# -*- coding: utf-8 -*-
"""Duga — ПРОШИВАЛЬЩИК. Сам находит платы, сам смотрит версии, сам обновляет.

=====================================================================
ЧТО ОН ДЕЛАЕТ И ЧЕГО НАРОЧНО НЕ ДЕЛАЕТ
=====================================================================
ДЕЛАЕТ: перебирает все порты, узнаёт по опознавателю USB, что за
плата на каждом, спрашивает у живой платы номер сборки, сравнивает с
тем, что лежит рядом в папке «proshivki», и заливает только то, что
устарело.

НЕ ДЕЛАЕТ: не стирает флеш. Стирание сносит вместе с прошивкой
настройки, контакты, почту и КНИЖКУ КЛЮЧЕЙ — то есть все знакомства,
ради которых люди стояли рядом и диктовали друг другу восемь цифр.
Обновление, которое молча разрывает все знакомства, — это не
обновление, а сброс. Стирание есть, но только по прямой просьбе
(--s-nulya), и тогда об этом сказано вслух.

=====================================================================
КАК ОТЛИЧАЮТСЯ ТРИ ПЛАТЫ
=====================================================================
Не по номеру порта — номера меняются от того, в какое гнездо воткнули.
По производителю в опознавателе USB:

    10C4:EA60   Silicon Labs CP2102 — переходник, через него говорит V3
    303A:1001   Espressif          — встроенный USB кристалла, это V4
    239A:*      Adafruit           — свой USB nRF52840, это T114

=====================================================================
ПОЧЕМУ У T114 ЗАЛИВКА СОВСЕМ ДРУГАЯ
=====================================================================
У ESP32 прошивка заливается по тому же проводу, по которому плата
разговаривает: кристалл сам умеет принимать её в загрузчике.

У nRF52840 дорога иная: плата после двойного сброса показывается
компьютеру как ФЛЕШКА, и прошивка на неё просто копируется файлом.
Двойной сброс раньше делали пальцем. Здесь его делает сам
прошивальщик: если открыть порт платы на скорости 1200 и тут же
закрыть, загрузчик Adafruit понимает это как «просят в загрузчик» и
уходит туда сам. Приём известный и общий для всех плат Adafruit; мы
его не выдумывали.

У ЗАГРУЗЧИКА ДВЕ ДОРОГИ, И ВТОРАЯ НУЖНА ЧАЩЕ, ЧЕМ КАЗАЛОСЬ.
На части плат (наши T114 из них) касание на 1200 бод срабатывает —
загрузчик поднимается, — а ФЛЕШКИ не появляется: том система не
назначает. Раньше прошивальщик в этом месте упирался и просил нажать
RESET дважды, хотя плата уже стояла в загрузчике и ждала.

Поэтому смотрим оба признака сразу:
    появилась флешка          -> копируем .uf2, как раньше;
    появился порт 239A:0071   -> заливаем по проводу, пакетом .zip
                                 (это и есть загрузчик, только другой
                                 его лик).

Не вышло ни то ни другое — скажем человеку нажать дважды кнопку RESET,
и подождём. Молча сдаваться нельзя: человек решит, что плата умерла.

=====================================================================
ЗАПУСК
=====================================================================
    py -3 duga_proshivalshchik.py            обновить всё, что устарело
    py -3 duga_proshivalshchik.py --smotret  только посмотреть, ничего не трогать
    py -3 duga_proshivalshchik.py --vsegda   залить даже если версия та же
    py -3 duga_proshivalshchik.py --s-nulya  стереть флеш и залить с нуля
                                             (СНОСИТ знакомства и настройки)
"""
import io
import os
import re
import shutil
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
sys.stdout.reconfigure(line_buffering=True)

try:
    import serial
    from serial.tools import list_ports
except ImportError:
    print("нет модуля pyserial:  py -3 -m pip install pyserial")
    sys.exit(1)

ЗДЕСЬ = os.path.dirname(os.path.abspath(__file__))
ПАПКА_ПРОШИВОК = os.path.join(ЗДЕСЬ, "proshivki")

FEND, FESC, TFEND, TFESC = 0xC0, 0xDB, 0xDC, 0xDD
CMD_NASTROYKA = 0xA5
N_SBORKA = 34
N_NOMER  = 45
N_ZNAKOMYH = 74

ит = []


def шаг(имя, вышло, пояснение=""):
    ит.append((имя, bool(вышло)))
    print("  %s %s%s" % ("✅" if вышло else "🔴", имя,
                         ("   — " + пояснение) if пояснение else ""))


def сказать(текст):
    print("     " + текст)


# ---------------------------------------------------------------------
#  ПРОВОД: открыть порт, не перезагрузив плату
#
#  У плат со своим USB (V4, T114) линии DTR/RTS заведены на сброс
#  ВНУТРИ кристалла. pyserial по умолчанию поднимает DTR при открытии —
#  то есть каждый наш взгляд перезагружал бы плату. Опускаем их ДО
#  открытия: после — уже поздно.
# ---------------------------------------------------------------------
def открыть(порт, скорость=115200, срок=0.6):
    п = serial.Serial()
    п.port = порт
    п.baudrate = скорость
    п.timeout = срок
    п.write_timeout = 3.0
    п.dtr = True
    п.rts = False
    п.open()
    return п


def кадр(cmd, тело=b""):
    v = bytearray([FEND, cmd])
    for b in тело:
        if b == FEND:   v += bytes([FESC, TFEND])
        elif b == FESC: v += bytes([FESC, TFESC])
        else:           v.append(b)
    v.append(FEND)
    return bytes(v)


def кадры(сырое):
    из, тек, в_кадре, экр = [], bytearray(), False, False
    for b in сырое:
        if b == FEND:
            if в_кадре and тек:
                из.append(bytes(тек))
            тек, в_кадре, экр = bytearray(), True, False
            continue
        if not в_кадре:
            continue
        if экр:
            if b == TFEND: b = FEND
            elif b == TFESC: b = FESC
            экр = False
        elif b == FESC:
            экр = True
            continue
        тек.append(b)
    return из


def настройка(п, номер, попыток=3):
    for _ in range(попыток):
        try:
            п.reset_input_buffer()
            п.write(кадр(CMD_NASTROYKA, bytes([номер])))
        except Exception:
            return None
        time.sleep(0.4)
        try:
            сырое = п.read(8192)
        except Exception:
            return None
        for k in кадры(сырое):
            if k and k[0] == CMD_NASTROYKA and len(k) >= 6 and k[1] == номер:
                v = int.from_bytes(k[2:6], "big")
                return None if v == 0xFFFFFFFF else v
    return None


# ---------------------------------------------------------------------
#  ОПОЗНАНИЕ ПЛАТЫ ПО USB
# ---------------------------------------------------------------------
def какая_плата(п):
    имя = ((п.description or "") + " " + str(п.hwid or "")).upper()
    if "239A" in имя or "ADAFRUIT" in имя:
        return "T114"
    if "303A" in имя:
        return "V4"
    if any(x in имя for x in ("10C4", "CP210", "SILICON LABS", "1A86", "CH340")):
        return "V3"
    return None


# ---------------------------------------------------------------------
#  КАКИЕ ПРОШИВКИ ЛЕЖАТ РЯДОМ
#
#  Версия читается ИЗ ИМЕНИ ФАЙЛА, а не из отдельного списка. Список
#  пришлось бы держать в согласии с файлами вручную, и однажды он бы
#  разошёлся: файл новый, в списке старый номер — и прошивальщик
#  уверенно объявил бы плату свежей.
# ---------------------------------------------------------------------
def прошивки_рядом():
    найдено = {}
    if not os.path.isdir(ПАПКА_ПРОШИВОК):
        return найдено
    for имя in os.listdir(ПАПКА_ПРОШИВОК):
        м = re.match(r"duga_(t114|v3|v4)_sborka(\d+)\.(uf2|bin)$", имя, re.I)
        if not м:
            continue
        плата = {"t114": "T114", "v3": "V3", "v4": "V4"}[м.group(1).lower()]
        версия = int(м.group(2))
        путь = os.path.join(ПАПКА_ПРОШИВОК, имя)
        если_лучше = плата not in найдено or версия > найдено[плата]["версия"]
        if если_лучше:
            найдено[плата] = {"версия": версия, "путь": путь, "имя": имя}

    # ПАКЕТ ДЛЯ ЗАЛИВКИ ПО ПРОВОДУ — отдельным файлом и с тем же номером
    # сборки. Номер сверяем, а не берём любой лежащий: разъедься они,
    # плата получила бы одно, а проверка ждала бы другого — и мы бы
    # объявили исправную заливку неудачной.
    for имя in os.listdir(ПАПКА_ПРОШИВОК):
        м = re.match(r"duga_t114_sborka(\d+)_dfu\.zip$", имя, re.I)
        if not м:
            continue
        если = найдено.get("T114")
        if если and если["версия"] == int(м.group(1)):
            если["dfu"] = os.path.join(ПАПКА_ПРОШИВОК, имя)
    return найдено


# ---------------------------------------------------------------------
#  ЗАЛИВКА ESP32 (V3 и V4)
# ---------------------------------------------------------------------
def выполнить(команда, срок=600):
    try:
        r = subprocess.run(команда, capture_output=True, text=True,
                           encoding="utf-8", errors="replace", timeout=срок)
        return (r.stdout or "") + (r.stderr or ""), r.returncode
    except Exception as е:
        return str(е), 1


def залить_esp(порт, плата, путь_прошивки, с_нуля):
    осн = [sys.executable, "-m", "esptool", "--chip", "esp32s3", "--port", порт]

    if с_нуля:
        сказать("стираю флеш целиком — вместе с настройками и знакомствами")
        текст, код = выполнить(осн + ["erase-flash"], срок=180)
        if код != 0:
            текст, код = выполнить(осн + ["erase_flash"], срок=180)
        if код != 0:
            return False, текст.strip().splitlines()[-1] if текст.strip() else "не стёрлось"

    части = []
    if с_нуля:
        # Первая заливка: нужны все четыре части, иначе плата не поднимется.
        низ = "v4_" if плата == "V4" else "v3_"
        for адрес, файл in (("0x0", низ + "bootloader.bin"),
                            ("0x8000", низ + "partitions.bin"),
                            ("0xe000", "boot_app0.bin")):
            п = os.path.join(ПАПКА_ПРОШИВОК, файл)
            if not os.path.isfile(п):
                return False, "нет файла " + файл
            части += [адрес, п]
    части += ["0x10000", путь_прошивки]

    # У esptool пятой версии подкоманда пишется через дефис, у четвёртой —
    # через подчёркивание. Пробуем обе: узнать версию отдельным запуском
    # можно, но это ещё один запуск, который тоже умеет не получиться.
    текст = ""
    for под in ("write-flash", "write_flash"):
        текст, код = выполнить(осн + [под] + части, срок=600)
        if код == 0:
            сверок = текст.count("Hash of data verified")
            return True, "сверок хеша: %d" % сверок
        низ = текст.lower()
        if "unknown" not in низ and "invalid choice" not in низ:
            break     # ошибка настоящая, а не в написании подкоманды
    хвост = [с for с in текст.strip().splitlines() if с.strip()]
    return False, хвост[-1] if хвост else "заливка не удалась"


# ---------------------------------------------------------------------
#  ЗАЛИВКА T114 — ЧЕРЕЗ ФЛЕШКУ ЗАГРУЗЧИКА
# ---------------------------------------------------------------------
def диски_с_загрузчиком():
    """Все диски, на которых лежит INFO_UF2.TXT — признак загрузчика UF2."""
    найдено = []
    if os.name == "nt":
        for буква in "DEFGHIJKLMNOPQRSTUVWXYZ":
            корень = буква + ":\\"
            try:
                if os.path.isfile(os.path.join(корень, "INFO_UF2.TXT")):
                    найдено.append(корень)
            except Exception:
                pass
    else:
        for где in ("/media", "/mnt", "/run/media"):
            for корень, папки, файлы in os.walk(где):
                if "INFO_UF2.TXT" in файлы:
                    найдено.append(корень)
                if корень.count(os.sep) - где.count(os.sep) > 2:
                    папки[:] = []
    return найдено


def порты_загрузчика():
    """Порты, на которых сидит ЗАГРУЗЧИК nRF52, а не прошивка.

    Загрузчик объявляется тем же производителем 239A, что и рабочая
    плата, но ДРУГИМ изделием — 0071. По этой паре его и узнаём.
    Полагаться на номер COM нельзя: при уходе в загрузчик плата
    отключается и подключается заново, и номер меняется.
    """
    найдено = []
    for п in list_ports.comports():
        опознание = (str(п.hwid or "") + " " + (п.description or "")).upper()
        if "239A" in опознание and "0071" in опознание:
            найдено.append(п.device)
    return найдено


def nrfutil_gotov():
    """Есть ли чем говорить с загрузчиком по проводу.

    ИМЯ МОДУЛЯ И ИМЯ ПАКЕТА РАЗНЫЕ, и это ловушка: ставится
    adafruit-nrfutil, а зовётся nordicsemi. Проверять по имени пакета
    значит решить, что ничего не установлено, при установленном.

    Ставим сами и говорим об этом вслух. Молча трогать чужой Python
    нехорошо, а заставлять человека искать команду в справке ради
    одной заливки — ещё хуже. Не поставилось — не беда: вернёмся к
    дороге через флешку.
    """
    _т, код = выполнить([sys.executable, "-m", "nordicsemi", "--help"], срок=60)
    if код == 0:
        return True
    сказать("ставлю adafruit-nrfutil — им плата заливается по проводу, "
            "без нажатия кнопки. Установка разовая, несколько секунд")
    текст, код = выполнить([sys.executable, "-m", "pip", "install",
                            "--quiet", "adafruit-nrfutil"], срок=420)
    if код != 0:
        хвост = [с for с in текст.strip().splitlines() if с.strip()]
        сказать("поставить не вышло: " + (хвост[-1] if хвост else "молча"))
        return False
    _т, код = выполнить([sys.executable, "-m", "nordicsemi", "--help"], срок=60)
    return код == 0


def залить_t114_po_provodu(порт_загрузчика, путь_dfu):
    """DFU по проводу: то же, что копирование на флешку, только без неё.

    --singlebank обязателен. У T114 памяти под ДВЕ прошивки разом нет,
    и без этого ключа загрузчик отказывается принимать пакет, ссылаясь
    на нехватку места, — а звучит это как поломка платы.
    """
    текст, код = выполнить([sys.executable, "-m", "nordicsemi", "dfu",
                            "serial", "-pkg", путь_dfu, "-p", порт_загрузчика,
                            "-b", "115200", "--singlebank"], срок=600)
    if код == 0 and "Device programmed" in текст:
        return True, "залито по проводу, кнопку нажимать не пришлось"
    хвост = [с for с in текст.strip().splitlines() if с.strip()]
    return False, хвост[-1] if хвост else "заливка по проводу не удалась"


def в_загрузчик_1200(порт):
    """Открыть порт на 1200 и закрыть — загрузчик Adafruit уходит в UF2.

    ПОЧЕМУ ИМЕННО 1200. Это уговор, а не физика: скорость 1200 никем не
    используется для данных, поэтому её и выбрали признаком «просят в
    загрузчик». Так делают все платы Adafruit и вся Arduino на SAMD.
    Здесь DTR НАДО поднять — именно перепад его и ловит загрузчик.
    """
    try:
        п = serial.Serial()
        п.port = порт
        п.baudrate = 1200
        п.dtr = True
        п.open()
        time.sleep(0.25)
        п.dtr = False
        п.close()
        return True
    except Exception:
        return False


def залить_t114(порт, путь_прошивки, путь_dfu=None):
    было_дисков = set(диски_с_загрузчиком())
    были_порты  = set(порты_загрузчика())
    if было_дисков:
        сказать("плата уже в загрузчике: " + ", ".join(sorted(было_дисков)))
    elif были_порты:
        сказать("плата уже в загрузчике (по проводу): "
                + ", ".join(sorted(были_порты)))
    else:
        сказать("прошу плату уйти в загрузчик (открываю порт на 1200)...")
        в_загрузчик_1200(порт)

    диск = None
    загр_порт = None
    ждали = 0.0
    просили_руками = False
    while ждали < 40.0:
        # ФЛЕШКА — ПЕРВЫМ ДЕЛОМ. Где она есть, там ничего ставить не
        # надо и заливка идёт копированием файла.
        диски = set(диски_с_загрузчиком())
        новые_д = (диски - было_дисков) if было_дисков else диски
        if новые_д:
            диск = sorted(новые_д)[0]
            break
        # Флешки нет — но, может быть, загрузчик всё-таки поднялся и
        # ждёт по проводу. На наших платах бывает именно так.
        порты = set(порты_загрузчика())
        новые_п = (порты - были_порты) if были_порты else порты
        if новые_п:
            загр_порт = sorted(новые_п)[0]
            break
        time.sleep(0.5)
        ждали += 0.5
        if ждали > 8.0 and not просили_руками:
            просили_руками = True
            print()
            print("     ┌───────────────────────────────────────────────────┐")
            print("     │  Плата сама в загрузчик не ушла.                  │")
            print("     │  НАЖМИТЕ ДВА РАЗА ПОДРЯД кнопку RESET на плате.   │")
            print("     │  Жду ещё полминуты.                               │")
            print("     └───────────────────────────────────────────────────┘")
            print()

    # ---- дорога вторая: по проводу ----
    if загр_порт and not диск:
        сказать("загрузчик отозвался по проводу на " + загр_порт
                + " (флешки эта плата не показывает)")
        if not путь_dfu:
            return False, ("плата в загрузчике, но рядом нет пакета "
                           "duga_t114_sborka<N>_dfu.zip — заливать по проводу нечем")
        if not nrfutil_gotov():
            return False, ("плата в загрузчике, но заливать по проводу нечем. "
                           "Поставьте вручную:  py -3 -m pip install adafruit-nrfutil")
        return залить_t114_po_provodu(загр_порт, путь_dfu)

    if not диск:
        return False, ("ни флешки загрузчика, ни его порта так и не появилось. "
                       "Двойной сброс — две быстрых нажатия RESET подряд")

    сказать("загрузчик на диске " + диск)
    try:
        shutil.copyfile(путь_прошивки, os.path.join(диск, "duga.uf2"))
    except Exception as е:
        # Плата перезагружается ПОСРЕДИ копирования — это нормально:
        # загрузчик принимает файл кусками и уходит в новую прошивку,
        # не дожидаясь, пока Windows допишет последний блок. Windows
        # при этом честно сообщает об ошибке записи.
        текст = str(е).lower()
        if "устройство" in текст or "device" in текст or "not ready" in текст:
            сказать("диск исчез посреди копирования — так и должно быть: "
                    "загрузчик уже ушёл в новую прошивку")
        else:
            return False, str(е)
    return True, "прошивка скопирована на флешку загрузчика"


def дождаться_порта(имя, срок=30):
    """Порт исчезает при сбросе и появляется заново. Ждём его."""
    прошло = 0.0
    while прошло < срок:
        if any(п.device == имя for п in list_ports.comports()):
            return True
        time.sleep(0.5)
        прошло += 0.5
    return False


# ---------------------------------------------------------------------
def главное():
    смотреть = "--smotret" in sys.argv
    всегда   = "--vsegda"  in sys.argv
    с_нуля   = "--s-nulya" in sys.argv

    print("=" * 70)
    print("DUGA — ПРОШИВАЛЬЩИК")
    print("=" * 70)
    if смотреть:
        print("  только смотрю: ни одна плата не будет тронута")
    if с_нуля:
        print("  С НУЛЯ: флеш будет стёрт. Настройки, контакты, почта и")
        print("  КНИЖКА КЛЮЧЕЙ (все знакомства) пропадут безвозвратно.")
    print()

    есть = прошивки_рядом()
    if есть:
        print("ПРОШИВКИ РЯДОМ")
        print("-" * 70)
        for плата in sorted(есть):
            print("  %-5s сборка %-4d %s" % (плата, есть[плата]["версия"], есть[плата]["имя"]))
        print()
    шаг("прошивки найдены", len(есть) > 0,
        "папка «proshivki» рядом с этим файлом" if not есть else "видов плат: %d" % len(есть))

    порты = list(list_ports.comports())
    платы = []
    for п in порты:
        вид = какая_плата(п)
        if not вид:
            continue
        запись = {"порт": п.device, "вид": вид, "сборка": None, "номер": None,
                  "знакомых": None, "живая": False}
        try:
            с = открыть(п.device)
        except Exception:
            платы.append(запись)
            continue
        time.sleep(1.6)
        запись["сборка"] = настройка(с, N_SBORKA)
        if запись["сборка"] is not None:
            запись["живая"] = True
            запись["номер"] = настройка(с, N_NOMER)
            запись["знакомых"] = настройка(с, N_ZNAKOMYH)
        с.close()
        платы.append(запись)

    print()
    print("ПЛАТЫ НА ПРОВОДЕ")
    print("-" * 70)
    for з in платы:
        что = ("сборка %s" % з["сборка"]) if з["живая"] else "по-нашему молчит (новая или чужая прошивка)"
        доп = ""
        if з["живая"]:
            доп = ", номер %s" % з["номер"]
            if з["знакомых"] is not None:
                доп += ", знакомых %d" % з["знакомых"]
        print("  %-6s %-5s %s%s" % (з["порт"], з["вид"], что, доп))
    шаг("платы найдены", len(платы) > 0, "нашлось %d" % len(платы))
    if not платы or not есть:
        return итог()

    print()
    print("ЧТО НАДО СДЕЛАТЬ")
    print("-" * 70)
    работа = []
    for з in платы:
        если_есть = есть.get(з["вид"])
        if not если_есть:
            print("  %-6s %-5s прошивки для этой платы рядом нет — пропускаю"
                  % (з["порт"], з["вид"]))
            continue
        новая = если_есть["версия"]
        своя = з["сборка"]
        if с_нуля or всегда or своя is None or своя < новая:
            почему = ("с нуля" if с_нуля else
                      "просили залить всегда" if всегда else
                      "по-нашему молчит — заливаем впервые" if своя is None else
                      "сборка %d старее %d" % (своя, новая))
            print("  %-6s %-5s ЗАЛИТЬ  (%s)" % (з["порт"], з["вид"], почему))
            работа.append((з, если_есть))
        else:
            print("  %-6s %-5s свежая (сборка %d) — не трогаю"
                  % (з["порт"], з["вид"], своя))

    if смотреть:
        print()
        print("  (только смотрел; чтобы залить — запустите без --smotret)")
        return итог()

    if not работа:
        шаг("обновлять нечего", True, "все платы свежие")
        return итог()

    # Плата, которую заливаем впервые, ДОЛЖНА получить все четыре части:
    # у неё нет ни загрузчика, ни разметки.
    for з, ф in работа:
        print()
        print("ЗАЛИВАЮ %s (%s) — %s" % (з["порт"], з["вид"], ф["имя"]))
        print("-" * 70)
        if з["вид"] == "T114":
            вышло, слова = залить_t114(з["порт"], ф["путь"], ф.get("dfu"))
        else:
            впервые = с_нуля or (з["сборка"] is None)
            вышло, слова = залить_esp(з["порт"], з["вид"], ф["путь"], впервые)
        шаг("%s (%s) залита" % (з["порт"], з["вид"]), вышло, слова)
        if not вышло:
            continue

        сказать("жду, пока плата поднимется...")
        дождаться_порта(з["порт"], 35)
        time.sleep(2.0)
        стало = None
        try:
            с = открыть(з["порт"])
            time.sleep(1.8)
            стало = настройка(с, N_SBORKA)
            с.close()
        except Exception:
            pass
        шаг("%s отвечает после заливки" % з["порт"], стало == ф["версия"],
            ("сборка %s" % стало) if стало is not None else
            "молчит — если это T114, дайте ей секунду и запустите ещё раз")

    return итог()


def итог():
    print()
    print("=" * 70)
    вышло = sum(1 for _, в in ит if в)
    print("ИТОГ: %d из %d" % (вышло, len(ит)))
    for имя, в in ит:
        if not в:
            print("  не вышло: %s" % имя)
    print("=" * 70)
    return 0 if вышло == len(ит) else 1


if __name__ == "__main__":
    sys.exit(главное())
