import claripy
import logging

from . import MemoryMixin
from ... import sim_options as o

l = logging.getLogger(__name__)

class UnderconstrainedMixin(MemoryMixin):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._unconstrained_range = 1024

    def load(self, addr, **kwargs):
        self._constrain_underconstrained_index(addr)
        return super().load(addr, **kwargs)

    def store(self, addr, data, **kwargs):
        self._constrain_underconstrained_index(addr)
        super().store(addr, data, **kwargs)

    def _default_value(self, addr, size, name='mem', key=None, inspect=True, events=True, **kwargs):
        if o.UNDER_CONSTRAINED_SYMEXEC in self.state.options and type(addr) is int:
            if self.category == 'mem':
                alloc_depth = self.state.uc_manager.get_alloc_depth(addr)
                uc_alloc_depth = 0 if alloc_depth is None else alloc_depth + 1
            else:
                uc_alloc_depth = 0

            bits = size * self.state.arch.byte_width
            return self.state.solver.Unconstrained(name, bits, key=key, inspect=inspect, events=events, uc_alloc_depth=uc_alloc_depth)

        return super()._default_value(addr, size, name=name, key=key, inspect=inspect, events=events, **kwargs)

    def _constrain_underconstrained_index(self, addr):
        if o.UNDER_CONSTRAINED_SYMEXEC in self.state.options and \
                isinstance(addr, claripy.ast.Base) and \
                addr.uninitialized and \
                addr.uc_alloc_depth is not None:
            if not self.state.uc_manager.is_bounded(addr) or self.state.solver.max_int(addr) - self.state.solver.min_int(addr) >= self._unconstrained_range:
                # in under-constrained symbolic execution, we'll assign a new memory region for this address
                mem_region = self.state.uc_manager.assign(addr)

                # ... but only if it's not already been constrained to something!
                if self.state.solver.solution(addr, mem_region):
                    self.state.add_constraints(addr == mem_region)
                l.debug('Under-constrained symbolic execution: assigned a new memory region @ %s to %s', mem_region, addr)

