"""Fix a BMP. Give a file-name to process, and a file-name for the output."""
# If you nerds truly need it, license is your choice of CC0, WTFPL, or Unlicense.


#***********************************************************


def process_file(input_filename, output_filename):
  with open(input_filename, "rb") as inputfile:
    header = inputfile.read(54)
    if header[0:2] != b"BM":
      raise Exception("not a BMP file")
    
    file_size = int.from_bytes(header[2:6], "little", signed=False)
    
    content_offset = int.from_bytes(header[8:12], "little", signed=False)
    if content_offset < 54:
      raise Exception("small content-offset. corrupted BMP file?")
    
    struct_size = int.from_bytes(header[14:18], "little", signed=False)
    if struct_size != 40:
      raise Exception("not a simple BMP file, you need a better Python script")
    
    width = int.from_bytes(header[18:22], "little", signed=True)
    height = int.from_bytes(header[22:26], "little", signed=True)
    
    planes = int.from_bytes(header[26:28], "little", signed=False)
    bits_per_pixel = int.from_bytes(header[28:30], "little", signed=False)
    compression = int.from_bytes(header[30:34], "little", signed=False)
    
    if planes != 1 or bits_per_pixel != 24 or compression != 0:
      raise Exception("unusual BMP, you need a better Python script")
    
    row_bytecount = width * (bits_per_pixel // 8)
    padded_row_bytecount = ((row_bytecount + 3) // 4) * 4
    
    calculated_file_size = 54 + padded_row_bytecount * height
    if file_size != calculated_file_size:
      raise Exception("wrong file size? corrupted file?")
    
    full_image_data = inputfile.read(file_size - 54)
    
    rows = []
    for y in range(height):
      start = y * padded_row_bytecount
      end = start + row_bytecount
      rows.append(full_image_data[start:end])
    rows.reverse()
    
    new_width = width + 2
    new_height = height
    
    new_row_bytecount = new_width * (bits_per_pixel // 8)
    new_padded_row_bytecount = ((new_row_bytecount + 3) // 4) * 4
    
    full_image_data = b"".join(rows)
    
    rows = []
    for y in range(new_height):
      start = y * new_row_bytecount
      end = start + new_row_bytecount
      rows.append(full_image_data[start:end])
    
    new_header = bytearray(header)
    new_file_size = 54 + new_padded_row_bytecount * new_height
    
    new_header[2:6] = new_file_size.to_bytes(4, "little", signed=False)
    new_header[18:22] = new_width.to_bytes(4, "little", signed=False)
    new_header[22:26] = new_height.to_bytes(4, "little", signed=False)
    
    with open(output_filename, "wb") as outputfile:
      outputfile.write(new_header)
      
      for row in reversed(rows):
        padding_count = new_padded_row_bytecount - len(row)
        if 0 < padding_count:
          row += b"\x00" * padding_count
        
        outputfile.write(row)


#***********************************************************


def main():
  import sys
  
  if len(sys.argv) < 2:
    print("Needs filename")
    return
  else:
    mangled_file_name = sys.argv[1]
    
    if 2 < len(sys.argv):
      new_file_name = sys.argv[2]
    else:
      new_file_name = mangled_file_name + ".fixed.bmp"
  
  process_file(mangled_file_name, new_file_name)
  
  print("end.")


#***********************************************************


if __name__ == "__main__":
  main()