"""VendorDmsMasterRestriction model."""

from typing import List

from abacus_common_logic.connectors.database import db
from abacus_common_logic.models.base import CRUDMixin
from sqlalchemy.sql import case


class VendorDmsMasterRestriction(db.Model, CRUDMixin):
    """VendorDmsMasterRestriction model.

    This model maps `vendor_dms_master_restriction` table from art_relations.
    """

    __tablename__ = 'vendor_dms_master_restriction'
    restriction_id = db.Column(db.BigInteger, primary_key=True)
    customer_master_master_id = db.Column(
        db.SmallInteger,
        db.ForeignKey('customer_master_master.customer_master_master_id'),
        nullable=False,
    )
    distribution_type_id = db.Column(
        db.SmallInteger, db.ForeignKey('distribution_type.id'), nullable=False
    )
    vendor_contract_id = db.Column(
        db.Integer, db.ForeignKey('vendor_contract.id'), nullable=False
    )

    @classmethod
    def get_by_criteria(
        cls,
        distribution_type_id: int,
        vendor_contract_id: int,
    ) -> List['VendorDmsMasterRestriction']:
        """Get vendor_dms_master_restriction's by distribution_type_id and vendor_contract_id.

        Args:
            distribution_type_id (int): Id of distribution_type
            vendor_contract_id (int): Id of vendor_contract
        """  # noqa: E501
        return cls.query.filter(
            cls.distribution_type_id == distribution_type_id,
            cls.vendor_contract_id == vendor_contract_id,
        ).all()

    @classmethod
    def delete_by_ids(cls, restriction_ids: List[str]):
        """Delete VendorDmsMasterRestriction records by restriction_id's.

        Args:
            restriction_ids (List[str]): List of vendor_dms_master_restriction ids
        """
        db.session.query(VendorDmsMasterRestriction).filter(
            VendorDmsMasterRestriction.restriction_id.in_(restriction_ids)
        ).delete(synchronize_session='fetch')
        db.session.commit()

    @classmethod
    def update_by_ids(cls, payload: dict):
        """Update VendorDmsMasterRestriction records by restriction_id's.

        Args:
            payload (dict): dict of restriction ids that needs to be updated
            (payload is in format of {restriction_id: new_customer_master_master_id})
        """
        db.session.query(VendorDmsMasterRestriction).filter(
            VendorDmsMasterRestriction.restriction_id.in_(payload)
        ).update(
            {
                VendorDmsMasterRestriction.customer_master_master_id: case(
                    payload, value=VendorDmsMasterRestriction.restriction_id
                )
            },
            synchronize_session=False,
        )
        db.session.commit()