import json
import sqlite3
import uuid
from datetime import datetime


class SQLiteManager:
    def __init__(self, db_path=":memory:"):
        self.connection = sqlite3.connect(db_path, check_same_thread=False)
        self._create_shortterm_memories_table()

    def _create_shortterm_memories_table(self):
        with self.connection:
            self.connection.execute(
                """
                CREATE TABLE IF NOT EXISTS shortterm_memories (
                  id TEXT PRIMARY KEY,
                  content TEXT,
                  user_id TEXT,
                  agent_id TEXT,
                  run_id TEXT,
                  table_name TEXT,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """
            )
            self.connection.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_shortterm_memories on shortterm_memories("user_id", "agent_id", "run_id", "created_at")
            """
            )

    def add(
        self,
        content,
        user_id,
        agent_id,
        run_id,
        table_name,
        created_at=None,
    ):
        with self.connection:
            self.connection.execute(
                """
                INSERT INTO shortterm_memories (id, content, user_id, agent_id, run_id, table_name, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
                (
                    str(uuid.uuid4()),
                    content.json(),
                    user_id,
                    agent_id,
                    run_id,
                    table_name,
                    created_at or datetime.now(),
                ),
            )

    def get(self, table_name, run_id, timestamp):
        cursor = self.connection.execute(
            """
            SELECT id, content, user_id, agent_id, run_id, created_at
            FROM shortterm_memories
            WHERE run_id = ? and table_name = ? and created_at <= ?
            ORDER BY created_at ASC
        """,
            (
                run_id,
                table_name,
                timestamp,
            ),
        )
        rows = cursor.fetchall()
        return [
            {
                "id": row[0],
                "content": json.loads(row[1]),
                "user_id": row[2],
                "agent_id": row[3],
                "run_id": row[4],
                "created_at": row[5],
            }
            for row in rows
        ]

    def get_runids(self, timestamp):
        cursor = self.connection.execute(
            """
            SELECT DISTINCT run_id
            FROM shortterm_memories where created_at <= ?
        """,
            (timestamp,),
        )
        rows = cursor.fetchall()
        return [
            {
                "run_id": row[0],
            }
            for row in rows
        ]

    def delete(self, id):
        with self.connection:
            self.connection.execute(
                """DELETE from shortterm_memories where id = ?""",
                (id,),
            )

    def reset(self):
        with self.connection:
            self.connection.execute("DROP INDEX IF EXISTS idx_shortterm_memories")
            self.connection.execute("DROP TABLE IF EXISTS shortterm_memories")
            self._create_shortterm_memories_table()
