import os
import time
import hid
import glfw
from OpenGL.GL import *
import imgui
import logging
from imgui.integrations.glfw import GlfwRenderer
from image import load_image
from image import load_image_rgba
from image import render_image
from proximity import reset_proximity
from proximity import proximity_read_proximity
from proximity import proximity_write_to_headset
from proximity import proximity_threshold_write_to_headset
from alignment import align_init
from alignment import align_find_hmd_horizon
from alignment import align_get_hmd_horizon
from alignment import align_destroy
from alignment import align_show_horizon
from alignment import align_minimize_pitch
from alignment import align_minimize_yaw
from alignment import align_minimize_roll
from alignment import align_get_left_config_string
from alignment import align_get_right_config_string
from alignment import align_set_best_roll
from alignment import confirm_distortion
from alignment import confirm_tracking_calibration
from hmd_config import hmd_config_write_serial
# from mes_logger import mes_add

LEFT_EYE = 0
RIGHT_EYE = 1

window_width = 1200
window_height = 1000
hid_device = hid.device()

pitch_left = 0.0
pitch_right = 0.0
yaw_left = 0.0
yaw_right = 0.0
roll_left = 0.0
roll_right = 0.0

idle_time = 0.0  # Used to automatically advance past confirmation screens
max_idle_time = 3.0  # Seconds

canting_override = 0.0
input_ipd = 0.0
serial_number = ""
should_check_red_distortion = False

# MES values for proximity
MES_OP = "20101883"
MES_STATION_ID = "T060 VAP01"
MES_SW_VER = "0.0.1"
uut_start = 0
uut_stop = 0


class STATES:
    PROMPT = 0
    CONNECT = 1
    PROXIMITY_RESET = 2
    PROXIMITY_NO_PAD = 3
    PROXIMITY_NO_PAD_CALC = 4
    PROXIMITY_RESET_BEFORE_PAD = 5
    PROXIMITY_PAD = 6
    PROXIMITY_PAD_CALC = 7
    PROXIMITY_REJECTED = 8
    PROXIMITY_WRITE_THRESHOLD = 9
    ALIGN_INIT = 10
    ALIGN_HORIZON = 11
    ALIGN_PITCH = 12
    ALIGN_YAW = 13
    ALIGN_ROLL = 14
    ALIGN_SHOW_RESULTS = 15
    CALIBRATION_REJECTED = 16


def centered(width):
    imgui.set_cursor_pos_x((window_width - width) * 0.5)


def text_centered(text):
    text_width = imgui.calc_text_size(text).x
    imgui.set_cursor_pos_x((window_width - text_width) * 0.5)
    imgui.text(text)


def render_ui():
    imgui.end()
    imgui.end_frame()
    imgui_renderer.process_inputs()
    imgui.render()
    imgui_renderer.render(imgui.get_draw_data())
    glfw.poll_events()
    glfw.swap_buffers(glfw_window)
    glClearColor(1.0, 1.0, 1.0, 1.0)
    glClear(GL_COLOR_BUFFER_BIT)
    imgui.new_frame()
    imgui.set_next_window_size(imgui.get_io().display_size.x,
                               imgui.get_io().display_size.y)
    imgui.set_next_window_position(0, 0)
    imgui.begin("Bigscreen VAP Tool", True, imgui.WINDOW_NO_COLLAPSE | imgui.WINDOW_NO_RESIZE)


# Before calibration, we prompt for input. In production, this value should
# always be the serial number, but for ease of testing we can also input
# a manual canting angle or manual IPD directly.
def parse_input(input_to_parse):
    global canting_override
    global input_ipd
    global serial_number

    try:
        number = int(input_to_parse)
    except ValueError:
        number = None

    match number:
        case num if num is not None and num < 20:
            print("Input less than 20; interpreting as canting angle.")
            canting_override = num
        case num if num is not None and 20 <= num < 100:
            print("Input number greater than 20; interpreting as IPD to determine canting.")
            input_ipd = num
        case _:
            print("Input serial number; extracting IPD to determine canting.")
            serial_number = input_to_parse
            input_ipd = int(serial_number[3:5])  # slice the serial number to get the IPD


