from transformers import pipeline
import psycopg2

# Connect to the School database

def summarize_articles(summfunc):
    """ Connect to the PostgreSQL database server """

    conn = psycopg2.connect(
        dbname="crowsnest",
        user="crow",
        password="4LlKpnQ2RZPzL13BSpkW4k",
        host="192.168.0.2"
    )

    try:
        # create a cursor
        cur = conn.cursor()

        for i in range(20):
            print(i)
            # execute a statement
            cur.execute("""
                SELECT id, content
                FROM articles 
		        WHERE aisummary = ''
                ORDER BY publishDate DESC
		        LIMIT 20
            """)

            # display the PostgreSQL database server version
            keys, contents = list(zip(*cur.fetchall()))
            
            for key, summary in zip(keys, summfunc(list(contents))):
                cur.execute("UPDATE articles SET aisummary = %s WHERE id = %s", (summary, key))
        
            conn.commit()

        # close the communication with the PostgreSQL
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        raise error
    finally:
        if conn is not None:
            conn.close()
            print('Database connection closed.')


if __name__ == '__main__':
    pipe = pipeline("summarization", model="Shahm/bart-german")
    summarize_articles(lambda x: [str(e["summary_text"]) for e in pipe(x)])