"""CLI tool for controlling the ITECH IT6500 Series DC Power Supply via USB."""

from __future__ import annotations

import argparse
import csv
import sys
import time

from it6500.driver import IT6500


def find_instrument(address: str | None) -> str:
    """Find the USB-TMC instrument address or use the provided one."""
    if address:
        return address

    import pyvisa

    rm = pyvisa.ResourceManager()
    resources = rm.list_resources()
    rm.close()

    # Look for ITECH by USB VID/PID
    itech = [r for r in resources if "2EC7" in r.upper() or "6522" in r.upper()]
    if itech:
        print(f"Found ITECH IT6500: {itech[0]}")
        return itech[0]

    usb = [r for r in resources if "USB" in r]
    if not usb:
        print("No ITECH IT6500 or USB instruments found. Available resources:")
        for r in resources:
            print(f"  {r}")
        sys.exit(1)

    if len(usb) == 1:
        print(f"Found USB instrument: {usb[0]}")
        return usb[0]

    print("Multiple USB instruments found:")
    for i, r in enumerate(usb):
        print(f"  [{i}] {r}")
    choice = input("Select instrument number: ")
    return usb[int(choice)]


def _safe_query(psu: IT6500, method, fallback="N/A"):
    """Call a query method, returning fallback on timeout/error."""
    try:
        return method()
    except Exception:
        return fallback


def cmd_identify(psu: IT6500, _args: argparse.Namespace) -> None:
    """Print instrument identity and status."""
    print(f"Identity:    {psu.idn()}")
    print(f"SCPI Ver:    {_safe_query(psu, psu.get_version)}")
    print(f"Output:      {'ON' if psu.get_output_state() else 'OFF'}")
    print(f"Voltage set: {psu.get_voltage():.4f} V")
    print(f"Current set: {psu.get_current():.4f} A")

    op = _safe_query(psu, psu.get_operation_status, {})
    if isinstance(op, dict):
        mode = "CV" if op.get("CV") else ("CC" if op.get("CC") else "---")
        print(f"Mode:        {mode}")

    qs = _safe_query(psu, psu.get_questionable_status, {})
    if isinstance(qs, dict):
        active = [k for k, v in qs.items() if v]
        print(f"Protection:  {', '.join(active) if active else 'None'}")

    err_code, err_msg = psu.get_error()
    if err_code != 0:
        print(f"Error:       ({err_code}) {err_msg}")


def cmd_measure(psu: IT6500, _args: argparse.Namespace) -> None:
    """Take a single measurement."""
    result = psu.measure_all()
    print(f"  Voltage = {result.voltage:>10.4f} V")
    print(f"  Current = {result.current:>10.4f} A")
    print(f"  Power   = {result.power:>10.4f} W")


def cmd_monitor(psu: IT6500, args: argparse.Namespace) -> None:
    """Continuously monitor measurements at an interval."""
    interval = args.interval

    writer = None
    outfile = None
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(["timestamp", "voltage_V", "current_A", "power_W"])
        print(f"Logging to {args.output}")

    print(f"{'Time':>10s}  {'Voltage(V)':>12s}  {'Current(A)':>12s}  {'Power(W)':>12s}")
    print("-" * 52)

    try:
        count = 0
        while args.count == 0 or count < args.count:
            result = psu.measure_all()
            ts = time.strftime("%H:%M:%S")
            print(
                f"{ts:>10s}  {result.voltage:>12.4f}  "
                f"{result.current:>12.4f}  {result.power:>12.4f}"
            )

            if writer:
                writer.writerow([
                    time.strftime("%Y-%m-%d %H:%M:%S"),
                    f"{result.voltage:.4f}",
                    f"{result.current:.4f}",
                    f"{result.power:.4f}",
                ])
                outfile.flush()

            count += 1
            if args.count == 0 or count < args.count:
                time.sleep(interval)
    except KeyboardInterrupt:
        print("\nMonitoring stopped.")
    finally:
        if outfile:
            outfile.close()
            print(f"Data saved to {args.output}")


