#!/bin/bash

# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

COMMON_SH=/usr/lib/crosutils/common.sh
. "${COMMON_SH}" || exit 1

# Flags
DEFINE_string board "${DEFAULT_BOARD}" \
	"The board to run debugger on."
DEFINE_string image "" \
	"Full pathname of the EC firmware image to flash."
DEFINE_string offset "0" \
	"Offset where to program the image from."
DEFINE_integer port 9999 \
	"Port to communicate to servo on."
DEFINE_boolean ro "${FLAGS_FALSE}" \
	"Write only the read-only partition"
DEFINE_boolean unprotect "${FLAGS_FALSE}" \
	"Clear the protect flag."

# Parse command line
FLAGS_HELP="usage: $0 [flags]"
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"
check_flags_only_and_allow_null_arg "$@" && set --

set -e

SERVO_TYPE=servo

# reset the EC
toad_hard_reset() {
	info "you probably need to hard-reset your EC with Refresh+Power"
}

servo_hard_reset() {
	dut_control cold_reset:on
	dut_control cold_reset:off
}

ec_reset() {
	eval ${SERVO_TYPE}_hard_reset
}

# force the EC to boot in serial monitor mode
toad_boot0() {
	dut_control boot_mode:yes
}

servo_boot0() {
	dut_control spi1_vref:pp3300
}

ec_enable_boot0() {
	eval ${SERVO_TYPE}_boot0
}

# Put back the servo and the system in a clean state at exit
cleanup() {
	if [ -n "${save}" ]; then
		info "Restoring servo settings..."
		servo_restore "$save"
	fi

	ec_reset
}
trap cleanup EXIT

BOARD=${FLAGS_board}
BOARD_ROOT=/build/${BOARD}

# Possible default EC images
if [ "${FLAGS_ro}" = ${FLAGS_TRUE} ] ; then
	EC_FILE=ec.RO.flat
else
	EC_FILE=ec.bin
fi
EMERGE_BUILD=${BOARD_ROOT}/firmware/${EC_FILE}
LOCAL_BUILD=${SRC_ROOT}/platform/ec/build/${BOARD}/${EC_FILE}

# Find the EC image to use
function ec_image() {
	# No image specified on the command line, try default ones
	if [[ -n "${FLAGS_image}" ]] ; then
		if [ -f "${FLAGS_image}" ]; then
			echo "${FLAGS_image}"
			return
		fi
		die "Invalid image path : ${FLAGS_image}"
	else
		if [ -f "${LOCAL_BUILD}" ]; then
			echo "${LOCAL_BUILD}"
			return
		fi
		if [ -f "${EMERGE_BUILD}" ]; then
			echo "${EMERGE_BUILD}"
			return
		fi
	fi
	die "no EC image found : build one or specify one."
}

DUT_CONTROL_CMD="dut-control --port=${FLAGS_port}"

# Find the EC UART on the servo v2
function ec_uart() {
	SERVOD_FAIL="Cannot communicate with servo. is servod running ?"
	($DUT_CONTROL_CMD ec_uart_pty || \
            die "${SERVOD_FAIL}") | cut -d: -f2
}

# Servo variables management

servo_VARS="ec_uart_en ec_uart_parity ec_uart_baudrate \
jtag_buf_on_flex_en jtag_buf_en spi1_vref"
toad_VARS="ec_uart_parity ec_uart_baudrate boot_mode"

function dut_control() {
	$DUT_CONTROL_CMD "$1" >/dev/null
}

function servo_save() {
	SERVO_VARS_NAME=${SERVO_TYPE}_VARS
	$DUT_CONTROL_CMD ${!SERVO_VARS_NAME}
}

function servo_restore() {
	echo "$1" | while read line
	do
		dut_control "$line"
	done
}

function free_pty() {
	pids=$(lsof -F p 2>/dev/null -- $1 | cut -d'p' -f2)
	if [ "${pids}" != "" ]; then
		kill -9 ${pids}
		info "You'll need to re-launch console on $1"
	fi
}

# Board specific flashing scripts

function flash_daisy() {
	TOOL_PATH="${SCRIPT_LOCATION}/../build/${BOARD}/util:$PATH"
	STM32MON=$(PATH="${TOOL_PATH}" which stm32mon)
	if [ ! -x "$STM32MON" ]; then
		die "no stm32mon util found."
	fi

	if [ "${FLAGS_unprotect}" = ${FLAGS_TRUE} ] ; then
		# Unprotect exists, but isn't needed because erasing pstate is
		# implicit in writing the entire image
		die "--unprotect not supported for this board."
	fi

	info "Using serial flasher : ${STM32MON}"
	free_pty ${EC_UART}

	if [ "${SERVO_TYPE}" = "servo" ] ; then
		dut_control ec_uart_en:on
	fi
	dut_control ec_uart_parity:even
	dut_control ec_uart_baudrate:115200
	# Force the EC to boot in serial monitor mode
	ec_enable_boot0
	# Reset the EC
	ec_reset
	# Unprotect flash, erase, and write
	${STM32MON} -d ${EC_UART} -u -e -w ${IMG}
}

function flash_link() {
	OCD_CFG="servo_v2_slower.cfg"
	OCD_PATH="${SRC_ROOT}/platform/ec/chip/lm4/openocd"
	OCD_CMDS="init; flash_lm4 ${IMG} ${FLAGS_offset};"
	if [ "${FLAGS_unprotect}" = ${FLAGS_TRUE} ] ; then
		info "Clearing write protect flag."
		OCD_CMDS="${OCD_CMDS} unprotect_link;"
	fi
	OCD_CMDS="${OCD_CMDS} shutdown;"

	dut_control jtag_buf_on_flex_en:on
	dut_control jtag_buf_en:on

	sudo openocd -s "${OCD_PATH}" -f "${OCD_CFG}" -c "${OCD_CMDS}" || \
	die "Failed to program ${IMG}"
}

function flash_slippy() {
	OCD_CFG="servo_v2_slower.cfg"
	OCD_PATH="${SRC_ROOT}/platform/ec/chip/lm4/openocd"
	OCD_CMDS="init; flash_lm4 ${IMG} ${FLAGS_offset};"
	if [ "${FLAGS_unprotect}" = ${FLAGS_TRUE} ] ; then
		# Unprotect exists, but isn't needed because erasing pstate is
		# implicit in writing the entire image
		die "--unprotect not supported for this board."
	fi
	OCD_CMDS="${OCD_CMDS} shutdown;"

	dut_control jtag_buf_on_flex_en:on
	dut_control jtag_buf_en:on

	sudo openocd -s "${OCD_PATH}" -f "${OCD_CFG}" -c "${OCD_CMDS}" || \
	die "Failed to program ${IMG}"
}

IMG="$(ec_image)"
info "Using EC image : ${IMG}"

EC_UART="$(ec_uart)"
info "EC UART pty : ${EC_UART}"

if dut_control uart_mux 2>/dev/null ; then
	SERVO_TYPE=toad
	info "Using a TOAD cable"
fi

save="$(servo_save)"

case "${BOARD}" in
	puppy | daisy | snow | spring | pit ) flash_daisy ;;
	link ) flash_link ;;
	slippy | falco | peppy ) flash_slippy ;;
	*) die "board ${BOARD} not supported" ;;
esac

info "Flashing done."