import os
import sys
import glob
import logging
import tempfile
import io

from contextlib import contextmanager

try:
    # nullcontext is not present before Python 3.7
    from contextlib import nullcontext  # type: ignore
except ImportError:

    class nullcontext:  # type: ignore
        def __init__(self, enter_result=None):
            self.enter_result = enter_result

        def __enter__(self):
            return self.enter_result

        def __exit__(self, *excinfo):
            pass


import time
import os
import errno
import threading
from ..utils import config
from ..utils.helpers import StateSerializer, PickleSerializer
from .smtlib.solver import SelectedSolver
from .state import StateBase
from ..exceptions import ManticoreError

from typing import Any, Dict, List, IO, Optional, Union

logger = logging.getLogger(__name__)

consts = config.get_group("workspace")
consts.add(
    "prefix", default="mcore_", description="The prefix to use for output and workspace directories"
)
consts.add("dir", default=".", description="Location of where to create workspace directories")


class Testcase:
    """
    A `WorkspaceTestCase` represents a test case input generated by Manticore.
    """

    def __init__(self, workspace: "ManticoreOutput", prefix: str):
        self._num = workspace._increment_id()
        self._prefix = prefix
        self._ws = workspace

    @property
    def prefix(self) -> str:
        return self._prefix

    @property
    def num(self) -> int:
        return self._num

    def open_stream(self, suffix: str, binary: bool = False):
        stream_name = f"{self._prefix}_{self._num:08x}.{suffix}"
        return self._ws.save_stream(stream_name, binary=binary)


class Store:
    """
    A `Store` can save arbitrary keys/values (including states) and file streams.
    Used for generating output, state saving and state loading.

    In subclasses:

     * Implement either save_value/load_value, or save_stream/load_stream, or both.
     * Define a `store_type` class variable of type str.
       * This is used as a prefix for a store descriptor
    """

    @classmethod
    def fromdescriptor(cls, desc):
        """
        Create a :class:`~manticore.core.workspace.Store` instance depending on the descriptor.

        Valid descriptors:
          * fs:<path>
          * redis:<hostname>:<port>
          * mem:

        :param str desc: Store descriptor
        :return: Store instance
        """
        type_, uri = ("fs", None) if desc is None else desc.split(":", 1)
        for subclass in cls.__subclasses__():
            if subclass.store_type == type_:
                return subclass(uri)
        raise NotImplementedError(f"Storage type '{type_}' not supported.")

    def __init__(self, uri: Optional[str], state_serialization_method: str = "pickle"):
        assert (
            self.__class__ != Store
        ), "The Store class can not be instantiated (create a subclass)"

        self.uri = uri
        self._sub: List = []

        if state_serialization_method == "pickle":
            self._serializer = PickleSerializer()
        else:
            raise NotImplementedError(
                f"Pickling method '{state_serialization_method}' not supported."
            )

    # save_value/load_value and save_stream/load_stream are implemented in terms of each other. A backing store
    # can choose the pair it's best optimized for.
    def save_value(self, key: str, value) -> None:
        """
        Save an arbitrary, serializable `value` under `key`.

        :param str key: A string identifier under which to store the value.
        :param value: A serializable value
        """
        with self.save_stream(key) as s:
            s.write(value)

    def load_value(self, key: str, binary: bool = False):
        """
        Load an arbitrary value identified by `key`.

        :param str key: The key that identifies the value
        :return: The loaded value
        """
        with self.load_stream(key, binary=binary) as s:
            return s.read()

    @contextmanager
    def save_stream(self, key: str, binary: bool = False):
        """
        Return a managed file-like object into which the calling code can write
        arbitrary data.

        :param key:
        :return: A managed stream-like object
        """
        s = io.BytesIO() if binary else io.StringIO()
        yield s
        self.save_value(key, s.getvalue())  # type: ignore

    @contextmanager
    def load_stream(self, key: str, binary: bool = False):
        """
        Return a managed file-like object from which the calling code can read
        previously-serialized data.

        :param key:
        :return: A managed stream-like object
        """
        value = self.load_value(key, binary=binary)
        yield io.BytesIO(value) if binary else io.StringIO(value)

    def save_state(self, state: StateBase, key: str) -> None:
        """
        Save a state to storage.

        :param manticore.core.StateBase state:
        :param key: the key to save under
        """
        with self.save_stream(key, binary=True) as f:
            self._serializer.serialize(state, f)

    def load_state(self, key: str, delete: bool = True) -> StateBase:
        """
        Load a state from storage.

        :param key: the key to load from
        """

        with self.load_stream(key, binary=True) as f:
            state = self._serializer.deserialize(f)
            if delete:
                self.rm(key)

            return state

    def rm(self, key: str) -> None:
        """
        Remove value identified by `key` from storage.

        :param key: What to remove
        """
        raise NotImplementedError

    def ls(self, glob_str: str) -> List[str]:
        """
        List all keys in storage that match `glob_str`.
        """
        raise NotImplementedError

    def lock(self):
        raise NotImplementedError


