"""CollectionSociety Model.

This CollectionSociety model is used to get information
about collection societies
"""

from oto import response
from sqlalchemy import Boolean
from sqlalchemy import Column
from sqlalchemy.dialects.mysql import INTEGER

from prs.connectors import mysql


class CollectionSociety(mysql.ppb_model):
    """Table definition for collection society table."""

    __tablename__ = 'collection_society'
    society_id = Column(
        'id', INTEGER(unsigned=True), primary_key=True,
        autoincrement=True, nullable=False)
    country_id = Column(INTEGER(unsigned=True), nullable=False)
    is_society = Column(Boolean, nullable=False)


@mysql.wrap_db_errors
def get_society_id_by_country_id(country_id):
    """Get society id by country id.

    Args:
        country_id (int): Country id for which society id is required.

    Returns:
        response.Response: Response with society id or error
    """
    with mysql.ppb_db_session() as session:
        society = session.query(CollectionSociety.society_id).filter(
            CollectionSociety.country_id == country_id,
            CollectionSociety.is_society == 1).all()
        if not society:
            return response.create_not_found_response(
                'Society for country id:{} not found.'.format(country_id))
        return response.Response(message={'society_id': society[0].society_id})