#!/usr/bin/env python3

"""
Update autogenerated source files from yaml database.

Copyright (c) 2019, vit9696
"""

from __future__ import print_function

import update_products

import fnmatch
import json
import operator
import os
import unicodedata
import sys
import yaml

def remove_accents(input_str):
  nfkd_form = unicodedata.normalize('NFKD', input_str)
  return u"".join([c for c in nfkd_form if not unicodedata.combining(c)])

def load_db(dbpath):
  """
  Load yaml database and return in a list.
  """

  if not os.path.exists(dbpath):
    print("Cannot find %s directory, rerun from AppleModels directory!" % dbpath)
    sys.exit(1)

  db = []

  for root, dirs, files in os.walk(dbpath):
    for file in fnmatch.filter(files, '*.yaml'):
      path = os.path.join(root, file)
      with open(path, 'r') as fh:
        try:
            r = yaml.safe_load(fh)
            if r.get('SystemProductName', None) is None:
                print("WARN: Missing SystemProductName in %s, skipping!" % path)
                continue
            db.append(r)
        except yaml.YAMLError as e:
            print("Failed to parse file %s - %s" % (path, e))
            sys.exit(1)

  if len(db) == 0:
    print("Empty database!")
    sys.exit(1)

  # Sorting is required for fast lookup.
  return sorted(db, key=operator.itemgetter('SystemProductName'))

def gather_products(db, ptype='AppleModelCode', empty_valid=False, shared_valid=False, fatal=True):
  """
  Obtain all product codes from the database
  """
  products = []
  for info in db:
    pp = info.get(ptype, None)
    if pp is None:
      continue
    for p in pp:
      if p == '':
        if not empty_valid:
          print("ERROR: {} in contains empty {}, skipping!".format(info['SystemProductName'], ptype))
          if fatal: sys.exit(1)
        continue
      if p == '000' or p == '0000':
        print("WARN: {} in contains zero {}, skipping!".format(info['SystemProductName'], ptype))
        continue
      if p in products:
        if not shared_valid:
          print("ERROR: {} shares {} {} with other model!".format(info['SystemProductName'], ptype, p))
          if fatal: sys.exit(1)
        continue
      products.append(p)
  return products

def validate_products(db, dbpd):
  usedproducts  = gather_products(db)

  # FIXME: Empty is not valid, but we let it be for now.
  gather_products(db, 'AppleBoardCode', True, True, False)

  knownproducts = dbpd
  for product in usedproducts:
    if knownproducts.get(product, None) is None:
      print("ERROR: Model %s is used in DataBase but not present in Products!" % product)
      sys.exit(1)
    if knownproducts[product][update_products.KEY_STATUS] != update_products.STATUS_OK:
      print("ERROR: Model %s is used in DataBase but not valid in Products!" % product)
      sys.exit(1)

  to_add = {}

  for product in knownproducts:
    if knownproducts[product][update_products.KEY_STATUS] != update_products.STATUS_OK:
      continue

    name = knownproducts[product][update_products.KEY_NAME]
    if name.find('Mac') < 0 and name.find('Xserve') < 0:
      continue
    if name.find('M1') >= 0:
      continue
      
    if len(product) > 3 and product not in usedproducts:
      print("WARN: Model %s (%s) is known but is not used in DataBase!" % (product, name))

      if to_add.get(name, None) is None:
        to_add[name] = []
      to_add[name].append(product)
      continue

  if len(to_add) > 0:
    for sysname in to_add:
      for info in db:
        if sysname in info['Specifications']['SystemReportName']:
          print("New AppleModelCode for {}:".format(info['SystemProductName']))
          for model in to_add[sysname]:
            print("  - \"{}\"".format(model))

