#!/usr/bin/python
# 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.

import copy
import json
import os
import sys

# First sector we can use.
START_SECTOR = 64

class ConfigNotFound(Exception):
  pass
class PartitionNotFound(Exception):
  pass
class InvalidLayout(Exception):
  pass


def LoadPartitionConfig(filename):
  """Loads a partition tables configuration file into a Python object.

  Args:
    filename: Filename to load into object
  Returns:
    Object containing disk layout configuration
  """

  if not os.path.exists(filename):
    raise ConfigNotFound('Partition config %s was not found!' % filename)
  with open(filename) as f:
    config = json.load(f)

  metadata = config['metadata']
  metadata['block_size'] = int(metadata['block_size'])

  for layout_name, layout in config['layouts'].items():
    for part in layout:
      part['blocks'] = int(part['blocks'])
      part['bytes'] = part['blocks'] * metadata['block_size']

      if 'fs_blocks' in part:
        part['fs_blocks'] = int(part['fs_blocks'])
        part['fs_bytes'] = part['fs_blocks'] * metadata['fs_block_size']

        if part['fs_bytes'] > part['bytes']:
          raise InvalidLayout('Filesystem may not be larger than partition')

  return config


def GetTableTotals(config, partitions):
  """Calculates total sizes/counts for a partition table.

  Args:
    config: Partition configuration file object
    partitions: List of partitions to process
  Returns:
    Dict containing totals data
  """

  ret = {
    'expand_count': 0,
    'expand_min': 0,
    'block_count': START_SECTOR * config['metadata']['block_size']
  }

  # Total up the size of all non-expanding partitions to get the minimum
  # required disk size.
  for partition in partitions:
    if 'features' in partition and 'expand' in partition['features']:
      ret['expand_count'] += 1
      ret['expand_min'] += partition['blocks']
    else:
      ret['block_count'] += partition['blocks']

  # At present, only one expanding partition is permitted.
  # Whilst it'd be possible to have two, we don't need this yet
  # and it complicates things, so it's been left out for now.
  if ret['expand_count'] > 1:
    raise InvalidLayout('1 expand partition allowed, %d requested'
                        % ret['expand_count'])

  ret['min_disk_size'] = ret['block_count'] + ret['expand_min']

  return ret


def GetPartitionTable(config, image_type):
  """Generates requested image_type layout from a layout configuration.

  This loads the base table and then overlays the requested layout over
  the base layout.

  Args:
    config: Partition configuration file object
    image_type: Type of image eg base/test/dev/factory_install
  Returns:
    Object representing a selected partition table
  """

  partitions = config['layouts']['base']

  if image_type != 'base':
    for partition_t in config['layouts'][image_type]:
      for partition in partitions:
        if partition['type'] == 'blank' or partition_t['type'] == 'blank':
          continue
        if partition_t['num'] == partition['num']:
          for k, v in partition_t.items():
            partition[k] = v

  return partitions


def GetScriptShell():
  """Loads and returns the skeleton script for our output script.

  Returns:
    A string containg the skeleton script
  """

  script_shell_path = os.path.join(os.path.dirname(__file__), 'cgpt_shell.sh')
  with open(script_shell_path, 'r') as f:
    script_shell = "".join(f.readlines())

  # Before we return, insert the path to this tool so somebody reading the
  # script later can tell where it was generated.
  script_shell = script_shell.replace('@SCRIPT_GENERATOR@', script_shell_path)

  return script_shell


