from sqlalchemy.sql import text
from sqlalchemy import create_engine
import os

def create_trigger_function(conn):
    auto_dealid = text("""
    CREATE OR REPLACE FUNCTION auto_dealid()
    RETURNS TRIGGER AS $$
    DECLARE stub text;
            sqlstr text;
    BEGIN
      sqlstr:= 'UPDATE '|| TG_TABLE_NAME ||' SET %s WHERE id = %L AND dealid is NULL';

       IF (TG_TABLE_NAME = 'bonds') THEN
          stub := 'SC_';
          sqlstr := format(sqlstr, 'dealid = $1||upper(left(asset_class::text,3))||id,
identifier = COALESCE(identifier, cusip, isin)', NEW.id);
       ELSIF (TG_TABLE_NAME = 'cds') THEN
          stub := 'SCCDS';
          sqlstr := format(sqlstr, 'dealid = $1||id', NEW.id);
       ELSIF (TG_TABLE_NAME = 'repo') THEN
          stub := 'SC_REP';
          sqlstr := format(sqlstr, 'dealid = $1||id', NEW.id);
       ELSIF (TG_TABLE_NAME = 'swaptions') THEN
          stub := 'SWPTN';
          sqlstr := format(sqlstr, 'dealid = $1||id', NEW.id);
       END IF;
       EXECUTE sqlstr USING stub;
       RETURN NEW;
    END;
    $$ language plpgsql""")
    conn.execute(auto_dealid)

def create_triggers():
    trigger = db.DDL("""CREATE TRIGGER dealid
    AFTER INSERT ON %(fullname)s
    FOR EACH ROW
    EXECUTE PROCEDURE auto_dealid()""")
    for name, tb in db.metadata.tables.items():
        db.event.listen(tb, 'after_create', trigger)

if __name__ == "__main__":
    """ This script will create the tables and triggers for the Dawn app.
    If schema is not None, it will create it under a specific schema"""
    engine = create_engine('postgresql://qa@debian/dawndb')
    conn = engine.connect()
    schema = 'qa'
    #conn.execute("CREATE SCHEMA IF NOT EXISTS {}".format(schema))
    conn.execute("SET search_path TO %s", (schema,))
    create_trigger_function(conn)
    conn.close()

    os.environ['CONF'] = 'config.ini'
    from Dawn import db
    create_triggers()
    db.create_all()