#!/usr/bin/env python3
import argparse
from typing import Literal

import fn_gen.errors as fg_err
from fn_gen import DG2052

from common import close_output, get_postamplified, get_preamplified
from examples.discrete_sweep import AMPLIFICATION


def closeOutput(fg: DG2052, channel: Literal[1, 2]):
    fg.set_output(channel, False)
    print(f"Output{channel} State: {fg.get_output_state(channel)}")


def generateDCSignal(v_dc):
    fg = DG2052("TCPIP::192.168.1.11::INSTR")
    v_dc = get_preamplified(13.0, v_dc)
    channel = 2
    try:
        print(fg.whoami())
        print(f"\nOutput{channel} Impedance: {fg.get_output_impedance(channel)} Ohm")
        print(f"Output{channel} Load: {fg.get_output_load(channel)} Ohm")
        print(f"Output{channel} Voltage Limits: {fg.get_output_volt_limits(channel)} V")
        fg.set_dc(channel, v_dc)
        print(
            f"Output{channel}: {fg.get_output_signal(channel)} | {fg.get_output_state(channel)}"
        )
        fg.set_output(channel, True)
        print(f"Voltage: {get_postamplified(AMPLIFICATION, v_dc):.2f} V")
        print(f"Output{channel} State: {fg.get_output_state(channel)}")
        while True:
            pass
    except fg_err.ValueOutOfBoundsError as err:
        print(err)
    except fg_err.UndefinedValueError as err:
        print(err)
    except KeyboardInterrupt:
        close_output(fg, channel)
    finally:
        close_output(fg, channel)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="This program is for testing the DG2052 function generator library. It generates a DC signal based on the voltage value supplied."
    )
    parser.add_argument(
        "-v",
        "--voltage",
        type=float,
        required=True,
        help="The DC voltage supplied",
    )
    args = parser.parse_args()
    generateDCSignal(args.voltage)
