#!/usr/bin/env python3
"""spec_client.py — тестовый клиент к spec_server (порт 5002).

    python spec_client.py --host rp-f0xxxx.local                   оболочка
    python spec_client.py --host 192.168.1.50 status               одна команда
    python spec_client.py --host 192.168.1.50 --run test_sequence.spec

Оболочка почти прозрачная: разбор команд живёт в сервере, здесь только отправка
строки и печать ответа. Явные do_*-методы нужны ради help и автодополнения по
Tab; всё, что им не подошло, уходит на сервер как есть (см. default).

Сверх серверных команд клиент понимает несколько собственных директив — они
нужны в файлах последовательностей и работают также из оболочки:

    echo <текст>          напечатать
    wait <мс>             пауза
    waitidle [мс]         ждать busy=0 (событие отработало)
    waitbuf [мс]          ждать buf_valid=0 (команда ушла из буфера в тракт)
    waitseq [мс]          ждать seq_busy=0 (секвенсор доиграл последовательность)
    run [-k] <файл>       выполнить файл построчно, ВЖИВУЮ (как раньше)
    load <файл>           тот же файл, но не исполнять, а ЗАГРУЗИТЬ в
                           аппаратный секвенсор (seqbegin/.../seqend) —
                           дальше seqrun запускает всё разом на 125 МГц, без
                           сети между шагами

Про waitbuf. Маска и длительность TTL — немедленные параметры: ttl_controller
защёлкивает их тогда, когда старт до него доходит. Пока команда лежит в буфере
axi_opcode_ctrl, она ещё ничего не защёлкнула, поэтому менять mask/dur до
снятия buf_valid нельзя — следующее событие уедет с чужой маской.

Про load. Файл — тот же самый .spec, что и для run: сервер сам разворачивает
set/start/op/defaults не в живой тракт, а в память секвенсора, пока идёт
загрузка (seqbegin..seqend) — компилировать на клиенте нечего. А вот
директивы, которые опрашивают ЖИВОЙ статус (wait/waitidle/waitbuf/waitseq),
внутри load-файла смысла не имеют: секвенсор в этот момент ничего не
исполняет, он просто пишет в BRAM. Не используйте их там — используйте после
seqrun, в оболочке или в отдельном live-файле.

Обычный сценарий: load prog.spec, затем seqrun, затем waitseq и waitidle
(секвенсор мог доиграть выдачу команд раньше, чем доиграет само последнее
событие — ждать нужно оба).
"""

import argparse
import cmd
import os
import socket
import sys
import time

DEFAULT_PORT = 5001
DEFAULT_TIMEOUT = 5.0
MAX_RUN_DEPTH = 8          # защита от run-петли: файл, зовущий сам себя
POLL_INTERVAL = 0.005      # шаг опроса status в waitidle/waitbuf


class SpecError(Exception):
    pass


class Link:
    """Строчный обмен с сервером: одна команда — один ответ."""

    def __init__(self, host, port, timeout=DEFAULT_TIMEOUT):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.sock = None
        self.fh = None

    def connect(self):
        self.close()
        self.sock = socket.create_connection((self.host, self.port), self.timeout)
        self.sock.settimeout(self.timeout)
        # newline='\n', чтобы не пришлось чистить \r в ответах вручную
        self.fh = self.sock.makefile("rw", encoding="utf-8", newline="\n")
        return self.readline()          # приветствие сервера

    def close(self):
        for obj in (self.fh, self.sock):
            if obj is not None:
                try:
                    obj.close()
                except OSError:
                    pass
        self.fh = None
        self.sock = None

    def readline(self):
        line = self.fh.readline()
        if not line:
            raise SpecError("сервер закрыл соединение")
        return line.rstrip("\n")

    def command(self, line):
        if self.fh is None:
            raise SpecError("нет соединения (connect)")
        try:
            self.fh.write(line + "\n")
            self.fh.flush()
            return self.readline()
        except socket.timeout:
            raise SpecError(f"нет ответа за {self.timeout} с")
        except OSError as exc:
            raise SpecError(str(exc))


def status_fields(answer):
    """'OK status=0x0a busy=0 ...' -> {'status': '0x0a', 'busy': '0', ...}"""
    fields = {}
    for token in answer.split():
        if "=" in token:
            key, _, val = token.partition("=")
            fields[key] = val
    return fields


