import json
import sqlite3 as sql


class InMemoryStorer:
    """
    Mainly to be used for testing.

    .. warning::
        This storer cannot be used together with a distributed execution.
        Consider to use the :class:`DistributedFileStorer <connectome.builder.persistence.DistributedFileStorer>`
        instead.
    """
    def __init__(self):
        self._store = {}

    def store(self, model, result):
        """
        Store the result of the model.


        Parameters
        ----------
        model: Model
            The executed model.

        result: dict
            The execution result.
        """
        if model.name not in self._store:
            self._store[model.name] = []
        self._store[model.name].append(result)

    def fetch_all(self):
        """
        Returns
        -------
        all_results: list
            A list of dictionaries of results.
        """
        return self._store


class DistributedFileStorer:
    """
    Store sampling results.

    Parameters
    ----------

    db_path: str
        Path, where to store the results.
        Results are stored in a SQLite database.

    timeout: int
        Timeout in seconds for SQLite write.
        Set this to a high number if you execute many sampling tasks
        in parallel on a cluster.
    """
    def __init__(self, db_path, timeout=5):
        self.db_path = db_path
        self.timeout = timeout
        self._connection = None

    def __repr__(self):
        return "{}({})".format(self.__class__.__name__, self.db_path)

    def _create_tables(self):
        self._execute_and_commit_query("""CREATE TABLE IF NOT EXISTS objects(id INTEGER  PRIMARY  KEY,
                                                                             DATE DEFAULT CURRENT_TIMESTAMP,
                                                                             name TEXT,
                                                                             content TEXT);""")

    def _execute_and_commit_query(self, query, *args):
        self._connect()
        cursor = self._connection.execute(query, *args)
        self._connection.commit()
        result = cursor.fetchall()
        self._close()
        return result

    def _close(self):
        self._connection.close()

    def _connect(self):
        self._connection = sql.connect(self.db_path, timeout=self.timeout)

    def store(self, model, result):
        """
        Store the result of the model.


        Parameters
        ----------
        model: Model
            The executed model.

        result: dict
            The execution result.
        """
        self._create_tables()
        result_json = json.dumps(result)
        self._execute_and_commit_query("INSERT INTO objects(name, content) VALUES (?,?);", (model.name, result_json))

    def fetch_all(self):
        """
        Returns
        -------
        all_results: list
            A list of dictionaries of results.
        """
        db_content = self._execute_and_commit_query("SELECT * FROM objects;")
        result = []
        for id, date, name, value_encoded in db_content:
            value = json.loads(value_encoded)
            result.append({"id": id, "name": name, "value": value, "date": date})
        return result
