import sys

import numpy as np
import cv2
from mss.windows import MSS as mss
from PIL import Image


class Capturer:

	def __init__(self, w, h, resolution='1080p', fps=24.975):
		if resolution == '1080p':
			self.monitor = {
				'top': 1080 - h,
				'left': 1920 - w,
				'width': w,
				'height': h
			}

		elif resolution == '1440p':
			self.monitor = {
				'top': 1440 - h,
				'left': 2560 - w,
				'width': w,
				'height': h
			}

		else:
			raise ValueError('Unsupported monitor resolution')

		self.sct = mss()

		self.vid = cv2.VideoWriter(
			'output.avi',
			cv2.VideoWriter_fourcc(*'XVID'),
			fps=fps,
			frameSize=(
				self.monitor['width'],
				self.monitor['height']
			)
		)

	def mainloop(self):
		while True:
			sct_img = self.sct.grab(self.monitor)
			sct_img_size = (sct_img.width, sct_img.height)

			img = Image.frombytes(
				'RGB',
				sct_img_size,
				sct_img.rgb
			)

			frame = cv2.cvtColor(
				np.array(img),
				cv2.COLOR_BGR2RGB
			)

			self.vid.write(frame)


if __name__ == '__main__':
	if len(sys.argv) != 3:
		print(f'Usage: {sys.argv[0]} <WIDTH> <HEIGHT>')
		sys.exit()

	w, h = (int(x) for x in sys.argv[1:3])
	capturer = Capturer(w, h, fps=50.0)

	try:
		capturer.mainloop()
	except KeyboardInterrupt:
		print('Interrupted')