def restart_vap():
    global calibration, threshold, raw_config, idle_time
    print("Restarting VAP...")
    calibration = 0.0
    threshold = 0.0
    idle_time = 0.0
    raw_config = []
    hid_device.close()


# ENTRY ###########################################

# Setup logging
logger = logging.getLogger("VAP")
logger.setLevel(logging.INFO)

# Setup glfw window and IMGUI.
if not glfw.init():
    print("Failed to start glfw.")
    exit()
glfw.window_hint(glfw.RESIZABLE, 0)
glfw_window = glfw.create_window(window_width, window_height,
                                 "Bigscreen VAP Tool", None, None)
if not glfw_window:
    print("Failed to create glfw window.")
    exit()
glfw.make_context_current(glfw_window)
glViewport(0, 0, window_width, window_height)
glEnable(GL_BLEND)
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
glOrtho(0.0, float(window_width), float(window_height), 0.0, 0.0, 1.0)
imgui.create_context()
imgui.get_io().display_size = [window_width, window_height]

# Load assets
asset_path = "assets/"
if os.path.exists("_internal/") is True:  # If we're in a Pyinstaller build, find assets in the _internal folder
    asset_path = "_internal/" + asset_path

imgui.get_io().fonts.add_font_from_file_ttf(asset_path + "fonts/arialuni.ttf", 48.0, None,
                                            imgui.get_io().fonts.get_glyph_ranges_chinese())
imgui.get_io().fonts.get_tex_data_as_rgba32()
imgui_renderer = GlfwRenderer(glfw_window)

# Load images.
connect_image = load_image(asset_path + "/images/connect.png")
no_pad_image = load_image(asset_path + "/images/no_pad.png")
pad_image = load_image(asset_path + "/images/pad.png")
good_image = load_image_rgba(asset_path + "/images/good.png")
bad_image = load_image_rgba(asset_path + "/images/bad.png")

# Main loop.
progress = 0.0
state = STATES.PROMPT
raw_config = []
calibration = 0.0
threshold = 0.0
last_error = ""
last_key_states = [glfw.RELEASE] * 350

should_run_pitch = [True]
should_run_proximity = [True]
should_run_roll = [False]

keyboard_focus_set = False  # Prevent repetitive focus calls in PROMPT state

