try:
    import sqlite3
except ImportError:
    from pysqlite2 import dbapi2 as sqlite3

from landscape.package.store import with_cursor


class ManagerStore(object):

    def __init__(self, filename):
        self._db = sqlite3.connect(filename)
        ensure_schema(self._db)

    @with_cursor
    def get_graph(self, cursor, graph_id):
        cursor.execute(
            "SELECT graph_id, filename, user FROM graph WHERE graph_id=?",
            (graph_id,))
        return cursor.fetchone()

    @with_cursor
    def get_graphs(self, cursor):
        cursor.execute("SELECT graph_id, filename, user FROM graph")
        return cursor.fetchall()

    @with_cursor
    def add_graph(self, cursor, graph_id, filename, user):
        cursor.execute(
            "SELECT graph_id FROM graph WHERE graph_id=?",
            (graph_id,))
        if cursor.fetchone():
            cursor.execute(
                "UPDATE graph SET filename=?, user=? WHERE graph_id=?",
                (filename, user, graph_id))
        else:
            cursor.execute(
                "INSERT INTO graph (graph_id, filename, user) VALUES (?, ?, ?)",
                (graph_id, filename, user))

    @with_cursor
    def remove_graph(self, cursor, graph_id):
        cursor.execute("DELETE FROM graph WHERE graph_id=?", (graph_id,))

    @with_cursor
    def set_graph_accumulate(self, cursor, graph_id, timestamp, value):
        cursor.execute(
            "SELECT graph_id, graph_timestamp, graph_value FROM "
            "graph_accumulate WHERE graph_id=?", (graph_id,))
        graph_accumulate = cursor.fetchone()
        if graph_accumulate:
            cursor.execute(
                "UPDATE graph_accumulate SET graph_timestamp = ?, "
                "graph_value = ? WHERE graph_id=?",
                (timestamp, value, graph_id))
        else:
            cursor.execute(
                "INSERT INTO graph_accumulate (graph_id, graph_timestamp, "
                "graph_value) VALUES (?, ?, ?)", (graph_id, timestamp, value))

    @with_cursor
    def get_graph_accumulate(self, cursor, graph_id):
        cursor.execute(
            "SELECT graph_id, graph_timestamp, graph_value FROM "
            "graph_accumulate WHERE graph_id=?", (graph_id,))
        return cursor.fetchone()


def ensure_schema(db):
    cursor = db.cursor()
    try:
        cursor.execute("CREATE TABLE graph"
                       " (graph_id INTEGER PRIMARY KEY,"
                       " filename TEXT NOT NULL, user TEXT)")
        cursor.execute("CREATE TABLE graph_accumulate"
                       " (graph_id INTEGER PRIMARY KEY,"
                       " graph_timestamp INTEGER, graph_value FLOAT)")
    except sqlite3.OperationalError:
        cursor.close()
        db.rollback()
    else:
        cursor.close()
        db.commit()