# Copyright 2021 The Pigweed Authors
#
# Licensed 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
#
#     https://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.
"""Converts ST generated .icf linker files into basic .ld linker files"""

from typing import Dict, Optional, Tuple

import pathlib


def parse_icf(icf_file: str) -> Tuple[Dict, Dict]:
    """Parse ICF linker file.

    ST only provides .icf linker files for many products, so there is a need
    to generate basic GCC compatible .ld files for all products.
    This parses the basic features from the .icf format well enough to work
    for the ST's .icf files that exist in `cmsis_device`

    Args:
        icf_file: .icf linker file read into a string

    Returns:
        (regions, blocks) where
            `regions` is a map from region_name -> (start_hex, end_hex)
            `blocks` is a map from block_name -> {feature_1: val_1,...}

    Raises:
        IndexError if .icf is malformed (at least compared to how ST makes them)
    """
    symbols = {}
    regions = {}  # region: (start_addr, end_addr)
    blocks = {}
    for line in icf_file.split('\n'):
        line = line.strip()
        if line == '' or line.startswith('/*') or line.startswith('//'):
            continue
        tokens = line.split()
        if len(tokens) < 2:
            continue
        if tokens[0] == 'define':
            if tokens[1] == 'symbol':
                symbols[tokens[2]] = tokens[4].strip(';')
            elif tokens[1] == 'region':
                regions[tokens[2].split('_')[0]] = (tokens[5],
                                                    tokens[7].strip('];'))
            elif tokens[1] == 'block':
                blocks[tokens[2]] = {
                    tokens[4]: tokens[6].strip(','),
                    tokens[7]: tokens[9]
                }
    parsed_regions = {
        region: (symbols[start] if start in symbols else start,
                 symbols[end] if end in symbols else end)
        for region, (start, end) in regions.items()
    }

    parsed_blocks = {
        name:
        {k: symbols[v] if v in symbols else v
         for k, v in fields.items()}
        for name, fields in blocks.items()
    }

    return (parsed_regions, parsed_blocks)


def icf_regions_to_ld_regions(icf_regions: Dict) -> Dict:
    """Converts .icf regions to .ld regions

    The .icf format specifies the beginning and end of each region, while
    .ld expects the beginning and a length string.

    Args:
        icf_regions: icf_regions parsed with `parse_icf()`

    Returns:
        A map from `region_name` -> (start_hex, length_str)
    """
    ld_regions = {}
    for region, (start, end) in icf_regions.items():
        start_dec = int(start, 16)
        end_dec = int(end, 16)
        length = end_dec - start_dec + 1
        length_str = str(length)
        if length % 1024 == 0:
            length_str = f'{int(length/1024)}K'

        # Some icf scripts incorrectly have an exclusive region end.
        # This corrects for that.
        elif (length - 1) % 1024 == 0:
            length_str = f'{int((length-1)/1024)}K'

        # ST's gcc linker scripts generally use FLASH instead of ROM
        if region == 'ROM':
            region = 'FLASH'

        ld_regions[region] = (start, length_str)

    return ld_regions


