import json
import sqlite3
from Memory import Memory


class Storage(Memory):
    def __init__(self, size):
        super().__init__(1)
        self.size = size
        self.conn = sqlite3.connect("Storage.se")
        self.create_table()

    def create_table(self):
        cursor = self.conn.cursor()
        cursor.execute(
            """CREATE TABLE IF NOT EXISTS KeyValue (
                            key INTEGER PRIMARY KEY,
                            value TEXT
                        )"""
        )
        self.conn.commit()

    def __setitem__(self, key, value):
        cursor = self.conn.cursor()
        value = json.dumps([value])
        cursor.execute(
            """
            INSERT INTO KeyValue (key, value) VALUES (?, ?)
            ON CONFLICT(key) DO UPDATE SET value = excluded.value
        """,
            (key, value),
        )
        self.conn.commit()

    def __getitem__(self, key):
        cursor = self.conn.cursor()
        cursor.execute("SELECT value FROM KeyValue WHERE key = ?", (key,))
        result = cursor.fetchone()
        if result:
            return json.loads(result[0])[0]  # 返回 value
        return 0