#!/usr/bin/env python

import os
import sys
import urllib2
import Image
import yaml

BASE_URL = "http://localhost:4567/"
YAML_FILE = 'info.yml'

class Variation(object):

  def __init__(self, name, labels):
    self.name = name
    self.labels = list(labels)

  def __unicode__(self):
    return self.filename
  __repr__ = __unicode__
  
  @property
  def filename(self):
    return '_'.join(self.labels)

  @property    
  def url(self):
    return "%s?%s" % (BASE_URL, "%20".join(self.labels))


def product(*args, **kwds):
    "Given a list of iterables, calculate the possible combinations."
    # Based on itertools.product from Python 2.6
    pools = map(tuple, args) * kwds.get('repeat', 1)
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)


def load_variations(yaml_path):
  "Loads a list of Variations from the specified yaml file."
  data = open(yaml_path).read()
  parsed = yaml.load(data)
  variation_name = parsed['name'].lower().replace(r' ', '_')
  variation_categories = []
  variations_dict = parsed.get('variations', {})
  for category_name in sorted(variations_dict.keys()):
    variation_categories.append(variations_dict[category_name])
  variations = []
  for combo in product(*[category.keys() for category in variation_categories]):
    variations.append(Variation(variation_name, combo))
  return variations
  
  
if __name__ == '__main__':
  if not os.path.exists(YAML_FILE):
    print "\n** Can't find %s.  This script must be run from the Theme's root directory.\n" % YAML_FILE
    sys.exit(0)
  
  try:
    urllib2.urlopen(BASE_URL)
  except urllib2.URLError, e:
    print "\n** Couldn't connect to %s -- are you sure ThemeKit is running?\n" % BASE_URL
    sys.exit(0)
  
  # Determine the directory in which to save the screenshots.
  # If it already exists, clear it's files and remove the directory.
  dest = os.path.join(os.getcwd(), 'screenshots')
  if os.path.isdir(dest):
    print "** Removing existing screenshots..."
    for path in os.listdir(dest):
      os.remove(os.path.join(dest, path))
    os.rmdir(dest)
  os.mkdir(dest)

  from subprocess import call, PIPE
  variations = load_variations(YAML_FILE)
  print '** Grabbing screenshots for %s variations...' % len(variations)
  count = 1
  for v in variations:
    print "** %s) %s \t(%s)" % (count, " ".join(v.labels), v.url)
    count += 1
    call(['webkit2png', '--filename=%s' % v.filename, '--dir=%s' % dest, '--width=1024', '--height=768', '-F', v.url], stdout=open(os.devnull, 'w'))
    
  for img_path in os.listdir(dest):
    full_path = os.path.join(dest, img_path)
    img = Image.open(full_path)
    img = img.crop((0,0,1024,768))
    img.thumbnail((300,300), Image.ANTIALIAS)
    img.save(full_path.replace("-full", ""), "JPEG", quality=95)
    os.remove(full_path)
    
  print '** Screenshots have been saved in %s' % dest