import dataclasses
import json
import logging
import string
import struct
from typing import Dict, Iterator, Tuple

import click
import prettytable
from elftools.elf.constants import SH_FLAGS
from elftools.elf.elffile import ELFFile

_LOG = logging.getLogger(__name__)


@dataclasses.dataclass
class LogFormatInfo:
    filename: str
    line: int
    n_args: int
    fmt: str


LOG_FMT_SECTION_ELF_HEADER_MAGIC = 0x66474F4C  # 'LOGf'
LOG_FMT_SECTION_ELF_VERSION = 1
LOG_FMT_HDR_LEN = 8


@dataclasses.dataclass
class Header:
    magic: int
    version: int

    @classmethod
    def parse(cls, data: bytes) -> "Header":
        # Maps to sMemfaultLogFmtElfSectionHeader from compact_log_helpers.h
        magic, version = struct.unpack("<IBxxx", data[0:LOG_FMT_HDR_LEN])
        return cls(magic & 0xFFFFFFFF, version)


class LogFormatElfSectionParser:
    def find_section(self, elf: ELFFile) -> Tuple[Header, bytes]:
        for section in elf.iter_sections():
            if section.name != "log_fmt":
                continue

            sh_flags = section["sh_flags"]

            if sh_flags & SH_FLAGS.SHF_ALLOC != 0:
                raise ValueError(
                    "The 'log_fmt' section should not be allocated. Mark it as (INFO) in the ELF"
                )

            data = section.data()
            header = Header.parse(data)

            if header.magic != LOG_FMT_SECTION_ELF_HEADER_MAGIC:
                raise ValueError(f"Invalid 'log_fmt' section magic: 0x{header.magic:02x}")

            return header, data[LOG_FMT_HDR_LEN:]

        raise ValueError("Could not locate 'log_fmt' section")

    def parse(self, section: Tuple[Header, bytes]) -> Dict[int, LogFormatInfo]:
        header, body = section

        if header.version != LOG_FMT_SECTION_ELF_VERSION:
            raise ValueError(f"Unsupported 'log_fmt' version: {header.version}")

        return dict(self._parse(body))

    def _parse(self, body: bytes) -> Iterator[Tuple[int, LogFormatInfo]]:
        first_error = None
        log_fmt_id = LOG_FMT_HDR_LEN

        # Note: Each log_fmt entry should be separated by a NUL character because we are parsing a list
        # of strings generated by the C Compiler
        sep = b"\0"
        for part in body.split(sep):
            try:
                if part:
                    entry = self._parse_entry(part)
                    yield log_fmt_id, entry
            except ValueError as exc:
                first_error = exc, part
            log_fmt_id += len(part) + len(sep)

        # avoid spamming log output
        if first_error:
            _LOG.warning("Failed to parse: %s", first_error)

    def _parse_entry(self, entry: bytes) -> LogFormatInfo:
        # The format of a V1 entry: (n_args;__FILE__;__LINE__;fmt_str)
        # See MEMFAULT_LOG_FMT_ELF_SECTION_ENTRY macro in compact_log_helpers.h
        fields = entry.split(b";", maxsplit=3)
        n_args_bytes, filename, line, log_fmt = fields
        try:
            n_args = int(n_args_bytes)
        except ValueError:
            # Empirically, certain compilers (such as the TI ARM Compiler) may emit
            # junk (padding?) bytes between strings. Therefore, we strip any leading
            # non-decimal characters from n_args to work around the problem.
            digits = set(string.digits.encode())
            n_args_bytes = bytes(c for c in n_args_bytes if c in digits)
            n_args = int(n_args_bytes)

        return LogFormatInfo(
            filename=filename.decode(),
            line=int(line),
            n_args=int(n_args),
            fmt=log_fmt.decode(),
        )

    @classmethod
    def get_mapping_from_elf_file(cls, filepath: str) -> Dict[int, LogFormatInfo]:
        """Convenience method to loading mapping from an ELF file path"""
        with open(filepath, "rb") as f:
            elf = ELFFile(f)
            parser = cls()
            section = parser.find_section(elf)
            return parser.parse(section)


@click.command()
@click.argument("elf", type=click.Path(exists=True))
@click.option(
    "-i",
    "--log_id",
    default=None,
    help="log_id to filter on. If not provided entire list is dumped",
)
@click.option(
    "--json-output",
    is_flag=True,
    show_default=True,
    default=False,
    help="Ouput in JSON format instead of table",
)
def main(elf, log_id, json_output):
    """
    Given an ELF, looks for the log_fmt section & dumps all formatters found

    Example Usage:

    \b
    $ python format.py my_build.elf
        Found 3 log formats
        +--------+-----------+------+----------+----------+
        | Log Id | File Name | Line | Num Args |  Format  |
        +--------+-----------+------+----------+----------+
        |   8    |   main.c  | 151  |    1     | Hello %d |
        |                      ...                        |
        +--------+-----------+------+----------+----------+

    """
    mappings = LogFormatElfSectionParser.get_mapping_from_elf_file(elf)

    if json_output:
        output = {}
        for k, v in mappings.items():
            output[k] = dataclasses.asdict(v)

        # sort by log_id; should be sorted by default but just in case
        output = {k: v for k, v in sorted(output.items(), key=lambda item: item[0])}
        click.echo(json.dumps(output))
        return
    else:
        pt = prettytable.PrettyTable(["Log Id", "File Name", "Line", "Num Args", "Format"])

        log_ids = list(mappings.keys())
        log_ids.sort()
        print(f"Found {len(log_ids)} log formats")

        if log_id is not None:
            log_ids = [n for n in log_ids if n == int(log_id)]

        for k in log_ids:
            v = mappings[k]
            pt.add_row([k, v.filename, v.line, v.n_args, v.fmt])

        print(pt)


if __name__ == "__main__":
    main()