def WriteLayoutFunction(sfile, func_name, image_type, config):
  """Writes a shell script function to write out a given partition table.

  Args:
    sfile: File handle we're writing to
    func_name: Function name to write out for specified layout
    image_type: Type of image eg base/test/dev/factory_install
    config: Partition configuration file object
  """

  partitions = GetPartitionTable(config, image_type)
  partition_totals = GetTableTotals(config, partitions)

  sfile.write('%s() {\ncreate_image $1 %d %s\n' % (
      func_name, partition_totals['min_disk_size'],
      config['metadata']['block_size']))

  sfile.write('CURR=%d\n' % START_SECTOR)
  sfile.write('$GPT create $1\n')

  # Pass 1: Set up the expanding partition size.
  for partition in partitions:
    partition['var'] = partition['blocks']
    if partition['type'] != 'blank':

      if partition['num'] == 1:
        if 'features' in partition and 'expand' in partition['features']:
          sfile.write('if [ -b $1 ]; then\n')
          sfile.write('STATEFUL_SIZE=$(( $(numsectors $1) - %d))\n' %
            partition_totals['block_count'])
          sfile.write('else\n')
          sfile.write('STATEFUL_SIZE=%s\n' % partition['blocks'])
          sfile.write('fi\n')
          partition['var'] = '$STATEFUL_SIZE'

  sfile.write('STATEFUL_SIZE=$((STATEFUL_SIZE-(STATEFUL_SIZE %% %d)))\n' %
    config['metadata']['fs_block_size'])

  # Pass 2: Write out all the cgpt add commands.
  for partition in partitions:
    if partition['type'] != 'blank':
      sfile.write('$GPT add -i %d -b $CURR -s %s -t %s -l %s $1 && ' % (
          partition['num'], str(partition['var']), partition['type'],
          partition['label']))

    # Increment the CURR counter ready for the next partition.
    sfile.write('CURR=$(( $CURR + %s ))\n' % partition['var'])

  # Set default priorities on kernel partitions
  sfile.write('$GPT add -i 2 -S 0 -T 15 -P 15 $1\n')
  sfile.write('$GPT add -i 4 -S 0 -T 15 -P 0 $1\n')
  sfile.write('$GPT add -i 6 -S 0 -T 15 -P 0 $1\n')

  sfile.write('$GPT boot -p -b $2 -i 12 $1\n')
  sfile.write('$GPT show $1\n')
  sfile.write('}\n')


def GetPartitionByNumber(partitions, num):
  """Given a partition table and number returns the partition object.

  Args:
    partitions: List of partitions to search in
    num: Number of partition to find
  Returns:
    An object for the selected partition
  """
  for partition in partitions:
    if partition['type'] == 'blank':
      continue
    if partition['num'] == int(num):
      return partition

  raise PartitionNotFound('Partition not found')


def WritePartitionScript(image_type, layout_filename, sfilename):
  """Writes a shell script with functions for the base and requested layouts.

  Args:
    image_type: Type of image eg base/test/dev/factory_install
    layout_filename: Path to partition configuration file
    sfilename: Filename to write the finished script to
  """

  config = LoadPartitionConfig(layout_filename)

  with open(sfilename, 'w') as f:
    script_shell = GetScriptShell()
    f.write(script_shell)

    WriteLayoutFunction(f, 'write_base_table', 'base', config)
    WriteLayoutFunction(f, 'write_partition_table', image_type, config)


def GetBlockSize(layout_filename):
  """Returns the partition table block size.

  Args:
    layout_filename: Path to partition configuration file
  Returns:
    Block size of all partitions in the layout
  """

  config = LoadPartitionConfig(layout_filename)
  return config['metadata']['block_size']


def GetFilesystemBlockSize(layout_filename):
  """Returns the filesystem block size.

  This is used for all partitions in the table that have filesystems.

  Args:
    layout_filename: Path to partition configuration file
  Returns:
    Block size of all filesystems in the layout
  """

  config = LoadPartitionConfig(layout_filename)
  return config['metadata']['fs_block_size']


def GetPartitionSize(image_type, layout_filename, num):
  """Returns the partition size of a given partition for a given layout type.

  Args:
    image_type: Type of image eg base/test/dev/factory_install
    layout_filename: Path to partition configuration file
    num: Number of the partition you want to read from
  Returns:
    Size of selected partition in bytes
  """

  config = LoadPartitionConfig(layout_filename)
  partitions = GetPartitionTable(config, image_type)
  partition = GetPartitionByNumber(partitions, num)

  return partition['bytes']


