from sqlalchemy import Connection, func
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapper

import database
from database import db


class BalanceTransaction(db.Model):
    __tablename__ = 'balance_transactions'

    id = db.Column(db.Integer, primary_key=True, comment='ID транзакции')
    type = db.Column(db.Integer, nullable=False, comment='Тип транзакции')
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    user = db.relationship('User', backref='transactions')
    amount = db.Column(db.Double, nullable=False, comment='Сумма транзакции')
    description = db.Column(db.Text, nullable=False, comment='Описание')
    json_data = db.Column(JSONB, comment='JSON данные')
    created_at = db.Column(db.DateTime, nullable=False, comment='Дата и время создания транзакции')


def recalculate_user_balance(user_id: int):
    user_balance = (database.BalanceTransaction.query.filter_by(user_id=user_id)
                    .group_by(database.BalanceTransaction.user_id)
                    .with_entities(func.sum(database.BalanceTransaction.amount))
                    .scalar()) or 0

    (db.session.query(database.User)
     .filter_by(id=user_id)
     .update({'balance': user_balance}, synchronize_session='fetch'))


@db.event.listens_for(BalanceTransaction, 'after_insert')
def after_balance_transaction_insert(mapper: Mapper, connection: Connection, target: BalanceTransaction):
    recalculate_user_balance(target.user_id)


@db.event.listens_for(BalanceTransaction, 'after_delete')
def after_balance_transaction_delete(mapper: Mapper, connection: Connection, target: BalanceTransaction):
    recalculate_user_balance(target.user_id)