def create_ld(ld_regions: Dict, blocks: Dict) -> str:
    """Create .ld file from template.

    This creates a barebones .ld file that *should* work for most single core
    stm32 families. It only contains regions for RAM and FLASH.

    This template can be bypassed in GN if a more sophisticated linker file
    is required.

    Args:
        ld_regions: generated by `icf_regions_to_ld_regions()`
        blocks: generated by `parse_icf`

    Returns:
        a string linker file with the RAM/FLASH specified by the given reginos.

    Raises:
        KeyError if ld_regions does not contain 'RAM' and 'FLASH'
    """
    return f"""\
ENTRY(Reset_Handler)
_estack = ORIGIN(RAM) + LENGTH(RAM);

_Min_Heap_Size = {blocks['HEAP']['size']};
_Min_Stack_Size = {blocks['CSTACK']['size']};

MEMORY
{{
  RAM (xrw) : ORIGIN = {ld_regions['RAM'][0]}, LENGTH = {ld_regions['RAM'][1]}
  FLASH (rx) : ORIGIN = {ld_regions['FLASH'][0]}, LENGTH = {ld_regions['FLASH'][1]}
}}

SECTIONS
{{

  /* The ARMv8-M architecture requires this is at least aligned to 128 bytes,
   * and aligned to a power of two that is greater than 4 times the number of
   * supported exceptions. 512 has been selected as it accommodates most vector
   * tables.
   */
  .isr_vector :
  {{
    . = ALIGN(512);
    KEEP(*(.isr_vector))
    . = ALIGN(4);
  }} >FLASH

  .text :
  {{
    . = ALIGN(4);
    *(.text)
    *(.text*)
    *(.glue_7)
    *(.glue_7t)
    *(.eh_frame)

    KEEP (*(.init))
    KEEP (*(.fini))

    . = ALIGN(4);
    _etext = .;
  }} >FLASH

  .rodata :
  {{
    . = ALIGN(4);
    *(.rodata)
    *(.rodata*)
    . = ALIGN(4);
  }} >FLASH

  .ARM.extab   : {{
    . = ALIGN(4);
    *(.ARM.extab* .gnu.linkonce.armextab.*)
    . = ALIGN(4);
  }} >FLASH

  .ARM : {{
    . = ALIGN(4);
    __exidx_start = .;
    *(.ARM.exidx*)
    __exidx_end = .;
    . = ALIGN(4);
  }} >FLASH

  .preinit_array     :
  {{
    . = ALIGN(4);
    PROVIDE_HIDDEN (__preinit_array_start = .);
    KEEP (*(.preinit_array*))
    PROVIDE_HIDDEN (__preinit_array_end = .);
    . = ALIGN(4);
  }} >FLASH

  .init_array :
  {{
    . = ALIGN(4);
    PROVIDE_HIDDEN (__init_array_start = .);
    KEEP (*(SORT(.init_array.*)))
    KEEP (*(.init_array*))
    PROVIDE_HIDDEN (__init_array_end = .);
    . = ALIGN(4);
  }} >FLASH

  .fini_array :
  {{
    . = ALIGN(4);
    PROVIDE_HIDDEN (__fini_array_start = .);
    KEEP (*(SORT(.fini_array.*)))
    KEEP (*(.fini_array*))
    PROVIDE_HIDDEN (__fini_array_end = .);
    . = ALIGN(4);
  }} >FLASH

  _sidata = LOADADDR(.data);
  .data :
  {{
    . = ALIGN(4);
    _sdata = .;
    *(.data)
    *(.data*)
    . = ALIGN(4);
    _edata = .;
  }} >RAM AT> FLASH

  . = ALIGN(4);
  .bss :
  {{
    _sbss = .;
    __bss_start__ = _sbss;
    *(.bss)
    *(.bss*)
    *(COMMON)

    . = ALIGN(4);
    _ebss = .;
    __bss_end__ = _ebss;
  }} >RAM

  /* The ARMv7-M architecture may require 8-byte alignment of the stack pointer
   * rather than 4 in some contexts and implementations, so this region is
   * 8-byte aligned (see ARMv7-M Architecture Reference Manual DDI0403E
   * section B1.5.7).
   */
  ._user_heap_stack :
  {{
    . = ALIGN(8);
    PROVIDE ( end = . );
    PROVIDE ( _end = . );
    . = . + _Min_Heap_Size;
    . = . + _Min_Stack_Size;
    . = ALIGN(8);
  }} >RAM

  /DISCARD/ :
  {{
    libc.a ( * )
    libm.a ( * )
    libgcc.a ( * )
  }}

  .ARM.attributes 0 : {{ *(.ARM.attributes) }}
}}
    """


def icf_to_ld(icf_path: pathlib.Path, ld_path: Optional[pathlib.Path]):
    """Convert icf file into an ld file.

    Note: This only works for ST generated .icf files.

    Args:
        icf_path: path to .icf file to convert
        ld_path: path to write generated .ld file or None.
                 If None, the .ld file is written to stdout.
    """
    with open(icf_path, 'rb') as icf_file:
        icf_str = icf_file.read().decode('utf-8', errors='ignore')

    icf_regions, blocks = parse_icf(icf_str)
    ld_regions = icf_regions_to_ld_regions(icf_regions)
    ld_str = create_ld(ld_regions, blocks)

    if ld_path:
        with open(ld_path, 'w') as ld_file:
            ld_file.write(ld_str)
    else:
        print(ld_str)