#!/usr/bin/python

import subprocess
import json
import pathlib
import os
import re

def select_options(message, options):
    user_input = ''
    input_message = message + '\n'
    for index, item in enumerate(options):
        input_message += f'{index+1}) {item}\n'
    input_message += 'Your choice: '
    while user_input not in map(str, range(1, len(options) + 1)):
        user_input = input(input_message)
    return options[int(user_input) - 1]

# save needed information in a dictionary
i3_dict = {}
        
cmd = str(subprocess.check_output("xrandr", shell=True))
outputs = []
freqs = []

# Get avaiable outputs, resolutions, and freqs
need_res = False
for row in cmd.split('\\n'):
       if " connected" in row:
           need_res = True
           outputs.append(row.split(" ")[0])
           continue
       elif re.search('[0-9]x[0-9]', row):
           freqs.append(row.split())
resolutions = [item[0] for item in freqs]

print('Setup monitor resolutions and freqs...')
monitor_set = {}
for index, output in enumerate(outputs):
    user_res = select_options(f'Pick {output} resolution: ', resolutions)
    user_freq = None
    for freq_lst in freqs:
        if freq_lst[0] == user_res:
            user_freq = select_options(f'Pick refresh rate: ', freq_lst[1:])
    monitor_set[output] = {"resolution": user_res, "freq": user_freq}

print("Setup monitor locations...")
print("Avaiable monitors")
for output in outputs:
    print(output)

if len(outputs) == 1:
    print("skipping input only one monitor...")
    first_key = list(monitor_set)[0]
    os.system("xrandr --output %s --mode %s --rate %s --primary"%(
                  outputs[0], 
                  monitor_set[first_key]['resolution'],
                  monitor_set[first_key]['freq'])
              )

elif len(outputs) == 2:
    left = input('which monitor is on the left? ')
    if left not in outputs:
        # crash
        0/0
    outputs.remove(left)
    right = outputs[0]
    left_primary = input('is left primary? (y or n) ')

    left_pri_str = "--primary" if left_primary == 'y' else ""
    right_pri_str = "--primary" if left_primary == 'n' else ""

    cmd_str = "xrandr --output %s --mode %s --pos 0x0 --rate %s %s --output %s --mode %s --pos 2561x0 --rate %s %s --right-of %s"%(
                  left, 
                  monitor_set[left]['resolution'],
                  monitor_set[left]['freq'],
                  left_pri_str, 
                  right, 
                  monitor_set[right]['resolution'],
                  monitor_set[right]['freq'],
                  right_pri_str, 
                  left
              )
    print(cmd_str)
    os.system(cmd_str)

else:
    #crash
    0/0


os.system("autorandr --save current --force")

if len(list(monitor_set)) == 1:
    i3_dict["disp_pri"] = list(monitor_set)[0]
    i3_dict["disp_sec"] = list(monitor_set)[0]
else:
    i3_dict["disp_pri"] = left if left_primary == 'y' else right
    i3_dict["disp_sec"] = left if left_primary == 'n' else right

print(i3_dict["disp_pri"], i3_dict["disp_sec"])

# write result to json
p = pathlib.PosixPath('~/.local/share/chezmoi/home/.gen/i3.json').expanduser()
p.parent.mkdir(parents=True, exist_ok=True)
with open(p, 'w') as f:
  json.dump(i3_dict, f, indent=4, sort_keys=True)