import json
import uuid
from datetime import datetime

import psycopg2

# from psycopg2 import sql
# from psycopg2.extras import RealDictCursor


class PostgreSQLManager:
    def __init__(self, dbname: str, user: str, password: str, host: str = "localhost", port: int = 5432):
        """Initialize the PostgreSQLManager with a connection to the database.

        Args:
            dbname (str): Name of the PostgreSQL database.
            user (str): Database user.
            password (str): Database password.
            host (str): Database host. Defaults to "localhost".
            port (int): Database port. Defaults to 5432.
        """
        self.connection = psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port)
        self._create_shortterm_memories_table()

    def _create_shortterm_memories_table(self) -> None:
        """Create the shortterm memory table if does not exist"""
        with self.connection.cursor() as cursor:
            cursor.execute("""
               CREATE TABLE IF NOT EXISTS shortterm_memories (
               id UUID PRIMARY KEY,
               content TEXT,
               user_id TEXT,
               agent_id TEXT,
               run_id TEXT,
               table_name TEXT,
               created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
               );
            """)
            self.connection.commit()
            cursor.close()

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

    def get(self, table_name, run_id, timestamp):
        with self.connection.cursor() as cursor:
            cursor.execute(
                """
                    SELECT id, content, user_id, agent_id, run_id, created_at
                    FROM shortterm_memories
                    WHERE run_id = %s and table_name = %s and created_at <= %s
                    ORDER BY created_at ASC
                    """,
                (
                    run_id,
                    table_name,
                    timestamp,
                ),
            )
            self.connection.commit()
            rows = cursor.fetchall()
            cursor.close()
            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):
        with self.connection.cursor() as cursor:
            cursor.execute(
                """
               SELECT DISTINCT run_id
               FROM shortterm_memories where created_at <= %s
               """,
                (timestamp,),
            )
            self.connection.commit()
            rows = cursor.fetchall()
            cursor.close()
            return [
                {
                    "run_id": row[0],
                }
                for row in rows
            ]

    def delete(self, id):
        with self.connection.cursor() as cursor:
            cursor.execute(
                """DELETE from shortterm_memories where id = %s""",
                (id,),
            )
            self.connection.commit()
            cursor.close()

    def reset(self) -> None:
        """Reset the shortterm_memories table by dropping and recreating it."""
        with self.connection.cursor() as cursor:
            cursor.execute("DROP TABLE IF EXISTS shortterm_memories;")
            self.connection.commit()
            cursor.close()
            self._create_shortterm_memories_table()

    def close(self) -> None:
        """Close the database connection."""
        self.connection.close()
