"""手动设置阀门开度，按 --dt 周期写 CSV，终端每秒显示一次。

示例：
    python test_valve.py
    python test_valve.py --opening 60 --dt 0.01
"""

import argparse
import contextlib
import csv
from datetime import datetime
import io
import math
from pathlib import Path
import sys
import time

try:
    import msvcrt
except ImportError:  # 本项目的交互方式仅适用于 Windows 终端
    msvcrt = None

import config
from flow_control import opening_to_motor_position


SAMPLE_PERIOD_S = 1.0          # CSV 采样周期默认值（秒），可用 --dt 覆盖
PRINT_PERIOD_S = 1.0           # 终端打印周期（秒），与采样周期解耦
KEY_POLL_PERIOD_S = 0.05
INITIAL_OPENING_PCT = 100.0    # 初始开度默认值（%），可用 --opening 覆盖
OUTPUT_DIRECTORY = Path(__file__).resolve().parent / "test_valve_data"


class ManualValveRecorder:
    """逐次保存手动测试数据，并在结束时生成三联图。"""

    CSV_FIELDS = (
        "time_s",
        "timestamp",
        "flow_slm",
        "opening_pct",
        "pressure_kpa",
        "motor_position",
    )

    def __init__(self, output_directory):
        output_directory.mkdir(parents=True, exist_ok=True)
        run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.csv_path = output_directory / f"test_valve_{run_id}.csv"
        self.image_path = output_directory / f"test_valve_{run_id}.png"
        self._file = self.csv_path.open("w", newline="", encoding="utf-8-sig")
        self._writer = csv.DictWriter(self._file, fieldnames=self.CSV_FIELDS)
        self._writer.writeheader()
        self._file.flush()
        self._start_time = None
        self.sample_count = 0
        self._closed = False

    def record(self, sample_time, flow, opening, pressure, position):
        if self._start_time is None:
            self._start_time = sample_time
        self._writer.writerow(
            {
                "time_s": f"{sample_time - self._start_time:.6f}",
                "timestamp": datetime.now().isoformat(timespec="milliseconds"),
                "flow_slm": self._format_number(flow),
                "opening_pct": self._format_number(opening),
                "pressure_kpa": self._format_number(pressure),
                "motor_position": self._format_number(position),
            }
        )
        self._file.flush()
        self.sample_count += 1

    def finalize(self):
        self.close()
        if self.sample_count == 0:
            return None
        self._create_plot()
        return self.image_path

    def close(self):
        if not self._closed:
            self._file.flush()
            self._file.close()
            self._closed = True

    def _create_plot(self):
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        time_s = []
        flow = []
        opening = []
        pressure = []
        with self.csv_path.open("r", newline="", encoding="utf-8-sig") as file:
            for row in csv.DictReader(file):
                time_s.append(float(row["time_s"]))
                flow.append(self._float_or_nan(row["flow_slm"]))
                opening.append(self._float_or_nan(row["opening_pct"]))
                pressure.append(self._float_or_nan(row["pressure_kpa"]))

        figure, axes = plt.subplots(
            3,
            1,
            figsize=(12, 10),
            sharex=True,
            constrained_layout=True,
        )
        figure.suptitle("Manual Valve Test", fontsize=15)

        axes[0].plot(time_s, flow, color="#1565C0", linewidth=1.4)
        axes[0].set_ylabel("Flow (SLM)")
        axes[0].set_title("Flow")

        axes[1].plot(time_s, opening, color="#2E7D32", linewidth=1.4)
        axes[1].set_ylabel("Opening (%)")
        axes[1].set_title("Valve opening")
        axes[1].set_ylim(-2, 102)

        axes[2].plot(time_s, pressure, color="#6A1B9A", linewidth=1.4)
        axes[2].set_ylabel("Pressure (kPa)")
        axes[2].set_xlabel("Time (s)")
        axes[2].set_title("Pressure")

        for axis in axes:
            axis.grid(True, alpha=0.3, linestyle="--")
            axis.margins(x=0)

        figure.savefig(self.image_path, dpi=config.PLOT_DPI, bbox_inches="tight")
        plt.close(figure)

    @staticmethod
    def _format_number(value):
        if value is None:
            return ""
        try:
            number = float(value)
        except (TypeError, ValueError):
            return ""
        return f"{number:.6f}" if math.isfinite(number) else ""

    @staticmethod
    def _float_or_nan(value):
        try:
            return float(value)
        except (TypeError, ValueError):
            return math.nan


