import os
from contextlib import contextmanager
from typing import Iterable

from sqlalchemy import Column, Integer, String, create_engine, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import declarative_base, sessionmaker

DATABASE_URL = os.getenv("DATABASE_URL", "mysql+pymysql://ec2_mt5:8FmzXj4xcz3AiH2R@163.123.183.106:3306/ip_ops")

engine = create_engine(DATABASE_URL, pool_pre_ping=True)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
Base = declarative_base()


class IPOperation(Base):
    __tablename__ = "ip_operations"

    id = Column(Integer, primary_key=True, autoincrement=True)
    ip_address = Column(String(64), unique=True, nullable=False, index=True)
    note = Column(String(255), nullable=True)


def init_db() -> None:
    Base.metadata.create_all(bind=engine)


@contextmanager
def db_session():
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except SQLAlchemyError:
        session.rollback()
        raise
    finally:
        session.close()


def load_disallowed_ips() -> set[str]:
    with db_session() as session:
        rows: Iterable[IPOperation] = session.scalars(select(IPOperation.ip_address))
        return {row for row in rows}