class SpecShell(cmd.Cmd):
    prompt = "spec> "

    def __init__(self, link):
        super().__init__()
        self.link = link

    # ------------------------------------------------------ исполнение строки

    def exec_line(self, line, depth=0):
        """Выполнить строку: клиентскую директиву или команду сервера.

        Возвращает (успех, текст ответа). Печать — на вызывающем, чтобы run мог
        добавить к строке имя файла и номер.
        """
        parts = line.split(None, 1)
        if not parts:
            return True, ""
        verb, rest = parts[0], (parts[1].strip() if len(parts) > 1 else "")

        if verb == "echo":
            return True, rest

        if verb == "wait":
            try:
                ms = float(rest)
                if ms < 0:
                    raise ValueError
            except ValueError:
                return False, "! wait: нужно время в мс"
            time.sleep(ms / 1000.0)
            return True, f"пауза {ms:g} мс"

        if verb in ("waitidle", "waitbuf", "waitseq"):
            field, want, default_ms = {
                "waitidle": ("busy", "0", 5000.0),
                "waitbuf":  ("buf_valid", "0", 1000.0),
                # seq_busy обычно снимается быстро (это просто "выдал все
                # слова программы"), а не "доиграло последнее событие" —
                # после waitseq почти всегда нужен ещё и waitidle.
                "waitseq":  ("seq_busy", "0", 1000.0),
            }[verb]
            try:
                ms = float(rest) if rest else default_ms
            except ValueError:
                return False, f"! {verb}: нужен таймаут в мс"
            return self.wait_field(verb, field, want, ms)

        if verb == "run":
            # depth растёт именно здесь: run_script проверяет его, но сам вглубь
            # не спускается — без инкремента файл, зовущий сам себя, уходил в
            # RecursionError вместо аккуратной остановки.
            return self.do_run(rest, depth=depth + 1), ""

        if verb == "load":
            return self.do_load(rest, depth=depth + 1), ""

        try:
            answer = self.link.command(line)
        except SpecError as exc:
            return False, f"! {exc}"
        return not answer.startswith("ERR"), answer

    def wait_field(self, verb, field, want, timeout_ms):
        """Опрашивать status, пока поле не примет нужное значение."""
        started = time.monotonic()
        deadline = started + timeout_ms / 1000.0
        while True:
            try:
                answer = self.link.command("status")
            except SpecError as exc:
                return False, f"! {exc}"
            if answer.startswith("ERR"):
                return False, answer
            fields = status_fields(answer)
            if field not in fields:
                return False, f"! {verb}: в ответе нет поля {field}: {answer}"
            if fields[field] == want:
                waited = (time.monotonic() - started) * 1000.0
                return True, f"{field}={want} через {waited:.0f} мс"
            if time.monotonic() >= deadline:
                return False, (f"! {verb}: таймаут {timeout_ms:g} мс, "
                               f"{field}={fields[field]} ({answer})")
            time.sleep(POLL_INTERVAL)

    def send(self, line):
        ok, text = self.exec_line(line)
        if text:
            print(text)
        return ok

    def default(self, line):
        """Всё незнакомое уходит на сервер — он и рассудит."""
        if line == "EOF":
            return self.do_quit("")
        self.send(line)

    def emptyline(self):
        pass

    # --------------------------------------------------- файлы с командами

    def run_script(self, path, keep_going=False, depth=0):
        """Выполнить файл построчно. Возвращает True, если ошибок не было."""
        if depth > MAX_RUN_DEPTH:
            print(f"! run: слишком глубокая вложенность (> {MAX_RUN_DEPTH})")
            return False
        try:
            with open(path, "r", encoding="utf-8") as fh:
                lines = fh.readlines()
        except OSError as exc:
            print(f"! run: {exc}")
            return False

        name = os.path.basename(path)
        errors = 0
        executed = 0

        for lineno, raw in enumerate(lines, 1):
            line = raw.split("#", 1)[0].strip()
            if not line:
                continue
            executed += 1
            ok, text = self.exec_line(line, depth=depth)
            if line.split(None, 1)[0] == "echo":
                print(text)          # echo — заголовок раздела, колонки только мешают
            else:
                print(f"{name}:{lineno:<4} {line:28s} {text}".rstrip())
            if not ok:
                errors += 1
                if not keep_going:
                    print(f"! run: остановлено на {name}:{lineno} "
                          f"(-k — продолжать несмотря на ошибки)")
                    return False

        print(f"run: {name} — строк {executed}, ошибок {errors}")
        return errors == 0

    def do_run(self, arg, depth=0):
        """run [-k] <файл> — выполнить файл с командами построчно.

        Пустые строки и всё после "#" игнорируются. По умолчанию исполнение
        останавливается на первой ошибке; -k — продолжать до конца.
        """
        args = arg.split()
        keep_going = False
        while args and args[0] in ("-k", "--keep-going"):
            keep_going = True
            args.pop(0)
        if not args:
            print("! run: нужен путь к файлу")
            return False
        return self.run_script(" ".join(args), keep_going, depth)

    def do_load(self, arg, depth=0):
        """load <файл> — загрузить файл как программу в аппаратный
        секвенсор (seqbegin, построчно как в run, затем seqend).

        Компилировать на клиенте нечего: сервер сам разворачивает set/
        start/op/defaults не в живой тракт, а в память секвенсора, пока
        идёт загрузка — файл тот же самый, что и для run. Не запускает
        воспроизведение — для этого seqrun (или seqrun N с явным числом
        слов). Внутри файла не используйте wait/waitidle/waitbuf/waitseq:
        секвенсор в момент загрузки ничего не исполняет, эти директивы
        опрашивают живой статус и внутри load смысла не имеют.

        seqend вызывается даже при ошибке посреди файла — не оставляем
        сервер подвешенным в режиме загрузки.
        """
        args = arg.split()
        if not args:
            print("! load: нужен путь к файлу")
            return False
        path = " ".join(args)

        ok0, text0 = self.exec_line("seqbegin", depth=depth)
        if text0:
            print(text0)
        if not ok0:
            print("! load: seqbegin не удался, отменено")
            return False

        ok = self.run_script(path, keep_going=False, depth=depth)

        ok1, text1 = self.exec_line("seqend", depth=depth)
        if text1:
            print(text1)

        return ok and ok1

    # -------------------------------------------------------------- команды

    def do_status(self, arg):
        """status — состояние секвенсора (busy/ready/buf_valid/overflow)."""
        self.send("status " + arg)

    def do_params(self, arg):
        """params — теневые параметры, какими их помнит сервер."""
        self.send("params " + arg)

    def do_defaults(self, arg):
        """defaults — вернуть параметры тракта к умолчаниям."""
        self.send("defaults " + arg)

    def do_set(self, arg):
        """set dur|mask|freq|inc|phase|pulse|step|env <значения> — параметры события."""
        self.send("set " + arg)

    def do_start(self, arg):
        """start ttl|rf — запустить событие."""
        self.send("start " + arg)

    def do_seqbegin(self, arg):
        """seqbegin — начать загрузку последовательности в секвенсор
        (обнуляет указатель BRAM). Дальше set/start/op/defaults пишутся не в
        живой тракт, а в память секвенсора — до seqend. Обычно проще
        использовать директиву load <файл>, которая сама зовёт
        seqbegin/seqend вокруг файла.
        """
        self.send("seqbegin " + arg)

    def do_seqend(self, arg):
        """seqend — закончить загрузку, вернуться к живому режиму."""
        self.send("seqend " + arg)

    def do_seqrun(self, arg):
        """seqrun [N] — запустить загруженную последовательность.

        Без аргумента — сервер помнит, сколько слов было загружено с
        последнего seqbegin. N — сыграть только первые N слов (например, для
        отладки не всей последовательности целиком).
        """
        self.send("seqrun " + arg)

    def do_seqabort(self, arg):
        """seqabort — немедленно остановить проигрывание последовательности
        (уже идущее событие доиграет до конца, как обычно — останавливается
        только выдача ДАЛЬНЕЙШИХ слов программы)."""
        self.send("seqabort " + arg)

    def do_op(self, arg):
        """op <код> <параметр> — сырая команда, без проверки контрактов."""
        self.send("op " + arg)

    def do_peek(self, arg):
        """peek <смещение> — прочитать регистр (0x00 OPCODE, 0x04 PARAM, 0x08 STATUS)."""
        self.send("peek " + arg)

    def do_poke(self, arg):
        """poke <смещение> <значение> — записать регистр напрямую."""
        self.send("poke " + arg)

    def do_echo(self, arg):
        """echo <текст> — напечатать строку (полезно в файлах последовательностей)."""
        self.send("echo " + arg)

    def do_wait(self, arg):
        """wait <мс> — пауза."""
        self.send("wait " + arg)

    def do_waitidle(self, arg):
        """waitidle [мс] — ждать busy=0, то есть конца всех событий."""
        self.send("waitidle " + arg)

    def do_waitbuf(self, arg):
        """waitbuf [мс] — ждать buf_valid=0: команда ушла из буфера в тракт.

        После этого можно менять маску и длительность следующего события, не
        испортив уже поставленное в очередь.
        """
        self.send("waitbuf " + arg)

    def do_waitseq(self, arg):
        """waitseq [мс] — ждать seq_busy=0: секвенсор выдал все слова программы.

        Это НЕ то же самое, что "последнее событие доиграло" — оно может ещё
        идти. После seqrun обычно нужны оба: waitseq, затем waitidle.
        """
        self.send("waitseq " + arg)

    def do_pulse(self, arg):
        """pulse <мкс> [частота Гц] [маска] — настроить и выдать один РЧ-импульс.

        Составная команда клиента: окно TTL берётся с запасом 20 % от длительности
        огибающей, иначе сервер отвергнет старт по контракту 2 (РЧ не должно
        переживать окно TTL).
        """
        args = arg.split()
        if not args:
            print("! нужна длительность в мкс")
            return
        try:
            us = float(args[0])
            freq = float(args[1]) if len(args) > 1 else None
            mask = args[2] if len(args) > 2 else None
        except ValueError:
            print("! не разобраны аргументы")
            return

        window = max(1, int(us * 1.2 + 0.5))
        steps = [f"set dur {window}", f"set pulse {us}"]
        if freq is not None:
            steps.append(f"set freq {freq}")
        if mask is not None:
            steps.append(f"set mask {mask}")
        steps.append("start rf")

        for step in steps:
            ok, text = self.exec_line(step)
            print(f"{step:22s} -> {text}")
            if not ok:
                print("! последовательность прервана")
                return

    def do_connect(self, arg):
        """connect [хост [порт]] — переподключиться."""
        args = arg.split()
        if args:
            self.link.host = args[0]
        if len(args) > 1:
            self.link.port = int(args[1])
        try:
            print(self.link.connect())
        except (OSError, SpecError) as exc:
            print(f"! не подключиться к {self.link.host}:{self.link.port}: {exc}")

    def do_quit(self, arg):
        """quit — закрыть соединение и выйти."""
        try:
            print(self.link.command("quit"))
        except SpecError:
            pass
        self.link.close()
        return True

    do_exit = do_quit
    do_EOF = do_quit