def format_value(value, unit):
    """将传感器读数格式化；读取失败时显示 --。"""
    if value is None:
        return f"-- {unit}"
    try:
        number = float(value)
    except (TypeError, ValueError):
        return f"-- {unit}"
    if not math.isfinite(number):
        return f"-- {unit}"
    return f"{number:.2f} {unit}"


def motor_position_for(opening_pct):
    return opening_to_motor_position(
        opening_pct,
        config.MOTOR_OPEN_POSITION,
        config.MOTOR_CLOSED_POSITION,
    )


def set_opening(hardware, opening_pct):
    """校验并下发开度，成功时返回对应的电机行程。"""
    opening = float(opening_pct)
    if not math.isfinite(opening) or not 0.0 <= opening <= 100.0:
        raise ValueError("开度必须是 0~100 之间的有限数值")

    position = motor_position_for(opening)
    if not hardware.set_motor_position(
        position,
        channel=config.MOTOR_OUTPUT_CHANNEL,
    ):
        raise RuntimeError("电机位置写入失败")
    return opening, position


def read_sensors(hardware, quiet=False):
    """读取一次传感器；输入模式中丢弃底层硬件代码的终端输出。"""
    output = io.StringIO()
    stream = output if quiet else sys.stdout
    with contextlib.redirect_stdout(stream):
        flow = hardware.get_flow(config.FLOW_INPUT_CHANNEL)
        pressure = (
            None
            if config.PRESSURE_INPUT_CHANNEL is None
            else hardware.get_pressure(config.PRESSURE_INPUT_CHANNEL)
        )
    return flow, pressure


def parse_args(argv=None):
    """解析命令行参数：--dt 采样周期（秒）、--opening 初始开度（%）。"""
    parser = argparse.ArgumentParser(
        description="手动设置阀门开度：按 --dt 周期写 CSV，终端每秒显示一次。"
    )
    parser.add_argument(
        "--dt",
        type=float,
        default=SAMPLE_PERIOD_S,
        help=f"CSV 采样周期（秒），默认 {SAMPLE_PERIOD_S:g}",
    )
    parser.add_argument(
        "--opening",
        type=float,
        default=INITIAL_OPENING_PCT,
        help=f"初始阀门开度（%%），默认 {INITIAL_OPENING_PCT:g}",
    )
    args = parser.parse_args(argv)
    if args.dt <= 0:
        parser.error("--dt 必须大于 0")
    if not 0.0 <= args.opening <= 100.0:
        parser.error("--opening 必须位于 0~100 之间")
    return args