def cmd_live(psu: IT6500, args: argparse.Namespace) -> None:
    """Live monitor with real-time graph."""
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from collections import deque

    max_points = args.history
    interval_ms = int(args.interval * 1000)

    timestamps: deque[float] = deque(maxlen=max_points)
    voltages: deque[float] = deque(maxlen=max_points)
    currents: deque[float] = deque(maxlen=max_points)
    powers: deque[float] = deque(maxlen=max_points)
    t0 = time.time()

    writer = None
    outfile = None
    if args.output:
        outfile = open(args.output, "w", newline="")
        writer = csv.writer(outfile)
        writer.writerow(["timestamp", "voltage_V", "current_A", "power_W"])
        print(f"Logging to {args.output}")

    fig, axes = plt.subplots(3, 1, figsize=(12, 9), squeeze=False)
    fig.suptitle("ITECH IT6500 Live Monitor", fontsize=14, fontweight="bold")
    axes = axes.flatten()

    labels = [("Voltage", "V", voltages), ("Current", "A", currents), ("Power", "W", powers)]
    lines = []
    for ax, (title, unit, _) in zip(axes, labels):
        ax.set_ylabel(f"{title} ({unit})")
        ax.set_xlabel("Time (s)")
        ax.grid(True, alpha=0.3)
        ax.set_title(title, fontsize=11)
        line, = ax.plot([], [], linewidth=1.5)
        lines.append(line)
        ax.legend([title], loc="upper left", fontsize=9)

    fig.tight_layout()

    def update(_frame):
        try:
            result = psu.measure_all()
        except Exception as e:
            print(f"Read error: {e}")
            return lines

        now = time.time() - t0
        timestamps.append(now)
        voltages.append(result.voltage)
        currents.append(result.current)
        powers.append(result.power)

        ts = time.strftime("%H:%M:%S")
        print(
            f"{ts}  V={result.voltage:.4f}  "
            f"I={result.current:.4f}  P={result.power:.4f}"
        )

        if writer:
            writer.writerow([
                time.strftime("%Y-%m-%d %H:%M:%S"),
                f"{result.voltage:.4f}",
                f"{result.current:.4f}",
                f"{result.power:.4f}",
            ])
            outfile.flush()

        t_list = list(timestamps)
        for line, (_, _, data) in zip(lines, labels):
            line.set_data(t_list, list(data))

        for ax in axes:
            ax.relim()
            ax.autoscale_view()

        return lines

    _anim = FuncAnimation(fig, update, interval=interval_ms, blit=False, cache_frame_data=False)

    try:
        plt.show()
    except KeyboardInterrupt:
        pass
    finally:
        if outfile:
            outfile.close()
            print(f"Data saved to {args.output}")


def cmd_set(psu: IT6500, args: argparse.Namespace) -> None:
    """Set voltage and/or current."""
    if args.voltage is not None and args.current is not None:
        psu.apply(args.voltage, args.current)
        print(f"Applied: {args.voltage:.4f} V, {args.current:.4f} A")
    elif args.voltage is not None:
        psu.set_voltage(args.voltage)
        print(f"Voltage set: {args.voltage:.4f} V")
    elif args.current is not None:
        psu.set_current(args.current)
        print(f"Current set: {args.current:.4f} A")
    else:
        print("Specify --voltage and/or --current")
        sys.exit(1)


def cmd_output(psu: IT6500, args: argparse.Namespace) -> None:
    """Turn output ON or OFF."""
    if args.state.upper() in ("ON", "1"):
        psu.output_on()
        print("Output ON")
    else:
        psu.output_off()
        print("Output OFF")


def cmd_protection(psu: IT6500, args: argparse.Namespace) -> None:
    """Check or clear protection status."""
    if args.clear:
        try:
            psu.clear_ovp()
        except Exception:
            pass
        psu.clear_status()
        print("Protection cleared.")
    else:
        qs = psu.get_questionable_status()
        for name, active in qs.items():
            status = "TRIGGERED" if active else "OK"
            print(f"  {name}: {status}")

        ovp_tripped = _safe_query(psu, psu.get_ovp_tripped, False)
        print(f"  OVP tripped: {'YES' if ovp_tripped else 'NO'}")

        ovp_level = _safe_query(psu, psu.get_ovp_level, None)
        if ovp_level is not None:
            print(f"\n  OVP level:   {ovp_level:.2f} V")
        ovp_state = _safe_query(psu, psu.get_ovp_state, None)
        if ovp_state is not None:
            print(f"  OVP enabled: {'YES' if ovp_state else 'NO'}")
        ovp_delay = _safe_query(psu, psu.get_ovp_delay, None)
        if ovp_delay is not None:
            print(f"  OVP delay:   {ovp_delay:.3f} s")


def cmd_config(psu: IT6500, args: argparse.Namespace) -> None:
    """View or change instrument configuration."""
    if args.ovp is not None:
        psu.set_ovp_level(args.ovp)
        psu.set_ovp_state(True)
        print(f"OVP set to {args.ovp:.2f} V (enabled)")
    if args.rise is not None:
        psu.set_rise_time(args.rise)
        print(f"Rise time: {args.rise:.3f} s")
    if args.fall is not None:
        psu.set_fall_time(args.fall)
        print(f"Fall time: {args.fall:.3f} s")
    if args.avg is not None:
        psu.set_averaging(args.avg)
        print(f"Averaging: {args.avg}")
    if args.vrange is not None:
        psu.set_voltage_range(args.vrange)
        print(f"Voltage range: {args.vrange:.2f} V")

    # Show current config
    print(f"\nCurrent configuration:")
    print(f"  Voltage set:  {psu.get_voltage():.4f} V")
    print(f"  Current set:  {psu.get_current():.4f} A")
    for label, method in [
        ("V range", psu.get_voltage_range),
        ("Rise time", psu.get_rise_time),
        ("Fall time", psu.get_fall_time),
        ("OVP level", psu.get_ovp_level),
    ]:
        val = _safe_query(psu, method, None)
        if val is not None:
            print(f"  {label + ':':12s} {val}")
    ovp_state = _safe_query(psu, psu.get_ovp_state, None)
    if ovp_state is not None:
        print(f"  OVP enabled:  {'YES' if ovp_state else 'NO'}")
    avg = _safe_query(psu, psu.get_averaging, None)
    if avg is not None:
        print(f"  Averaging:    {avg}")
    beep = _safe_query(psu, psu.get_beeper, None)
    if beep is not None:
        print(f"  Beeper:       {'ON' if beep else 'OFF'}")


