#
# coff_file.py -- read/write COFF files
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#
# WARNING: This module implements a VERY strict subset of the full COFF file
#          format. Just enough to read the output of GPASM, perform some
#          transforms, and write a new COFF file.
#
# BASIC FILE STRUCTURE
#
#   FILE HEADER (20 bytes)
#   "OPTIONAL" HEADER (18 bytes)  note: always present
#   SECTION RECORDS (40 bytes each)
#   SECTION DATA
#   RELOCATION RECORDS (12 bytes each)
#   LINENO RECORDS (16 bytes each)
#   SYMBOL TABLE (20 bytes each)
#   STRINGS DATA
#
#


import struct

# NOTE: we only support v2
_MICROCHIP_MAGIC_v2 = 0x1240
_OPT_HDR_SIZE_v2 = 18
_OPT_MAGIC_v2 = 0x5678
_ASSEMBLER_VERSION = 1
_SYMBOL_SIZE_v2 = 20
_RELOC_SIZE = 12
_LINENO_SIZE = 16

# Symbol storage classes
C_NULL = 0       # null
C_EXT = 2        # external symbol
C_STAT = 3       # static
C_LABEL = 6      # label
C_FILE = 103     # file name
C_EOF = 107      # end of file
C_SECTION = 109  # section

### for now, just support 16f688
PIC16F688 = 0x6688  # proc_code


def read_file(fname):
  fp = open(fname, 'rb')

  magic = _read16(fp)
  if magic != _MICROCHIP_MAGIC_v2:
    raise BadMagic()

  numsec = _read16(fp)
  tm = _read32(fp)
  symptr = _read32(fp)
  numsym = _read32(fp)
  opthdr = _read16(fp)
  fileflags = _read16(fp)

  # Optional header must be present, and v2 sized.
  if opthdr != _OPT_HDR_SIZE_v2:
    raise BadOptionalHeaderSize()

  magic = _read16(fp)
  if magic != _OPT_MAGIC_v2:
    raise BadMagic()
  vstamp = _read32(fp)
  if vstamp != _ASSEMBLER_VERSION:
    raise BadAssemblerVersion()

  proc_code = _read32(fp)
  if proc_code != PIC16F688:
    raise UnimplementedFeature()

  rom_width = _read32(fp)  # opcode bit width
  assert rom_width == 14  ### for 16F688

  ram_width = _read32(fp)  # RAM bit width
  assert ram_width == 8  ### for all PIC processors

  # This is the start of the section headers. It is actually a fixed position,
  # based on the size of the FILE and OPT headers above. Meh. Use .tell()
  sechdrs = fp.tell()

  #print 'SYMPTR:',symptr
  fp.seek(symptr + _SYMBOL_SIZE_v2 * numsym)
  strings = fp.read()

  symbols = [ ]
  if numsym > 0:
    fp.seek(symptr)

    i = 0
    while i < numsym:
      i += 1  # increment now, for clarity.

      name = _readname(fp, strings)
      value, secnum, symtyp, cls, numaux = struct.unpack('<LHLBB', fp.read(12))
      #print '#%d:'% i, name, value, secnum, symtyp, cls
      if numaux > 0:
        ### for now, only accept FILE and SECTION aux records
        if (cls != C_FILE and cls != C_SECTION) or numaux != 1:
          raise UnimplementedFeature()
        # Aux is the next symbol entry.
        i += 1
        if cls == C_FILE:
          pos = _read32(fp)
          aux_fname = _readstr(strings, pos)
          aux_lineno = _read32(fp)
          aux_flags = _read8(fp)
          _ = fp.read(11)  # skip. 11 unused bytes.
          #print '  FILE:', aux_fname, aux_lineno, aux_flags
          symbols.append((name, value, secnum, symtyp, cls,
                          (aux_fname, aux_lineno, aux_flags)))
          symbols.append(None)  # padding, so len(symbols) == numsym
        else:
          # Note: this data is duplicated in the section header.
          aux_len = _read32(fp)
          aux_nreloc = _read16(fp)
          aux_nlineno = _read16(fp)
          _ = fp.read(12)  # skip. 12 unused bytes.
          #print '  SECTION:', aux_len, aux_nreloc, aux_nlineno
          symbols.append((name, value, secnum, symtyp, cls,
                          (aux_len, aux_nreloc, aux_nlineno)))
          symbols.append(None)  # padding, so len(symbols) == numsym
      else:
        symbols.append((name, value, secnum, symtyp, cls, ()))

  #print len(symbols)

  # Read in the sections
  fp.seek(sechdrs)
  sections = [ ]
  for i in range(numsec):
    name = _readname(fp, strings)
    (addr, virt, size, dataptr, relocptr, linenoptr, nreloc, nlineno,
     secflags) = struct.unpack('<LLLLLLHHL', fp.read(32))
    assert addr == virt
    #print 'SECTION:', name, addr, virt, size, dataptr, relocptr, linenoptr, \
    #    nreloc, nlineno, secflags
    data = reloc = lineno = ''
    if dataptr and size:
      data = _readblock(fp, dataptr, size)
    if relocptr and nreloc:
      reloc = _readblock(fp, relocptr, nreloc * _RELOC_SIZE)
    if linenoptr and nlineno:
      lineno = _readblock(fp, linenoptr, nlineno * _LINENO_SIZE)

    # Note: record section SIZE and any DATA separately. The section may
    #       be "uninitialized", so it has a size, but no data.
    sections.append((name, addr, size, data, reloc, lineno, secflags))

  return tm, fileflags, symbols, sections


