from contextlib import contextmanager
from sqlalchemy.schema import CreateTable

from db.models import Base

from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.mysql import insert
from sqlalchemy import create_engine, text, inspect
from env_provider import DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD, DB_SCHEMA_UPDATED, ENVIRONMENT
print(f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}")
engine = create_engine(f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}")
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

@contextmanager
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

def clear_db():
    with engine.connect() as conn:
        conn.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
        inspector = inspect(engine)
        for table_name in inspector.get_table_names():
            conn.execute(text(f"DROP TABLE IF EXISTS {table_name}"))
        conn.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))


def create_all():
    with engine.begin() as conn:
        Base.metadata.create_all(bind=conn)

def run_scripts():
    from db.models import table_models
    with get_db() as session:
        for model in table_models:
            if hasattr(model, "__scripts__"):
                scripts = model.__scripts__
                for script in scripts:
                    session.execute(script)
        session.commit()

def init_payload():
    from db.models import table_models
    with get_db() as session:
        session.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
        for model in table_models:
            print(str(CreateTable(model.__table__)))

            print(f"MODEL -- {model}, {hasattr(model, '__payload__')}")
            if hasattr(model, "__payload__"):
                payload = model.__payload__[ENVIRONMENT]
                print(f"- - [ - ] hasattr, {ENVIRONMENT} - {payload}")
                stmt = insert(model.__table__).values(payload).prefix_with("IGNORE")
                print(f"- - [ - ] {stmt}\n")
                session.execute(stmt)
        session.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))
        session.commit()


def setup_db():
    if DB_SCHEMA_UPDATED:
        clear_db()
        print("[ x ] db cleared")
    create_all()
    print("[ x ] db created")
    run_scripts()
    print("[ x ] db scripts executed")
    init_payload()
    print("[ x ] payload loaded")