while not glfw.window_should_close(glfw_window):
    glClearColor(1.0, 1.0, 1.0, 1.0)
    glClear(GL_COLOR_BUFFER_BIT)
    imgui.new_frame()
    imgui.set_next_window_size(imgui.get_io().display_size.x,
                               imgui.get_io().display_size.y)
    imgui.set_next_window_position(0, 0)
    imgui.begin("Bigscreen VAP Tool", True, imgui.WINDOW_NO_COLLAPSE | imgui.WINDOW_NO_RESIZE)

    # Process input
    enter_key_just_pressed = False

    enter_key_state = glfw.get_key(glfw_window, glfw.KEY_ENTER)
    kp_enter_key_state = glfw.get_key(glfw_window, glfw.KEY_KP_ENTER)

    if (enter_key_state == glfw.PRESS and last_key_states[glfw.KEY_ENTER] == glfw.RELEASE) or \
            (kp_enter_key_state == glfw.PRESS and last_key_states[glfw.KEY_KP_ENTER] == glfw.RELEASE):
        enter_key_just_pressed = True

    last_key_states[glfw.KEY_ENTER] = enter_key_state
    last_key_states[glfw.KEY_KP_ENTER] = kp_enter_key_state

    # UI to display.
    match state:
        case STATES.PROMPT:
            prox_checked, should_run_proximity[0] = imgui.checkbox("Run Proximity Calibration", should_run_proximity[0])
            va_checked, should_run_pitch[0] = imgui.checkbox("Run Pitch Alignment", should_run_pitch[0])
            roll_checked, should_run_roll[0] = imgui.checkbox("Run Roll Alignment", should_run_roll[0])
            imgui.dummy(0.0, 200.0)
            prompt_text = "Enter Serial #: "
            text_centered(prompt_text)
            imgui.push_item_width(400)
            centered(400)
            input_field_cleared_this_frame = False
            if serial_number != "" and serial_number[-1] == 'c':
                serial_number = ""
                input_field_cleared_this_frame = True
            else:
                _, serial_number = imgui.input_text('##sn', serial_number, 256, imgui.INPUT_TEXT_ENTER_RETURNS_TRUE)
            imgui.pop_item_width()
            if keyboard_focus_set is False:  # Set focus to text input on the first frame
                imgui.set_keyboard_focus_here(0)
                keyboard_focus_set = True
            if va_checked or prox_checked or roll_checked:  # If checkboxes were checked, bring focus back to text input
                imgui.set_keyboard_focus_here(0)
            if input_field_cleared_this_frame:
                keyboard_focus_set = False
            if enter_key_just_pressed:
                # parse_input(serial_number)
                state = STATES.CONNECT
                logger.handlers = []
                file_handler = logging.FileHandler(time.strftime(f"VAP_%Y_%m_%d") + f"_%s.log" % serial_number)
                formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
                file_handler.setFormatter(formatter)
                logger.addHandler(file_handler)
                logger.info("VAP process started. Serial number: %s" % serial_number)

        case STATES.CONNECT:
            try:
                hid_device.open(0x35bd, 257)
            except OSError:
                imgui.dummy(0.0, 20.0)
                text_centered("Connect linkbox to headset...")
                imgui.dummy(0.0, 20.0)
                text_centered("将串流盒连接到耳机")
                imgui.dummy(0.0, 20.0)
            else:
                hid_device.set_nonblocking(1)
                # raw_config = hmd_config_write_serial(hid_device, serial_number)
                if confirm_distortion() is False and should_check_red_distortion is True:
                    state = STATES.CALIBRATION_REJECTED
                    last_error = "Failed to detect optical calibration, run calibration upload again"
                elif confirm_tracking_calibration() is False:
                    state = STATES.CALIBRATION_REJECTED
                    last_error = "Failed to detect tracking calibration, run calibration upload again"
                elif should_run_proximity[0] is False:
                    state = STATES.ALIGN_INIT
                else:
                    state = STATES.PROXIMITY_RESET
        case STATES.CALIBRATION_REJECTED:
            err_str = last_error
            text_centered(err_str)
            logger.error(err_str)
            imgui.dummy(0.0, 60.0)
            text_centered("无法检测到校准，请重新运行校准上传")

            # If enter is pressed after rejection, restart VAP
            if enter_key_just_pressed:
                logger.info("Failed red distortion detection.")
                break
        case STATES.PROXIMITY_RESET:
            uut_start = time.time()
            raw_config = reset_proximity(hid_device)
            logger.info("Proximity values reset.")
            progress = 0.0
            # Reboot.
            hid_device.send_feature_report([0, 0x43])
            logger.info("Rebooting headset.")
            for p in range(10):
                imgui.dummy(0.0, 300.0)
                text_centered("Proximity calibration has begun, please wait...")
                imgui.dummy(0.0, 40.0)
                text_centered("接近校准已开始")
                text_centered("请稍等")
                imgui.dummy(0.0, 20.0)
                imgui.dummy(0.0, 40.0)
                centered(300)
                imgui.progress_bar(progress, (300, 20))
                progress += 0.1
                render_ui()
            hid_device.close()
            hid_device = hid.device()
            device_open = False
            while device_open is False:
                try:
                    hid_device.open(0x35bd, 257)
                except OSError:
                    time.sleep(0.1)
                else:
                    device_open = True
                    hid_device.set_nonblocking(1)
                    state = STATES.PROXIMITY_NO_PAD
        case STATES.PROXIMITY_NO_PAD:
            imgui.dummy(0.0, 40.0)
            text_centered("Place headset on side without the proximity pad, then click OK")
            imgui.dummy(0.0, 20.0)
            text_centered("将耳机放在没有接近垫的一侧，然后单击“好的”")
            imgui.dummy(0.0, 40.0)
            centered(200)
            if imgui.button("  OK 好的  ") or enter_key_just_pressed:
                state = STATES.PROXIMITY_NO_PAD_CALC
        case STATES.PROXIMITY_NO_PAD_CALC:
            # Flush out built up states.
            while True:
                data = hid_device.read(64)
                if not data:
                    break
            progress = 0
            min_value = float('inf')  # Tracks minimum read prox value for range calc
            max_value = float('-inf')  # Tracks maximum read prox value for range calc
            for p in range(10):
                imgui.dummy(0.0, 300.0)
                text_centered("Calculating...")
                text_centered("计算...")
                imgui.dummy(0.0, 40.0)
                read_proximity = proximity_read_proximity(hid_device)
                calibration += read_proximity
                print(read_proximity, flush=True)
                min_value = min(min_value, read_proximity)
                max_value = max(max_value, read_proximity)
                progress += 0.1
                centered(300)
                imgui.progress_bar(progress, (300, 20))
                render_ui()
            calibration /= 10.0
            print("Calibration: ", calibration, flush=True)
            logger.info("Calibration value: %f" % calibration)
            integer_calibration = int(calibration)

            if max_value - min_value > 150:
                state = STATES.PROXIMITY_REJECTED
                logger.error("Calibration value rejected. Range of read values exceeds 150.")
            else:
                raw_config = proximity_write_to_headset(hid_device, integer_calibration, raw_config)
                state = STATES.PROXIMITY_RESET_BEFORE_PAD
                logger.info("Calibration value written to headset: %f" % calibration)
        case STATES.PROXIMITY_RESET_BEFORE_PAD:
            progress = 0
            # Reboot.
            hid_device.send_feature_report([0, 0x43])
            logger.info("Rebooting headset.")
            for p in range(20):
                imgui.dummy(0.0, 300.0)
                text_centered("Calibrating, please wait...")
                imgui.dummy(0.0, 40.0)
                text_centered("正在校准，请稍候...")
                text_centered("请稍等")
                imgui.dummy(0.0, 60.0)
                centered(300)
                imgui.progress_bar(progress, (300, 20))
                progress += 0.05
                render_ui()
                time.sleep(0.1)
            hid_device.close()
            hid_device = hid.device()
            device_open = False
            while not device_open:
                try:
                    hid_device.open(0x35bd, 257)
                except OSError:
                    time.sleep(0.1)
                else:
                    device_open = True
                    hid_device.set_nonblocking(1)
                    state = STATES.PROXIMITY_NO_PAD
            state = STATES.PROXIMITY_PAD
        case STATES.PROXIMITY_PAD:
            imgui.dummy(0.0, 40.0)
            text_centered("Place headset on side with the proximity pad, then click OK")
            imgui.dummy(0.0, 20.0)
            text_centered("将耳机放在包括接近垫的一侧，然后单击“好的”")
            imgui.dummy(0.0, 40.0)
            centered(200)
            if imgui.button("  OK 好的  ") or enter_key_just_pressed:
                state = STATES.PROXIMITY_PAD_CALC
        case STATES.PROXIMITY_PAD_CALC:
            # Flush out built up states.
            while True:
                data = hid_device.read(64)
                if not data:
                    break
                time.sleep(0.1)
            calibrated_proximity = 0.0
            progress = 0
            min_value = float('inf')  # Tracks minimum read prox value for range calc
            max_value = float('-inf')  # Tracks maximum read prox value for range calc
            for p in range(10):
                imgui.dummy(0.0, 300.0)
                text_centered("Calculating...")
                text_centered("计算")
                imgui.dummy(0.0, 40.0)
                read_proximity = proximity_read_proximity(hid_device)
                threshold += read_proximity
                print(read_proximity, flush=True)
                min_value = min(min_value, read_proximity)
                max_value = max(max_value, read_proximity)
                progress += 0.1
                centered(300)
                imgui.progress_bar(progress, (300, 20))
                render_ui()
            threshold /= 10.0
            print(threshold, flush=True)
            logger.info("Threshold value: %f" % threshold)

            # Determine if the proximity threshold is acceptable
            if threshold < 500.0:
                state = STATES.PROXIMITY_REJECTED
                logger.error("Threshold value rejected. Value is below minimum acceptable threshold.")
            elif max_value - min_value > 150:
                logger.error("Threshold value rejected. Range of read values exceeds 150.")
            else:
                integer_threshold = int(threshold)
                raw_config = proximity_threshold_write_to_headset(hid_device, integer_threshold,
                                                                  raw_config)
                state = STATES.PROXIMITY_WRITE_THRESHOLD
                logger.info("Threshold value written to headset: %f" % threshold)
        case STATES.PROXIMITY_REJECTED:
            calibration_str = "Calibration: %i" % int(calibration)
            threshold_str = "Threshold: %i" % int(threshold)
            imgui.dummy(0.0, 20.0)
            text_centered(calibration_str)
            imgui.dummy(0.0, 20.0)
            text_centered(threshold_str)
            imgui.dummy(0.0, 60.0)
            text_centered("HMD did not pass the proximity calibration")
            imgui.dummy(0.0, 40.0)
            text_centered("没有通过接近校准")

            # If enter is pressed after rejection, restart VAP
            if enter_key_just_pressed:
                uut_stop = time.time()
                # formatted_string = mes_add(serial_number, "FAIL", uut_start, uut_stop, MES_SW_VER)
                #logger.info("Writing MES...")
                #logger.info(formatted_string)
                #logger.info("Failed proximity calibration.")
                break

        case STATES.PROXIMITY_WRITE_THRESHOLD:
            uut_stop = time.time()
            calibration_str = "Calibration: %i" % int(calibration)
            threshold_str = "Threshold: %i" % int(threshold)
            imgui.dummy(0.0, 10.0)
            text_centered(calibration_str)
            imgui.dummy(0.0, 10.0)
            text_centered(threshold_str)
            imgui.dummy(0.0, 10.0)
            text_centered("Successfully calibrated proximity")
            imgui.dummy(0.0, 10.0)
            centered(200)

            # Increment idle_time by the time since the last frame
            idle_time += imgui.get_io().delta_time

            if imgui.button("  OK 好的  ") or enter_key_just_pressed or idle_time > max_idle_time:
                idle_time = 0.0
                state = STATES.ALIGN_INIT
            imgui.dummy(0.0, 40.0)
            text_centered("成功校准接近度")
        case STATES.ALIGN_INIT:
            imgui.dummy(0.0, 300.0)
            text_centered("Starting alignment...")
            text_centered("开始对齐...")
            text_centered("请稍等")
            render_ui()
            logger.info("Starting alignment phase of VAP.")
            align_init(serial_number, input_ipd, canting_override)
            hid_device.send_feature_report([0, 0x70])  # Temporarily disable proximity sensor
            state = STATES.ALIGN_HORIZON
        case STATES.ALIGN_HORIZON:
            imgui.dummy(0.0, 300.0)
            text_centered("Finding HMD horizon...")
            text_centered("寻找 HMD 地平线...")
            text_centered("请稍等")
            render_ui()
            if align_find_hmd_horizon():
                if should_run_roll[0] is True:
                    state = STATES.ALIGN_ROLL
                elif should_run_pitch[0] is True:
                    state = STATES.ALIGN_PITCH
                else:
                    state = STATES.ALIGN_SHOW_RESULTS
                logger.debug("HMD Horizon: %s" % align_get_hmd_horizon())
        case STATES.ALIGN_PITCH:
            imgui.dummy(0.0, 300.0)
            text_centered("Finding best pitch...")
            text_centered("寻找最佳音高...")
            text_centered("请稍等")
            render_ui()
            pitch_left, pitch_right = align_minimize_pitch()
            logger.info("Left pitch value: %f" % pitch_left)
            logger.info("Right pitch value: %f" % pitch_right)
            logger.info("Left - Right: %f" % (pitch_left - pitch_right))
            print(align_get_left_config_string())
            logger.info(align_get_left_config_string())
            print(align_get_right_config_string())
            logger.info(align_get_right_config_string())
            state = STATES.ALIGN_SHOW_RESULTS
        case STATES.ALIGN_YAW:
            imgui.dummy(0.0, 300.0)
            text_centered("Finding best yaw...")
            text_centered("寻找最佳偏航...")
            text_centered("请稍等")
            render_ui()
            yaw_left = align_minimize_yaw(LEFT_EYE)
            yaw_right = align_minimize_yaw(RIGHT_EYE)
            print(yaw_left, flush=True)
            print(yaw_right, flush=True)
            state = STATES.ALIGN_ROLL
        case STATES.ALIGN_ROLL:
            imgui.dummy(0.0, 300.0)
            text_centered("Finding best roll...")
            text_centered("寻找最好的卷...")
            text_centered("请稍等")
            render_ui()
            roll_left = align_minimize_roll(LEFT_EYE)
            roll_right = align_minimize_roll(RIGHT_EYE)
            align_set_best_roll(-roll_left, -roll_right)
            if should_run_pitch[0] is True:
                state = STATES.ALIGN_PITCH
            else:
                state = STATES.ALIGN_SHOW_RESULTS
        case STATES.ALIGN_SHOW_RESULTS:
            l_pitch_str = "Pitch [left eye]: %f degrees" % pitch_left
            r_pitch_str = "Pitch [right eye]: %f degrees" % pitch_right
            pitch_diff_str = "L - R: %f degrees" % (pitch_left - pitch_right)
            l_roll_str = "Roll [left eye]: %f degrees" % roll_left
            r_roll_str = "Roll [right eye]: %f degrees" % roll_right
            imgui.dummy(0.0, 300.0)
            text_centered(l_pitch_str)
            text_centered(r_pitch_str)
            text_centered(pitch_diff_str)
            text_centered(l_roll_str)
            text_centered(r_roll_str)
            align_show_horizon()
            if enter_key_just_pressed:
                restart_vap()
                keyboard_focus_set = False
                state = STATES.PROMPT
                logger.info("VAP restarted after alignment.")

    imgui.end()
    imgui.end_frame()
    imgui_renderer.process_inputs()
    imgui.render()
    imgui_renderer.render(imgui.get_draw_data())

    # Image to display.
    match state:
        case STATES.CONNECT:
            render_image(connect_image.tex_id,
                         window_width / 2 - connect_image.width / 2,
                         window_height / 2 - connect_image.height / 2 + 150,
                         connect_image.width, connect_image.height)
        case STATES.PROXIMITY_NO_PAD:
            render_image(no_pad_image.tex_id,
                         window_width / 2 - no_pad_image.width / 2,
                         window_height / 2 - no_pad_image.height / 2 + 150,
                         no_pad_image.width, no_pad_image.height)
        case STATES.PROXIMITY_PAD:
            render_image(pad_image.tex_id,
                         window_width / 2 - pad_image.width / 2,
                         window_height / 2 - pad_image.height / 2 + 150,
                         pad_image.width, pad_image.height)
        case STATES.PROXIMITY_REJECTED:
            render_image(bad_image.tex_id,
                         window_width / 2 - bad_image.width / 2,
                         window_height / 2 - bad_image.height / 2 + 150,
                         bad_image.width, bad_image.height)
        case STATES.CALIBRATION_REJECTED:
            render_image(bad_image.tex_id,
                         window_width / 2 - bad_image.width / 2,
                         window_height / 2 - bad_image.height / 2 + 150,
                         bad_image.width, bad_image.height)
        case STATES.PROXIMITY_WRITE_THRESHOLD:
            render_image(good_image.tex_id,
                         window_width / 2 - good_image.width / 2,
                         window_height / 2 - good_image.height / 2 + 150,
                         good_image.width, good_image.height)

    glfw.poll_events()
    if glfw.get_key(glfw_window, glfw.KEY_ESCAPE):
        break
    glfw.swap_buffers(glfw_window)

# Cleanup.
align_destroy()
glfw.terminate()