#!/usr/bin/env python

# This file is part of Window-Switch.
# Copyright (c) 2009-2013 Antoine Martin <antoine@nagafix.co.uk>
# Window-Switch is released under the terms of the GNU GPL v3

import os
import thread

from winswitch.consts import MAX_SCREENSHOT_WIDTH, MAX_SCREENSHOT_HEIGHT
from winswitch.globals import OSX, WIN32
from winswitch.util.commands_util import which
from winswitch.util.common import is_valid_file, is_valid_exe, load_binary_file, delete_if_exists, is_our_display
from winswitch.util.simple_logger import Logger
from winswitch.util.main_loop import callFromThread, callLater

logger = Logger("capture_util", log_colour=Logger.HIGHLIGHTED_BROWN)

DEBUG_DISPLAY_ACCESS = False
NO_CAPTURE_COMMAND_WARNING_SHOWN = False


def warn_cannot_capture(cmd):
	global NO_CAPTURE_COMMAND_WARNING_SHOWN
	actual_command = cmd
	assert not WIN32
	if cmd is not None:
		try:
			if type(cmd)==list:
				actual_command = which(cmd[0])
			elif type(cmd)==str:
				actual_command = which(cmd.split()[0])
			else:
				actual_command = cmd	#hope for the best!
		except:
			pass
	if not is_valid_exe(actual_command):
		err = "screen_capture_command '%s' is not defined or invalid!" % str(cmd)
		if not NO_CAPTURE_COMMAND_WARNING_SHOWN:
			logger.serror(err)
			NO_CAPTURE_COMMAND_WARNING_SHOWN = True
		else:
			logger.sdebug(err)
		return err
	NO_CAPTURE_COMMAND_WARNING_SHOWN = False
	return None


def screen_capture(screen_capture_command, display, env, filename, ok_callback, err_callback):
	""" Take a screenshot of the given display, resizing it to a thumbnail stored in filename.
		Then pass the data to ok_callback """
	delete_if_exists(filename)
	try:
		if WIN32:
			c = PILCaptureDisplay(display, filename, ok_callback, err_callback)
		elif not OSX and is_our_display(display):
			c = PyGTKCaptureDisplay(filename, ok_callback, err_callback)
		else:
			c = ExecCaptureDisplay(screen_capture_command, display, env, filename, ok_callback, err_callback)
		c.capture()
	except Exception, e:
		logger.serr(None, e, display, filename)
		err_callback("exception: %s" % e)
		return




class ScreenCapture:
	""" Utility superclass for all capture implementations """

	def __init__(self, display, env, filename, ok_callback, err_callback, threaded):
		Logger(self, log_colour=Logger.HIGHLIGHTED_BROWN)
		self.display = display
		self.env = env
		self.filename = filename
		self.ok_callback = ok_callback
		self.err_callback = err_callback
		self.threaded = threaded

	def ok(self, data):
		callFromThread(self.ok_callback, data)

	def err(self, msg):
		callFromThread(self.err_callback, "%s failed for display %s: %s" % (self.__class__, self.display, msg))

	def capture(self):
		def _capture():
			try:
				delete_if_exists(self.filename)
				self.do_capture()
			except Exception, e:
				self.err("%s" % e)
		if self.threaded:
			thread.start_new_thread(_capture, ())
		else:
			_capture()

	def do_capture(self):
		raise Exception("override me!")

	def get_thumbnail_dimensions(self, w, h):
		""" keeps aspect ratio and fits picture within max dimensions """
		max_h = MAX_SCREENSHOT_HEIGHT
		max_w = MAX_SCREENSHOT_WIDTH
		if w<max_w and h<max_h:
			self.sdebug("no need to resize", w, h)
			return (w,h)
		if (w/h)>=(max_w/max_h):
			return (max_w, int(h*max_w/w))	#resize to max_w
		else:
			return	(int(w*max_h/h), h)		#resize to max_h

	def finish(self):
		if not is_valid_file(self.filename):
			self.err("capture file '%s' does not exist!" % self.filename)
			return
		data = load_binary_file(self.filename)
		if not data:
			self.err("capture file '%s' is empty!" % self.filename)
			return
		self.ok(data)

	def resize_file(self):
		#resize it
		img_data = load_binary_file(self.filename)
		import StringIO
		import PIL.Image
		image = PIL.Image.open(StringIO.StringIO(img_data))
		self.save_resized(image)

	def save_resized(self, image):
		""" save a PIL.Image after resizing it to thumbnail size """
		import PIL.Image
		(w,h) = image.size
		w,h = self.get_thumbnail_dimensions(w, h)
		image.thumbnail((w,h), PIL.Image.ANTIALIAS)
		image.save(self.filename, "PNG")