def GetFilesystemSize(image_type, layout_filename, num):
  """Returns the filesystem size of a given partition for a given layout type.

  If no filesystem size is specified, returns the partition size.

  Args:
    image_type: Type of image eg base/test/dev/factory_install
    layout_filename: Path to partition configuration file
    num: Number of the partition you want to read from
  Returns:
    Size of selected partition filesystem in bytes
  """

  config = LoadPartitionConfig(layout_filename)
  partitions = GetPartitionTable(config, image_type)
  partition = GetPartitionByNumber(partitions, num)

  if 'fs_bytes' in partition:
    return partition['fs_bytes']
  else:
    return partition['bytes']


def GetLabel(image_type, layout_filename, num):
  """Returns the label for a given partition.

  Args:
    image_type: Type of image eg base/test/dev/factory_install
    layout_filename: Path to partition configuration file
    num: Number of the partition you want to read from
  Returns:
    Label of selected partition, or 'UNTITLED' if none specified
  """

  config = LoadPartitionConfig(layout_filename)
  partitions = GetPartitionTable(config, image_type)
  partition = GetPartitionByNumber(partitions, num)

  if 'label' in partition:
    return partition['label']
  else:
    return 'UNTITLED'


def DoDebugOutput(image_type, layout_filename):
  """Prints out a human readable disk layout in on-disk order.

  This will round values larger than 1MB, it's exists to quickly
  visually verify a layout looks correct.

  Args:
    image_type: Type of image eg base/test/dev/factory_install
    layout_filename: Path to partition configuration file
  """
  config = LoadPartitionConfig(layout_filename)
  partitions = GetPartitionTable(config, image_type)

  for partition in partitions:
    if partition['bytes'] < 1024 * 1024:
      size = '%d bytes' % partition['bytes']
    else:
      size = '%d MB' % (partition['bytes'] / 1024 / 1024)
    if 'label' in partition:
      if 'fs_bytes' in partition:
        if partition['fs_bytes'] < 1024 * 1024:
          fs_size = '%d bytes' % partition['fs_bytes']
        else:
          fs_size = '%d MB' % (partition['fs_bytes'] / 1024 / 1024)
        print '%s - %s/%s' % (partition['label'], fs_size, size)
      else:
        print '%s - %s' % (partition['label'], size)
    else:
      print 'blank - %s' %  size


def main(argv):
  action_map = {
    'write': {
      'argc': 4,
      'usage': '<image_type> <partition_config_file> <script_file>',
      'func': WritePartitionScript,
    },
    'readblocksize': {
      'argc': 2,
      'usage': '<partition_config_file>',
      'func': GetBlockSize,
    },
    'readfsblocksize': {
      'argc': 2,
      'usage': '<partition_config_file>',
      'func': GetFilesystemBlockSize,
    },
    'readpartsize': {
      'argc': 4,
      'usage': '<image_type> <partition_config_file> <partition_num>',
      'func': GetPartitionSize,
    },
    'readfssize': {
      'argc': 4,
      'usage': '<image_type> <partition_config_file> <partition_num>',
      'func': GetFilesystemSize,
    },
    'readlabel': {
      'argc': 4,
      'usage': '<image_type> <partition_config_file> <partition_num>',
      'func': GetLabel,
    },
    'debug': {
      'argc': 3,
      'usage': '<image_type> <partition_config_file>',
      'func': DoDebugOutput,
    }
  }

  if len(sys.argv) < 2 or sys.argv[1] not in action_map:
    print 'Usage: %s <action>\n' % sys.argv[0]
    print 'Valid actions are:'
    for action in action_map:
      print '  %s %s' % (action, action_map[action]['usage'])
    sys.exit(1)
  else:
    action_name = sys.argv[1]
    action = action_map[action_name]
    if action['argc'] == len(sys.argv) - 1:
      print action['func'](*sys.argv[2:])
    else:
      sys.exit('Usage: %s %s %s' % (sys.argv[0], sys.argv[1], action['usage']))


if __name__ == '__main__':
  main(sys.argv)