# pylint:disable=arguments-differ
from typing import Tuple, Optional, Dict, List
import string

from archinfo import Endness

from ailment.expression import Const, StackBaseOffset
from ailment.statement import Call, Store

from .base import PeepholeOptimizationStmtBase


ASCII_PRINTABLES = set(string.printable)
ASCII_DIGITS = set(string.digits)


class InlinedStrcpy(PeepholeOptimizationStmtBase):
    """
    Simplifies inlined string copying logic into calls to strcpy.
    """

    __slots__ = ()

    NAME = "Simplifying inlined strcpy"
    stmt_classes = (Store,)

    def optimize(self, stmt: Store, stmt_idx: int = None, block=None, **kwargs):
        if isinstance(stmt.data, Const):
            r, s = self.is_integer_likely_a_string(stmt.data.value, stmt.data.size, stmt.endness)
            if r:
                # replace it with a call to strncpy
                str_id = self.kb.custom_strings.allocate(s.encode("ascii"))
                return Call(
                    stmt.idx,
                    "strncpy",
                    args=[
                        stmt.addr,
                        Const(None, None, str_id, stmt.addr.bits, custom_string=True),
                        Const(None, None, len(s), self.project.arch.bits),
                    ],
                    **stmt.tags,
                )

            # scan forward in the current block to find all consecutive constant stores
            if block is not None and stmt_idx is not None:
                all_constant_stores: Dict[int, Tuple[int, Optional[Const]]] = self.collect_constant_stores(
                    block, stmt_idx
                )
                if all_constant_stores:
                    offsets = sorted(all_constant_stores.keys())
                    next_offset = min(offsets)
                    stride = []
                    for offset in offsets:
                        if next_offset is not None and offset != next_offset:
                            next_offset = None
                            stride = []
                        stmt_idx_, v = all_constant_stores[offset]
                        if v is not None:
                            stride.append((offset, stmt_idx_, v))
                            next_offset = offset + v.size
                        else:
                            next_offset = None
                            stride = []

                    integer, size = self.stride_to_int(stride)
                    r, s = self.is_integer_likely_a_string(integer, size, Endness.BE)
                    if r:
                        # we remove all involved statements whose statement IDs are greater than the current one
                        for _, stmt_idx_, _ in reversed(stride):
                            if stmt_idx_ <= stmt_idx:
                                continue
                            block.statements[stmt_idx_] = None
                        block.statements = [ss for ss in block.statements if ss is not None]

                        str_id = self.kb.custom_strings.allocate(s.encode("ascii"))
                        return Call(
                            stmt.idx,
                            "strncpy",
                            args=[
                                stmt.addr,
                                Const(None, None, str_id, stmt.addr.bits, custom_string=True),
                                Const(None, None, len(s), self.project.arch.bits),
                            ],
                            **stmt.tags,
                        )

        return None

    @staticmethod
    def stride_to_int(stride: List[Tuple[int, int, Const]]) -> Tuple[int, int]:
        stride = sorted(stride, key=lambda x: x[0])
        n = 0
        size = 0
        for _, _, v in stride:
            size += v.size
            n <<= v.bits
            n |= v.value
        return n, size

    @staticmethod
    def collect_constant_stores(block, starting_stmt_idx: int) -> Dict[int, Tuple[int, Optional[Const]]]:
        r = {}
        for idx, stmt in enumerate(block.statements):
            if idx < starting_stmt_idx:
                continue
            if isinstance(stmt, Store) and isinstance(stmt.addr, StackBaseOffset) and isinstance(stmt.addr.offset, int):
                if isinstance(stmt.data, Const):
                    r[stmt.addr.offset] = idx, stmt.data
                else:
                    r[stmt.addr.offset] = idx, None

        return r

    @staticmethod
    def is_integer_likely_a_string(
        v: int, size: int, endness: Endness, min_length: int = 4
    ) -> Tuple[bool, Optional[str]]:
        # we need at least four bytes of printable characters

        chars = []
        if endness == Endness.LE:
            while v != 0:
                byt = v & 0xFF
                if chr(byt) not in ASCII_PRINTABLES:
                    return False, None
                chars.append(chr(byt))
                v >>= 8

        elif endness == Endness.BE:
            first_non_zero = False
            for _ in range(size):
                byt = v & 0xFF
                v >>= 8
                if byt == 0:
                    if first_non_zero:
                        return False, None
                    continue
                first_non_zero = True  # this is the first non-zero byte
                if chr(byt) not in ASCII_PRINTABLES:
                    return False, None
                chars.append(chr(byt))
            chars = chars[::-1]
        else:
            # unsupported endness
            return False, None

        if len(chars) >= min_length:
            if len(chars) <= 4 and all(ch in ASCII_DIGITS for ch in chars):
                return False, None
            return True, "".join(chars)
        return False, None
