"""postgresql_chat_message_history_demo.py — PostgresChatMessageHistory + RunnableWithMessageHistory"""

import os
from dotenv import load_dotenv
import psycopg
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_postgres import PostgresChatMessageHistory

load_dotenv()

POSTGRES_URL = os.getenv(
    "POSTGRES_URL",
    "postgresql://postgres:postgres@localhost:5432/chatdb",
)
TABLE_NAME = "chat_history"

model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

prompt = ChatPromptTemplate.from_messages([
    ("system", "Keep each reply to one short sentence."),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}"),
])

chain = prompt | model

conn = psycopg.connect(POSTGRES_URL)
PostgresChatMessageHistory.create_tables(conn, TABLE_NAME)

def get_session_history(session_id: str):
    return PostgresChatMessageHistory(
        TABLE_NAME,
        session_id,
        sync_connection=conn,
    )

chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
)

config = {"configurable": {"session_id": "demo-1"}}

turns = [
    "What does the HTML <a> tag do?",
    "What attribute opens the link in a new tab?",
]

for question in turns:
    reply = chain_with_history.invoke({"input": question}, config=config)
    print(f"Human: {question}")
    print(f"AI:   {reply.content}\n")

print(f"History stored in PostgreSQL table '{TABLE_NAME}'")