from sqlalchemy import Table, Column, Integer, String, MetaData, insert, select, delete
from typing import Literal

from src.Credenciales.postgres_credencial import PostgresCredencial
from src.ConexionSql.Postgres_conexion import PostgresConexion  # Usamos la clase específica
from llms.Memory.Base_MemoryConv import MemoryConvABC


class MemoryConvPostgres(MemoryConvABC):
    def __init__(self, nombre_tabla: str, k: int, credencial: PostgresCredencial):
        """
        :param nombre_tabla: Nombre de la tabla en PostgreSQL.
        :param k: Número de pares user/assistant a recuperar.
        :param credencial: Credencial para conectar a PostgreSQL.
        """
        super().__init__(k=k)
        self.conexion = PostgresConexion(credencial)  # Se instancia directamente con la credencial
        self.nombre_tabla = nombre_tabla

        self.metadata = MetaData()
        self.tabla = Table(
            self.nombre_tabla,
            self.metadata,
            Column("id", Integer, primary_key=True, autoincrement=True),
            Column("rol", String, nullable=False),
            Column("contenido", String, nullable=False),
        )

        # Crea la tabla si no existe
        self.metadata.create_all(self.conexion.engine)

    def guardar_turno(self, rol: Literal["user", "assistant"], contenido: str) -> None:
        stmt = insert(self.tabla).values(rol=rol, contenido=contenido)
        with self.conexion.get_session() as session:
            session.execute(stmt)
            session.commit()

    def cargar_historial_chat(self) -> list[dict]:
        stmt = (
            select(self.tabla.c.rol, self.tabla.c.contenido)
            .order_by(self.tabla.c.id.desc())
            .limit(self.k * 2)
        )
        with self.conexion.get_session() as session:
            result = session.execute(stmt).fetchall()
        return [{"role": r.rol, "content": r.contenido} for r in reversed(result)]

    def limpiar(self) -> None:
        stmt = delete(self.tabla)
        with self.conexion.get_session() as session:
            session.execute(stmt)
            session.commit()