class ExecCaptureDisplay(ScreenCapture):
	""" Captures the screen using screen_capture_command """

	def __init__(self, screen_capture_command, display, env, filename, ok_callback, err_callback):
		ScreenCapture.__init__(self, display, env, filename, ok_callback, err_callback, False)
		self.screen_capture_command = screen_capture_command

	def capture(self):
		err = warn_cannot_capture(self.screen_capture_command)
		if err:
			self.err_callback(err)
			return
		ScreenCapture.capture(self)

	def do_capture(self):
		cmd = ["sh", "-c", "%s%s" % (self.screen_capture_command, self.filename)]
		self.exec_capture(cmd)

	def exec_capture(self, cmd, wait_for_file=0):
		self.sdebug(None, cmd, wait_for_file)
		def log_fail_info(message):
			if OSX:
				self.sdebug("%s, maybe the screensaver is activated?" % message, cmd, wait_for_file)
			else:
				self.serror("%s, capture failed" % message, cmd, wait_for_file)
		def check_for_file(wait=0):
			if is_valid_file(self.filename):
				filesize = os.stat(self.filename).st_size
				if filesize==0:
					log_fail_info("file '%s' is empty" % self.filename)
				else:
					self.resize_file()
					self.finish()
			elif wait>0:
				callLater(1, check_for_file, wait-1)
			else:
				log_fail_info("file %s is not valid, maybe the screensaver is activated?" % self.filename)
				self.err("screenshot file still missing after %s seconds" % wait_for_file)
		def start_checking(*args):
			check_for_file(wait_for_file)
		from winswitch.util.process_util import twisted_exec
		twisted_exec(cmd, start_checking, err_callback=self.err, env=self.env, log_lifecycle=False)


class PILCaptureDisplay(ScreenCapture):
	""" Captures the screen using PIL.ImageGrab """

	def __init__(self, display, filename, ok_callback, err_callback):
		ScreenCapture.__init__(self, display, None, filename, ok_callback, err_callback, True)

	def	do_capture(self):
		#force py2exe to include these:
		import PIL.Image
		import PIL.GifImagePlugin
		import PIL.JpegImagePlugin
		import PIL.PngImagePlugin
		import PIL.ImageGrab
		assert PIL.Image and PIL.GifImagePlugin and PIL.JpegImagePlugin and PIL.PngImagePlugin and PIL.ImageGrab		#avoid pydev warnings
		screenshot = PIL.ImageGrab.grab()
		self.save_resized(screenshot)
		self.finish()


class PyGTKCaptureDisplay(ScreenCapture):

	def __init__(self, filename, ok_callback, err_callback):
		ScreenCapture.__init__(self, None, None, filename, ok_callback, err_callback, False)

	def do_capture(self):
		import gtk.gdk
		screen = gtk.gdk.screen_get_default()
		window = screen.get_root_window()
		w,h = window.get_size()
		pb = gtk.gdk.Pixbuf(gtk.gdk.COLORSPACE_RGB, False, 8, w, h)
		pb = pb.get_from_drawable(window, window.get_colormap(), 0, 0, 0, 0, w, h)
		self.sdebug("window=%s, w=%s h=%s pixbuf=%s" % (window, w, h, pb))
		tw, th = self.get_thumbnail_dimensions(w, h)
		scaled = pb.scale_simple(tw, th, gtk.gdk.INTERP_BILINEAR)
		if not scaled:
			self.err_callback("error scaling gtk screenshot")
			return
		scaled.save(self.filename, "png")
		thread.start_new_thread(self.finish, ())
