# A simple script to automatically take app screenshots
# Hopefully some day it'll grow into an actual testing framework

from uuid import UUID
from datetime import datetime, timezone
from time import sleep
import os
import png
import json

from pebble_tool.util.logs import PebbleLogPrinter
from pebble_tool.sdk.emulator import ManagedEmulatorTransport
from pebble_tool.commands.emucontrol import send_data_to_qemu
from libpebble2.communication import PebbleConnection
from libpebble2.communication.transports.qemu.protocol import QemuButton, QemuTimeFormat
from libpebble2.services.install import AppInstaller
from libpebble2.services.appmessage import AppMessageService, Uint32
from libpebble2.services.screenshot import Screenshot
from libpebble2.protocol.system import TimeMessage, SetUTC

package_json = None
with open('package.json', 'r') as file:
  package_json = json.load(file)

app_uuid = UUID(package_json['pebble']['uuid'])
appname = package_json['name']
target_platforms = package_json['pebble']['targetPlatforms']
msg_keys = {
  key: 10000 + i
  for i, key in enumerate(package_json['pebble']['messageKeys'])
}

def emulator(platform):
  print(f'Launching `{platform}` emulator')
  pebble = PebbleConnection(ManagedEmulatorTransport(platform))
  pebble.connect()
  pebble.run_async()
  print(f'Installing `{appname}.pbw`')
  AppInstaller(pebble, f'./build/{appname}.pbw').install()
  sleep(1)
  return pebble

def for_all_platforms(fn, logs=True):
  for platform in target_platforms:
    pebble = emulator(platform)
    if logs:
      logger = PebbleLogPrinter(pebble)
      fn(pebble)
      logger.stop()
    else:
      fn(pebble)


def is_color(pebble):
  return pebble.watch_platform in ['basalt', 'chalk', 'emery']

def send_app_message(pebble, msg):
  ams = AppMessageService(pebble)
  tid = ams.send_message(app_uuid, msg)

  # unfortunately the following doesn't work
  # wait_for_event assumes only one value is returned
  # but ack returns 2 :(

  # tid2, uuid = ams.wait_for_event('ack')
  # assert uuid == app_uuid
  # assert tid2 == tid

def set_time(pebble, hour, minute, year=2013, month=1, day=23, tzinfo=timezone.utc):
  dt = int(datetime(year, month, day, hour, minute, tzinfo=tzinfo).timestamp())
  pebble.send_packet(TimeMessage(message=SetUTC(unix_time=dt, utc_offset=0, tz_name="UTC")))

def press_buttons(pebble, buttons, duration=0.1):
  send_data_to_qemu(pebble.transport, QemuButton(state=buttons))
  sleep(duration)
  send_data_to_qemu(pebble.transport, QemuButton(state=0))
  sleep(0.1)

def press_back(pebble, duration=0.1):
  press_buttons(pebble, QemuButton.Button.Back, duration)
def press_up(pebble, duration=0.1):
  press_buttons(pebble, QemuButton.Button.Up, duration)
def press_select(pebble, duration=0.1):
  press_buttons(pebble, QemuButton.Button.Select, duration)
def press_down(pebble, duration=0.1):
  press_buttons(pebble, QemuButton.Button.Down, duration)

def set_24_hour_mode(pebble, is_24_hour=True):
  send_data_to_qemu(pebble.transport, QemuTimeFormat(is_24_hour=is_24_hour))

def screenshot(pebble, name):
  filename = f'{pebble.watch_platform}/{name}.png'
  path = './screenshots/' + filename
  print(f'Taking screenshot `{filename}`')
  os.makedirs(os.path.dirname(path), exist_ok=True)
  image = Screenshot(pebble).grab_image()
  png.from_array(image, 'RGB').save(path)

# ============================================================
# Code above this line is not specific to the watchface and
# should probably be eventually split off into its own library
# ============================================================

def set_colors(pebble, background, digits, frame):
  def wrap_color(color):
    if isinstance(color, int):
      color = hex(color)
    color = color.removeprefix('0x').removeprefix('#')
    if len(color) == 3:
      r, g, b = color
      color = r + r + g + g + b + b
    return Uint32(int(color, 16))
  send_app_message(pebble, {
    msg_keys["BACKGROUND_COLOR"]: wrap_color(background),
    msg_keys["DIGIT_COLOR"]: wrap_color(digits),
    msg_keys["FRAME_COLOR"]: wrap_color(frame),
  })

def run(pebble):
  set_24_hour_mode(pebble, True)
  set_time(pebble, 12, 34)
  set_colors(pebble, 'FFF', '000', '000')
  sleep(1.5)
  screenshot(pebble, '01_white-black-black')

  set_24_hour_mode(pebble, False)
  set_time(pebble, 16, 59)
  set_colors(pebble, '000', 'FFF', 'AAA')
  sleep(1.5)
  screenshot(pebble, '02_black-white-gray')

  if is_color(pebble):
    set_24_hour_mode(pebble, True)
    set_time(pebble, 4, 20)
    set_colors(pebble, 'FF5', '00F', '00A')
    sleep(1.5)
    screenshot(pebble, '03_yellow-blue-blue')


if __name__ == '__main__':
  for_all_platforms(run)