import os
import subprocess

PREFIX = "arm-none-eabi-"
BUILDER = "DEV"

if os.getenv("PEDAL"):
  PROJECT = "pedal"
  STARTUP_FILE = "stm32fx/startup_stm32f205xx.s"
  LINKER_SCRIPT = "stm32fx/stm32fx_flash.ld"
  APP_START_ADDRESS = "0x8004000"
  MAIN = "pedal/main.c"
  PROJECT_FLAGS = [
    "-mcpu=cortex-m3",
    "-msoft-float",
    "-DSTM32F2",
    "-DSTM32F205xx",
    "-O2",
    "-DPEDAL",
  ]
  if os.getenv("PEDAL_USB"):
    PROJECT = "pedal_usb"
    PROJECT_FLAGS.append("-DPEDAL_USB")

elif os.getenv("PANDA_H7"):
  PROJECT = "panda_h7"
  STARTUP_FILE = "stm32h7/startup_stm32h7x5xx.s"
  LINKER_SCRIPT = "stm32h7/stm32h7x5_flash.ld"
  APP_START_ADDRESS = "0x8020000"
  MAIN = "main.c"
  PROJECT_FLAGS = [
    "-mcpu=cortex-m7",
    "-mhard-float",
    "-DSTM32H7",
    "-DSTM32H725xx",
    "-mfpu=fpv5-d16",
    "-fsingle-precision-constant",
    "-Os",
    "-g",
    "-DPANDA",
  ]

else:
  PROJECT = "panda"
  STARTUP_FILE = "stm32fx/startup_stm32f413xx.s"
  LINKER_SCRIPT = "stm32fx/stm32fx_flash.ld"
  APP_START_ADDRESS = "0x8004000"
  MAIN = "main.c"
  PROJECT_FLAGS = [
    "-mcpu=cortex-m4",
    "-mhard-float",
    "-DSTM32F4",
    "-DSTM32F413xx",
    "-mfpu=fpv4-sp-d16",
    "-fsingle-precision-constant",
    "-Os",
    "-g",
    "-DPANDA",
  ]


def get_version(builder, build_type):
  try:
    git = subprocess.check_output(["git", "rev-parse", "--short=8", "HEAD"], encoding='utf8').strip()
  except subprocess.CalledProcessError:
    git = "unknown"
  return f"{builder}-{git}-{build_type}"


def to_c_uint32(x):
  nums = []
  for _ in range(0x20):
    nums.append(x % (2**32))
    x //= (2**32)
  return "{" + 'U,'.join(map(str, nums)) + "U}"


def get_key_header(name):
  from Crypto.PublicKey import RSA

  public_fn = File(f'../certs/{name}.pub').srcnode().abspath
  rsa = RSA.importKey(open(public_fn).read())
  assert(rsa.size_in_bits() == 1024)

  rr = pow(2**1024, 2, rsa.n)
  n0inv = 2**32 - pow(rsa.n, -1, 2**32)

  r = [
    f"RSAPublicKey {name}_rsa_key = {{",
    f"  .len = 0x20,",
    f"  .n0inv = {n0inv}U,",
    f"  .n = {to_c_uint32(rsa.n)},",
    f"  .rr = {to_c_uint32(rr)},",
    f"  .exponent = {rsa.e},",
    f"}};",
  ]
  return r


def objcopy(source, target, env, for_signature):
    return '$OBJCOPY -O binary %s %s' % (source[0], target[0])


linkerscript_fn = File(LINKER_SCRIPT).srcnode().abspath

flags = [
  "-Wall",
  "-Wextra",
  "-Wstrict-prototypes",
  "-Werror",
  "-mlittle-endian",
  "-mthumb",
  "-nostdlib",
  "-fno-builtin",
  f"-T{linkerscript_fn}",
  "-std=gnu11",
] + PROJECT_FLAGS


if os.getenv("RELEASE"):
  BUILD_TYPE = "RELEASE"
  cert_fn = os.getenv("CERT")
  assert cert_fn is not None, 'No certificate file specified. Please set CERT env variable'
  assert os.path.exists(cert_fn), 'Certificate file not found. Please specify absolute path'
else:
  BUILD_TYPE = "DEBUG"
  cert_fn = File("../certs/debug").srcnode().abspath
  flags += ["-DALLOW_DEBUG"]

includes = [
  "stm32fx/inc",
  "stm32h7/inc",
  "..",
  ".",
]

panda_env = Environment(
  ENV=os.environ,
  CC=PREFIX + 'gcc',
  AS=PREFIX + 'gcc',
  OBJCOPY=PREFIX + 'objcopy',
  OBJDUMP=PREFIX + 'objdump',
  ASCOM="$AS $ASFLAGS -o $TARGET -c $SOURCES",
  CFLAGS=flags,
  ASFLAGS=flags,
  LINKFLAGS=flags,
  CPPPATH=includes,
  BUILDERS={
    'Objcopy': Builder(generator=objcopy, suffix='.bin', src_suffix='.elf')
  }
)

# Common autogenerated includes
version = f'const uint8_t gitversion[] = "{get_version(BUILDER, BUILD_TYPE)}";'
gitversion = panda_env.Textfile("obj/gitversion.h", [version, ""])
Ignore('bootstub.o', gitversion)
Requires('bootstub.o', gitversion)
Ignore('main.o', gitversion)
Requires('main.o', gitversion)

certs = [get_key_header(n) for n in ["debug", "release"]]
certheader = panda_env.Textfile("obj/cert.h", certs + [""])

startup = panda_env.Object(STARTUP_FILE)

# Bootstub
crypto = ["../crypto/rsa.c", "../crypto/sha.c"]
bootstub_elf = panda_env.Program(f"obj/bootstub.{PROJECT}.elf", [startup] + crypto + ["bootstub.c"])
bootstub_bin = panda_env.Objcopy(f"obj/bootstub.{PROJECT}.bin", bootstub_elf)

# Build main
main_elf = panda_env.Program(f"obj/{PROJECT}.elf", [startup, MAIN],
  LINKFLAGS=[f"-Wl,--section-start,.isr_vector={APP_START_ADDRESS}"] + flags)
main_bin = panda_env.Objcopy(f"obj/{PROJECT}.bin", main_elf)

# Sign main
sign_py = File("../crypto/sign.py").srcnode().abspath
panda_bin_signed = panda_env.Command(f"obj/{PROJECT}.bin.signed", main_bin, f"SETLEN=1 {sign_py} $SOURCE $TARGET {cert_fn}")