def main():
    ap = argparse.ArgumentParser(description="тестовый клиент spec_server")
    ap.add_argument("--host", default="rp-f0xxxx.local", help="адрес платы")
    ap.add_argument("--port", type=int, default=DEFAULT_PORT)
    ap.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    ap.add_argument("--run", metavar="ФАЙЛ",
                    help="выполнить файл с командами и выйти")
    ap.add_argument("-k", "--keep-going", action="store_true",
                    help="с --run: не останавливаться на ошибках")
    ap.add_argument("command", nargs="*",
                    help="если задана — выполнить одну команду и выйти")
    args = ap.parse_args()

    link = Link(args.host, args.port, args.timeout)
    try:
        hello = link.connect()
    except (OSError, SpecError) as exc:
        print(f"! не подключиться к {args.host}:{args.port}: {exc}", file=sys.stderr)
        return 1

    shell = SpecShell(link)
    try:
        if args.run:
            print(hello)
            return 0 if shell.run_script(args.run, args.keep_going) else 1

        if args.command:
            line = " ".join(args.command)
            ok, text = shell.exec_line(line)
            print(text)
            return 0 if ok else 2

        shell.intro = f"{hello}\nhelp — команды сервера, ? — команды клиента"
        shell.cmdloop()
        return 0
    finally:
        link.close()


if __name__ == "__main__":
    sys.exit(main())