def export_db_macinfolib(db, path, year=0):
  """
  Export yaml database to MacInfoLib format.
  TODO: use jinja2?
  """

  with open(path, 'w') as fh:
    print('// DO NOT EDIT! This is an autogenerated file.', file=fh)
    print('#include "MacInfoInternal.h"', file=fh)
    print('CONST MAC_INFO_INTERNAL_ENTRY gMacInfoModels[] = {', file=fh)

    for info in db:
      if max(info['AppleModelYear']) < year:
        continue

      sb_model = info.get('AppleModelId')
      if sb_model:
        sb_model = '"{}"'.format(sb_model.lower())
      else:
        sb_model = 'NULL'

      print(' {\n'
        '  .SystemProductName = "%s",\n'
        '  .BoardProduct = "%s",\n'
        '  .BoardRevision = %s,\n'
        '  .SmcRevision = {%s},\n'
        '  .SmcBranch = {%s},\n'
        '  .SmcPlatform = {%s},\n'
        '  .BIOSVersion = "%s",\n'
        '  .BIOSReleaseDate = "%s",\n'
        '  .SystemVersion = "%s",\n'
        '  .SystemSKUNumber = "%s",\n'
        '  .SystemFamily = "%s",\n'
        '  .BoardVersion = "%s",\n'
        '  .BoardAssetTag = "%s",\n'
        '  .BoardLocationInChassis = "%s",\n'
        '  .SmcGeneration = 0x%X,\n'
        '  .BoardType = 0x%X,\n'
        '  .ChassisType = 0x%X,\n'
        '  .MemoryFormFactor = 0x%X,\n'
        '  .PlatformFeature = %s,\n'
        '  .ChassisAssetTag = "%s",\n'
        '  .FirmwareFeatures = 0x%XULL,\n'
        '  .FirmwareFeaturesMask = 0x%XULL,\n'
        '  .SecureBootModel = %s,\n'
        ' },' % (
          info['SystemProductName'],
          info['BoardProduct'][0] if isinstance(info['BoardProduct'], list) else info['BoardProduct'],
          '0x{:X}'.format(info['BoardRevision']) if 'BoardRevision' in info else 'MAC_INFO_BOARD_REVISION_MISSING',
          ', '.join(map(str, info.get('SmcRevision', [0x00]))),
          ', '.join(map(str, info.get('SmcBranch', [0x00]))),
          ', '.join(map(str, info.get('SmcPlatform', [0x00]))),
          info['BIOSVersion'],
          info['BIOSReleaseDate'],
          info['SystemVersion'],
          info['SystemSKUNumber'],
          info['SystemFamily'],
          info['BoardVersion'],
          info['BoardAssetTag'],
          info['BoardLocationInChassis'],
          info['SmcGeneration'],
          info['BoardType'],
          info['ChassisType'],
          info['MemoryFormFactor'],
          '0x{:X}'.format(info['PlatformFeature']) if 'PlatformFeature' in info else 'MAC_INFO_PLATFORM_FEATURE_MISSING',
          info['ChassisAssetTag'],
          info.get('ExtendedFirmwareFeatures', info.get('FirmwareFeatures', 0)),
          info.get('ExtendedFirmwareFeaturesMask', info.get('FirmwareFeaturesMask', 0)),
          sb_model
        ), file=fh)

    print('};', file=fh)

    print('CONST UINTN gMacInfoModelCount = ARRAY_SIZE (gMacInfoModels);', file=fh)
    print('CONST UINTN gMacInfoDefaultModel = 0;', file=fh)