def cmd_send(psu: IT6500, args: argparse.Namespace) -> None:
    """Send a raw SCPI command."""
    command = " ".join(args.raw_command)
    result = psu.send(command)
    if result is not None:
        print(f"Response: {result}")
    else:
        print("OK")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="ITECH IT6500 Series DC Power Supply USB Control Tool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
examples:
  %(prog)s identify
  %(prog)s measure
  %(prog)s monitor --interval 1.0 --output data.csv
  %(prog)s live --interval 0.5
  %(prog)s set --voltage 12.0 --current 5.0
  %(prog)s output on
  %(prog)s output off
  %(prog)s protection
  %(prog)s protection --clear
  %(prog)s config --ovp 60.0 --rise 0.1 --fall 0.1
  %(prog)s send "*IDN?"
  %(prog)s send "VOLT 24.0"
""",
    )
    parser.add_argument(
        "-a", "--address",
        help="VISA resource address (e.g., USB0::0x2EC7::0x6522::...::INSTR). Auto-detects if omitted.",
    )
    parser.add_argument(
        "--timeout", type=int, default=5000,
        help="Communication timeout in ms (default: 5000)",
    )

    sub = parser.add_subparsers(dest="command", required=True)

    # identify
    sub.add_parser("identify", help="Show instrument identity and status")

    # measure
    sub.add_parser("measure", help="Take a single measurement")

    # monitor
    p_mon = sub.add_parser("monitor", help="Continuously monitor measurements")
    p_mon.add_argument("-i", "--interval", type=float, default=1.0, help="Seconds between readings (default: 1.0)")
    p_mon.add_argument("-n", "--count", type=int, default=0, help="Number of readings (0=infinite)")
    p_mon.add_argument("-o", "--output", help="CSV output file path")

    # live
    p_live = sub.add_parser("live", help="Live monitor with real-time graph")
    p_live.add_argument("-i", "--interval", type=float, default=1.0, help="Seconds between readings (default: 1.0)")
    p_live.add_argument("-o", "--output", help="CSV output file path")
    p_live.add_argument("--history", type=int, default=300, help="Max data points to display (default: 300)")

    # set
    p_set = sub.add_parser("set", help="Set voltage and/or current")
    p_set.add_argument("-v", "--voltage", type=float, help="Output voltage in Volts")
    p_set.add_argument("-c", "--current", type=float, help="Output current in Amps")

    # output
    p_out = sub.add_parser("output", help="Turn output ON or OFF")
    p_out.add_argument("state", choices=["on", "off", "ON", "OFF", "1", "0"], help="ON or OFF")

    # protection
    p_prot = sub.add_parser("protection", help="Check or clear protection status")
    p_prot.add_argument("--clear", action="store_true", help="Clear protection flags")

    # config
    p_cfg = sub.add_parser("config", help="View or change instrument configuration")
    p_cfg.add_argument("--ovp", type=float, help="Set OVP level (Volts)")
    p_cfg.add_argument("--rise", type=float, help="Set voltage rise time (seconds)")
    p_cfg.add_argument("--fall", type=float, help="Set voltage fall time (seconds)")
    p_cfg.add_argument("--avg", type=int, help="Set measurement averaging count (0-15)")
    p_cfg.add_argument("--vrange", type=float, help="Set voltage range upper limit (Volts)")

    # send
    p_send = sub.add_parser("send", help="Send raw SCPI command")
    p_send.add_argument("raw_command", nargs="+", help="Command string (queries auto-detected by '?')")

    args = parser.parse_args()
    address = find_instrument(args.address)

    dispatch = {
        "identify": cmd_identify,
        "measure": cmd_measure,
        "monitor": cmd_monitor,
        "live": cmd_live,
        "set": cmd_set,
        "output": cmd_output,
        "protection": cmd_protection,
        "config": cmd_config,
        "send": cmd_send,
    }

    with IT6500(address, timeout_ms=args.timeout) as psu:
        psu.remote()
        dispatch[args.command](psu, args)


if __name__ == "__main__":
    main()