def _read8(fp):
  return ord(fp.read(1))
def _read16(fp):
  return struct.unpack('<H', fp.read(2))[0]
def _read32(fp):
  return struct.unpack('<L', fp.read(4))[0]
def _readstr(strings, pos):
  idx = strings.index('\0', pos)
  return strings[pos:idx]
def _readname(fp, strings):
  name = fp.read(8)
  if name.startswith('\0\0\0\0'):
    pos = struct.unpack('<L', name[4:])[0]
    return _readstr(strings, pos)
  return name.rstrip('\0')
def _readblock(fp, pos, amt):
  was = fp.tell()
  fp.seek(pos)
  data = fp.read(amt)
  fp.seek(was)
  return data


def write_file(fname, tm, fileflags, symbols, sections):
  fp = open(fname, 'wb')

  # Aggregate all the data/records from the sections
  secdata = ''
  relocdata = ''
  linenodata = ''
  for sec in sections:
    secdata += sec[3]
    relocdata += sec[4]
    linenodata += sec[5]
  #print len(sections), len(secdata), len(relocdata), len(linenodata)

  symptr = (20 + 18 + (40 * len(sections)) + len(secdata) +
            len(relocdata) + len(linenodata))
  #print 'SYMPTR:', symptr

  # Write the FILE header
  fp.write(struct.pack('<HHLLLHH',
                       _MICROCHIP_MAGIC_v2, len(sections), tm, symptr,
                       len(symbols), _OPT_HDR_SIZE_v2, fileflags))

  # Write the OPTIONAL header
  fp.write(struct.pack('<HLLLL', _OPT_MAGIC_v2, _ASSEMBLER_VERSION,
                       PIC16F688, 14, 8))

  # Current position, list of strings.
  stringptr = symptr + (20 * len(symbols))
  #print 'STRINGPTR: 0x%04x' % stringptr
  strings = [4, []]

  # Write the SECTION records, section data, reloc data, and lineno data
  dataptr = 20 + 18 + (40 * len(sections))
  relocptr = dataptr + len(secdata)
  linenoptr = relocptr + len(relocdata)
  for sec in sections:
    name = _formatname(sec[0], strings)
    this_data = len(sec[3]) and dataptr or 0
    this_reloc = len(sec[4]) and relocptr or 0
    this_lineno = len(sec[5]) and linenoptr or 0
    fp.write(struct.pack('<8sLLLLLLHHL', name, sec[1], sec[1], sec[2],
                         this_data, this_reloc, this_lineno,
                         len(sec[4])/_RELOC_SIZE, len(sec[5])/_LINENO_SIZE,
                         sec[6]))
    dataptr += len(sec[3])
    relocptr += len(sec[4])
    linenoptr += len(sec[5])
  fp.write(secdata)
  fp.write(relocdata)
  fp.write(linenodata)

  # Write the SYMBOL TABLE
  for sym in symbols:
    if not sym:
      continue
    name = _formatname(sym[0], strings)
    if sym[4] == C_FILE:
      fp.write(struct.pack('<8sLHLBB', name, sym[1], sym[2], sym[3], sym[4], 1))
      pos = strings[0]
      strings[0] += len(sym[5][0]) + 1  # + NUL terminator
      strings[1].append(sym[5][0] + '\0')
      fp.write(struct.pack('<LLB11s', pos, sym[5][1], sym[5][2], ''))
    elif sym[4] == C_SECTION:
      fp.write(struct.pack('<8sLHLBB', name, sym[1], sym[2], sym[3], sym[4], 1))
      fp.write(struct.pack('<LHH12s', sym[5][0], sym[5][1], sym[5][2], ''))
    else:
      fp.write(struct.pack('<8sLHLBB', name, sym[1], sym[2], sym[3], sym[4], 0))

  # Write the STRINGS
  stringdata = ''.join(strings[1])
  fp.write(struct.pack('<L', len(stringdata) + 4))
  fp.write(stringdata)


def _formatname(name, strings):
  if len(name) <= 8:
    return name  # struct.pack() will append NUL chars to read 8 chars.
  value = '\0\0\0\0' + struct.pack('<L', strings[0])
  strings[0] += len(name) + 1  # + NUL terminator
  strings[1].append(name + '\0')
  return value


class COFFFileException(Exception):
  pass
class UnimplementedFeature(COFFFileException):
  pass
class BadMagic(COFFFileException):
  pass
class BadOptionalHeaderSize(COFFFileException):
  pass
class BadAssemblerVersion(COFFFileException):
  pass