def export_db_macserial(db, dbpd, path, year=0):
  """
  Export yaml database to macserial format.
  TODO: use jinja2?
  """

  with open(path, 'w') as fh:
    print('#ifndef GENSERIAL_MODELINFO_AUTOGEN_H', file=fh)
    print('#define GENSERIAL_MODELINFO_AUTOGEN_H\n', file=fh)
    print('// DO NOT EDIT! This is an autogenerated file.\n', file=fh)
    print('#include "macserial.h"\n', file=fh)

    print('typedef enum {', file=fh)

    for info in db:
      print('  {}, // {}'.format(
          info['SystemProductName'].replace(',', '_'),
          info['Specifications']['CPU'][0]
        ), file=fh)

    print('} AppleModel;\n', file=fh)
    print('#define APPLE_MODEL_MAX {}\n'.format(len(db)), file=fh)

    print('static PLATFORMDATA ApplePlatformData[] = {', file=fh)
    for info in db:
      print('  {{ "{}", "{}" }},'.format(
          info['SystemProductName'],
          info['SystemSerialNumber']
        ), file=fh)

    print('};\n', file=fh)

    print('#define APPLE_MODEL_CODE_MAX {}'.format(max(len(info['AppleModelCode']) for info in db)), file=fh)
    print('static const char *AppleModelCode[][APPLE_MODEL_CODE_MAX] = {', file=fh)

    for info in db:
      print('  /* {:14} */ {{"{}"}},'.format(
          info['SystemProductName'],
          '", "'.join(info['AppleModelCode'])
        ), file=fh)

    print('};\n', file=fh)

    print('#define APPLE_BOARD_CODE_MAX {}'.format(max(len(info['AppleBoardCode']) for info in db)), file=fh)
    print('static const char *AppleBoardCode[][APPLE_BOARD_CODE_MAX] = {', file=fh)

    for info in db:
      print('  /* {:14} */ {{"{}"}},'.format(
          info['SystemProductName'],
          '", "'.join(info['AppleBoardCode'])
        ), file=fh)

    print('};\n', file=fh)

    print('#define APPLE_MODEL_YEAR_MAX {}'.format(max(len(info['AppleModelYear']) for info in db)), file=fh)
    print('static uint32_t AppleModelYear[][APPLE_MODEL_YEAR_MAX] = {', file=fh)
    for info in db:
      print('  /* {:14} */ {{{}}},'.format(
          info['SystemProductName'],
          ', '.join(str(year) for year in info['AppleModelYear'])
        ), file=fh)

    print('};\n', file=fh)

    print('static uint32_t ApplePreferredModelYear[] = {', file=fh)
    for info in db:
      print('  /* {:14} */ {},'.format(
          info['SystemProductName'],
          info.get('MacserialModelYear', 0)
        ), file=fh)

    print('};\n', file=fh)

    print('static APPLE_MODEL_DESC AppleModelDesc[] = {', file=fh)

    models = sorted(dbpd.keys())
    models.sort(key=len)

    for model in models:
      if dbpd[model][update_products.KEY_STATUS] == update_products.STATUS_OK:
        print(' {{"{}", "{}"}},'.format(
            model,
            remove_accents(dbpd[model][update_products.KEY_NAME])
          ), file=fh)

    print('};\n', file=fh)

    print('#endif // GENSERIAL_MODELINFO_AUTOGEN_H', file=fh)

def export_mlb_boards(db, boards):
  l = {}
  for info in db:
    if len(info['SystemSerialNumber']) == 12:
      models = [info['BoardProduct']] if not isinstance(info['BoardProduct'], list) else info['BoardProduct']

      for model in models:
        if info['MaximumOSVersion'] is None:
          l[model] = 'latest'
        else:
          l[model] = info['MaximumOSVersion']

  with open(boards, 'w') as fh:
    json.dump(l, fh, indent=1)

if __name__ == '__main__':
  db   = load_db('DataBase')
  dbpd = update_products.load_products()
  # Run test phase to validate the library
  validate_products(db, dbpd)
  export_db_macinfolib(db, os.devnull)
  export_db_macserial(db, dbpd, os.devnull)
  # Export new models
  export_db_macinfolib(db, '../Library/OcMacInfoLib/AutoGenerated.c')
  export_db_macserial(db, dbpd, '../Utilities/macserial/modelinfo_autogen.h')
  # Export MLB models
  export_mlb_boards(db, '../Utilities/macrecovery/boards.json')