def main(argv=None):
    args = parse_args(argv)

    if msvcrt is None:
        print("本程序需要在 Windows 终端中运行。")
        return 1

    try:
        from PcControl import MT2AM8Client
    except ModuleNotFoundError as exc:
        print(f"无法导入硬件客户端: {exc}")
        return 2

    hardware = MT2AM8Client(
        host=config.MT2AM8_HOST,
        port=config.MT2AM8_PORT,
        slave_id=config.MT2AM8_SLAVE_ID,
        pressure_range=config.PRESSURE_RANGE_KPA,
        flow_range=config.FLOW_METER_RANGE_SLM,
    )

    connected = False
    opening = args.opening
    position = motor_position_for(opening)
    input_mode = False
    input_buffer = ""
    recorder = None
    exit_code = 0

    try:
        print(
            f"正在连接 MT2-AM8 {config.MT2AM8_HOST}:{config.MT2AM8_PORT} ..."
        )
        connected = bool(hardware.connect())
        if not connected:
            print("连接失败。")
            return 3

        try:
            opening, position = set_opening(hardware, opening)
        except (ValueError, RuntimeError) as exc:
            print(f"初始开度设置失败：{exc}")
            return 4

        recorder = ManualValveRecorder(OUTPUT_DIRECTORY)
        print("手动阀门控制已启动：按 S 输入新开度，按 Ctrl+C 退出。")
        print(
            f"初始开度已下发：{opening:.2f}%"
            f"（行程 {position:.1f}）。"
        )
        print(f"CSV 将保存到：{recorder.csv_path}")

        next_sample_time = time.perf_counter()
        next_print_time = time.perf_counter()
        last_flow = None
        last_pressure = None

        while True:
            while msvcrt.kbhit():
                ch = msvcrt.getch()

                if not input_mode:
                    if ch == b"\x03":
                        raise KeyboardInterrupt
                    if ch.decode("ascii", errors="ignore").lower() == "s":
                        input_mode = True
                        input_buffer = ""
                        print(
                            f"当前开度={opening:.2f}%，请输入新开度 0~100 "
                            "(回车确认, Esc 取消): ",
                            end="",
                            flush=True,
                        )
                    continue

                if ch in (b"\r", b"\n"):
                    print()
                    input_mode = False
                    raw = input_buffer.strip()
                    if not raw:
                        print("未输入开度，保持原开度。")
                        continue
                    try:
                        opening, position = set_opening(hardware, float(raw))
                    except (ValueError, RuntimeError) as exc:
                        print(f"开度切换失败：{exc}；保持原开度。")
                    else:
                        print(
                            f"开度已切换为 {opening:.2f}%"
                            f"（行程 {position:.1f}）。"
                        )
                elif ch == b"\x1b":
                    input_mode = False
                    print(" (已取消，保持原开度)")
                elif ch in (b"\x08", b"\x7f"):
                    if input_buffer:
                        input_buffer = input_buffer[:-1]
                        print("\b \b", end="", flush=True)
                elif ch == b"\x03":
                    raise KeyboardInterrupt
                else:
                    text = ch.decode("ascii", errors="ignore")
                    if text and text.isprintable():
                        input_buffer += text
                        print(text, end="", flush=True)

            now = time.perf_counter()

            if now >= next_sample_time:
                last_flow, last_pressure = read_sensors(
                    hardware, quiet=input_mode
                )
                recorder.record(
                    now, last_flow, opening, last_pressure, position
                )
                # 不追补错过的采样；输入期间的样本只写 CSV，不打印终端。
                next_sample_time = time.perf_counter() + args.dt

            if not input_mode and now >= next_print_time:
                timestamp = time.strftime("%H:%M:%S")
                print(
                    f"[{timestamp}] "
                    f"流量={format_value(last_flow, 'SLM')}  "
                    f"压力={format_value(last_pressure, 'kPa')}  "
                    f"开度={opening:.2f}%  行程={position:.1f}"
                )
                next_print_time = time.perf_counter() + PRINT_PERIOD_S

            time.sleep(min(KEY_POLL_PERIOD_S, args.dt))

    except KeyboardInterrupt:
        if input_mode:
            print()
        print(
            "收到 Ctrl+C，正在保存数据并退出；退出前不再下发开度命令。"
            "断开连接后的阀门位置由硬件决定。"
        )
    except Exception as exc:
        if input_mode:
            print()
        print(f"程序异常：{exc}")
        exit_code = 4
    finally:
        if connected:
            try:
                hardware.disconnect()
            except Exception as exc:
                print(f"警告：断开设备失败：{exc}")
                exit_code = max(exit_code, 5)
        if recorder is not None:
            try:
                image_path = recorder.finalize()
                print(f"CSV 已保存：{recorder.csv_path}")
                if image_path is not None:
                    print(f"三联图已保存：{image_path}")
            except Exception as exc:
                recorder.close()
                print(f"警告：生成图表失败：{exc}")
                exit_code = max(exit_code, 6)

    return exit_code


if __name__ == "__main__":
    raise SystemExit(main())