class FilesystemStore(Store):
    """
    A directory-backed Manticore workspace
    """

    store_type = "fs"

    def __init__(self, uri: Optional[str] = None):
        """
        :param uri: The path to on-disk workspace, or None.
        """
        self._tlock = threading.Lock()
        if not uri:
            uri = os.path.abspath(tempfile.mkdtemp(prefix=consts.prefix, dir=consts.dir))

        if os.path.exists(uri):
            assert os.path.isdir(uri), "Store must be a directory"
        else:
            os.mkdir(uri)

        super().__init__(uri)

    def _uri_of_key(self, key: str) -> str:
        """
        Produce a URI for the given key, designed to be usable as a filename.
        """
        assert self.uri is not None
        key = key.replace("/", "_")
        return os.path.join(self.uri, key)

    @contextmanager
    def lock(self):
        lockfile = self._uri_of_key(".lock")
        with self._tlock:
            while True:
                try:
                    fd = os.open(lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
                except OSError as e:
                    if e.errno != errno.EEXIST:
                        raise
                    time.sleep(0.05)
                else:
                    yield
                    os.close(fd)
                    os.unlink(lockfile)
                    break

    @contextmanager
    def stream(self, key: str, mode: str = "r", lock: bool = False):
        """
        Yield a file object representing `key`

        :param key: The file to save to
        :param mode: mode is an optional string that specifies the mode in which the file is opened
        :param lock: exclusive access if True
        :return:
        """
        fname = self._uri_of_key(key)
        if lock:
            with self.lock():
                with open(fname, mode) as f:
                    yield f
        else:
            with open(fname, mode) as f:
                yield f

    @contextmanager
    def save_stream(self, key: str, binary: bool = False, lock: bool = False):
        """
        Yield a file object representing `key`

        :param str key: The file to save to
        :param bool binary: Whether we should treat it as binary
        :param lock: exclusive access if True
        :return:
        """
        mode = "wb" if binary else "w"
        with self.stream(key, mode, lock) as f:
            yield f

    @contextmanager
    def load_stream(self, key: str, binary: bool = False, lock: bool = False):
        """
        :param str key: name of stream to load
        :param bool binary: Whether we should treat it as binary
        :param lock: exclusive access if True
        :return:
        """
        mode = "rb" if binary else "r"
        with self.stream(key, mode, lock) as f:
            yield f

    def rm(self, key: str) -> None:
        """
        Remove file identified by `key`.

        :param key: The file to delete
        """
        path = self._uri_of_key(key)
        if os.path.exists(path):
            os.remove(path)

    def ls(self, glob_str: str) -> List[str]:
        """
        Return just the filenames that match `glob_str` inside the store directory.

        :param glob_str: A glob string, e.g. 'state_*'
        :return: list of matched keys
        """
        path = self._uri_of_key(glob_str)
        return [os.path.split(s)[1] for s in glob.glob(path)]


class MemoryStore(Store):
    """
    An in-memory (dict) Manticore workspace.

    NOTE: This is mostly used for experimentation and testing functionality.
    Can not be used with multiple workers!
    """

    store_type = "mem"

    # TODO(yan): Once we get a global config store, check it to make sure
    # we're executing in a single-worker or test environment.

    def __init__(self, uri: Optional[str] = None):
        self._lock = threading.RLock()
        self._data: Dict = {}
        super().__init__(None)

    def save_value(self, key, value):
        self._data[key] = value

    def load_value(self, key, binary=False):
        return self._data.get(key)

    def rm(self, key):
        del self._data[key]

    def ls(self, glob_str):
        return list(self._data)

    @contextmanager
    def lock(self):
        with self._lock:
            yield

    @contextmanager
    def stream(self, key, mode="r", lock=False):
        if lock:
            raise ManticoreError("mem: does not support concurrency")
        if "b" in mode:
            s = io.BytesIO(self._data.get(key, b""))
        else:
            s = io.StringIO(self._data.get(key, ""))
        yield s
        self._data[key] = s.getvalue()


class RedisStore(Store):
    """
    A redis-backed Manticore workspace
    """

    store_type = "redis"

    def __init__(self, uri=None):
        """
        :param uri: A url for redis
        """

        # Local import to avoid an explicit dependency
        import redis

        hostname, port = uri.split(":")
        self._client = redis.StrictRedis(host=hostname, port=int(port), db=0)

        super().__init__(uri)

    @contextmanager
    def lock(self):
        with self._client.lock(".lock"):
            yield

    def save_value(self, key, value):
        """
        Save an arbitrary, serializable `value` under `key`.

        :param str key: A string identifier under which to store the value.
        :param value: A serializable value
        :return:
        """
        return self._client.set(key, value)

    def load_value(self, key):
        """
        Load an arbitrary value identified by `key`.

        :param str key: The key that identifies the value
        :return: The loaded value
        """
        return self._client.get(key)

    def rm(self, key):
        self._client.delete(key)

    def ls(self, glob_str):
        return self._client.keys(glob_str)


class Workspace:
    """
    A workspace maintains a list of states to run and assigns them IDs.
    """

    def __init__(self, store_or_desc: Union[None, Store, str] = None):
        if isinstance(store_or_desc, Store):
            self._store: Store = store_or_desc
        else:
            self._store = Store.fromdescriptor(store_or_desc)
        self._serializer: StateSerializer = PickleSerializer()
        self._prefix: str = "state_"
        self._suffix: str = ".pkl"

    @property
    def uri(self):
        return self._store.uri

    def try_loading_workspace(self):
        state_names = self._store.ls(f"{self._prefix}*")

        def get_state_id(name):
            return int(name[len(self._prefix) : -len(self._suffix)], 16)

        state_ids = list(map(get_state_id, state_names))

        if not state_ids:
            return []

        return state_ids

    def _get_id(self):
        """
        Get a unique state id.

        :rtype: int
        """
        with self._store.lock():
            try:
                with self._store.load_stream(".state_id") as f:
                    last_id = int(f.read())
            except Exception as e:
                last_id = 0
            else:
                last_id += 1
            with self._store.save_stream(".state_id") as f:
                f.write(f"{last_id}")
                f.flush()
        return last_id

    def load_state(self, state_id, delete=True):
        """
        Load a state from storage identified by `state_id`.

        :param state_id: The state reference of what to load
        :return: The deserialized state
        :rtype: State
        """
        return self._store.load_state(f"{self._prefix}{state_id:08x}{self._suffix}", delete=delete)

    def save_state(self, state: StateBase, state_id=None):
        """
        Save a state to storage, return identifier.

        :param state: The state to save
        :param int state_id: If not None force the state id potentially overwriting old states
        :return: New state id
        :rtype: int
        """
        assert isinstance(state, StateBase)
        if state_id is None:
            state_id = self._get_id()
        else:
            self.rm_state(state_id)

        self._store.save_state(state, f"{self._prefix}{state_id:08x}{self._suffix}")
        return state_id

    def rm_state(self, state_id):
        """
        Remove a state from storage identified by `state_id`.

        :param state_id: The state reference of what to load
        """
        return self._store.rm(f"{self._prefix}{state_id:08x}{self._suffix}")


class ManticoreOutput:
    """
    Functionality related to producing output. Responsible for generating state summaries,
    coverage information, etc.

    Invoked only from :class:`manticore.Manticore` from a single parent process, so
    locking is not required.
    """

    def __init__(self, desc=None):
        """
        Create an object capable of producing Manticore output.

        :param desc: A descriptor ('type:uri') of where to write output.
        """
        self._named_key_prefix = "test"
        self._descriptor = desc
        self._store = Store.fromdescriptor(desc)

    def testcase(self, prefix: str = "test") -> Testcase:
        return Testcase(self, prefix)

    @property
    def store(self):
        return self._store

    @property
    def descriptor(self):
        """
        Return a descriptor that created this workspace. Descriptors are of the
        format <type>:<uri>, where type signifies the medium. For example,
          fs:/tmp/workspace
          redis:127.0.0.1:6379

        :rtype: str
        """
        if self._descriptor is None:
            self._descriptor = f"{self._store.store_type}:{self._store.uri}"

        return self._descriptor

    def _increment_id(self) -> int:
        """
        Get a unique testcase id.
        """
        filename = ".testcase_id"
        with self._store.lock():
            try:
                with self._store.stream(filename, "r") as f:
                    last_id = int(f.read())
            except Exception as e:
                last_id = 0
            else:
                last_id += 1
            with self._store.stream(filename, "w") as f:
                f.write(f"{last_id}")
                f.flush()
        return last_id

    @property
    def _last_id(self):
        filename = ".testcase_id"
        try:
            with self._store.stream(filename, "r") as f:
                last_id = int(f.read())
        except Exception as e:
            last_id = 0
        return last_id

    def _named_key(self, suffix):
        return f"{self._named_key_prefix}_{self._last_id:08x}.{suffix}"

    def save_stream(self, key, *rest, **kwargs):
        return self._store.save_stream(key, *rest, **kwargs)

    @contextmanager
    def _named_stream(self, name, binary=False, lock=False):
        """
        Create an indexed output stream e.g. 'test_00000001.name'

        :param name: Identifier for the stream
        :param lock: exclusive access if True
        :return: A context-managed stream-like object
        """
        with self._store.save_stream(self._named_key(name), binary=binary, lock=lock) as s:
            yield s

    # Remove/move ...
    def save_testcase(self, state: StateBase, testcase: Testcase, message: str = ""):
        """
        Save the environment from `state` to storage. Return a state id
        describing it, which should be an int or a string.

        :param state: The state to serialize
        :param message: The message to add to output
        :return: A state id representing the saved state
        """

        # TODO / FIXME: We should move workspace to core/workspace and create a workspace for evm and native binaries
        # The workspaces should override `save_testcase` method
        #
        # Below is native-only
        self.save_summary(testcase, state, message)
        self.save_trace(testcase, state)
        self.save_constraints(testcase, state)
        self.save_input_symbols(testcase, state)

        for stream_name, data in state.platform.generate_workspace_files().items():
            with testcase.open_stream(stream_name, binary=True) as stream:
                if isinstance(data, str):
                    data = data.encode()
                stream.write(data)

        # self._store.save_state(state, self._named_key('pkl'))
        return testcase

    @staticmethod
    def save_summary(testcase, state, message):
        with testcase.open_stream("messages") as summary:
            summary.write(f"Command line:\n  '{' '.join(sys.argv)}'\n")
            summary.write(f"Status:\n  {message}\n\n")

            memories = set()
            for cpu in filter(None, state.platform.procs):
                idx = state.platform.procs.index(cpu)
                summary.write(f"================ PROC: {idx:02d} ================\n")
                summary.write("Memory:\n")
                if hash(cpu.memory) not in memories:
                    summary.write(str(cpu.memory).replace("\n", "\n  "))
                    memories.add(hash(cpu.memory))

                summary.write(f"CPU:\n{cpu}")

                if hasattr(cpu, "instruction") and cpu.instruction is not None:
                    i = cpu.instruction
                    summary.write(f"  Instruction: 0x{i.address:x}\t{i.mnemonic:s} {i.op_str:s})\n")
                else:
                    summary.write("  Instruction: {symbolic}\n")

    @staticmethod
    def save_trace(testcase, state: StateBase):
        with testcase.open_stream("trace") as f:
            if "trace" not in state.context:
                return
            for entry in state.context["trace"]:
                f.write(f"0x{entry:x}\n")

    @staticmethod
    def save_constraints(testcase, state: StateBase):
        # XXX(yan): We want to conditionally enable this check
        # assert solver.check(state.constraints)

        with testcase.open_stream("smt") as f:
            f.write(str(state.constraints))

    @staticmethod
    def save_input_symbols(testcase, state: StateBase):
        with testcase.open_stream("input") as f:
            bufs = state.solve_one_n_batched(state.input_symbols)
            for symbol, buf in zip(state.input_symbols, bufs):
                # TODO - constrain=False here, so the extra migration shouldn't cause problems, right?
                f.write(f"{symbol.name}: {buf!r}\n")
