Solo_bot/database/tariffs.py

import hashlib
from collections import defaultdict

from datetime import datetime

from sqlalchemy import delete, func, insert, select, update
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession

from database.models import Server, Tariff
from logger import logger


def create_subgroup_hash(subgroup_title: str, group_code: str) -> str:
    if not subgroup_title:
        return ""

    unique_key = f"{subgroup_title}:{group_code}"
    hash_object = hashlib.md5(unique_key.encode("utf-8"))
    return hash_object.hexdigest()[:8]


async def find_subgroup_by_hash(session: AsyncSession, subgroup_hash: str, group_code: str) -> str | None:
    result = await session.execute(
        select(Tariff.subgroup_title)
        .where(Tariff.group_code == group_code, Tariff.subgroup_title.isnot(None))
        .distinct()
    )
    subgroups = [row[0] for row in result.fetchall()]

    for subgroup_title in subgroups:
        if create_subgroup_hash(subgroup_title, group_code) == subgroup_hash:
            return subgroup_title

    return None


async def get_tariffs(session: AsyncSession, tariff_id: int = None, group_code: str = None, with_subgroup_weights: bool = False):
    try:
        if tariff_id:
            result = await session.execute(select(Tariff).where(Tariff.id == tariff_id))
        elif group_code:
            result = await session.execute(
                select(Tariff)
                .where(Tariff.group_code == group_code)
                .order_by(Tariff.sort_order, Tariff.id)
            )
        else:
            result = await session.execute(select(Tariff).order_by(Tariff.sort_order, Tariff.id))

        tariffs = [dict(r.__dict__) for r in result.scalars().all()]

        if with_subgroup_weights and group_code:
            tariffs_without_order = [t for t in tariffs if t.get("sort_order") is None]
            if tariffs_without_order:
                max_order = max((t.get("sort_order") or 0 for t in tariffs), default=0)
                for i, tariff in enumerate(tariffs_without_order, start=max_order + 1):
                    tariff["sort_order"] = i
                    await session.execute(
                        update(Tariff).where(Tariff.id == tariff["id"]).values(sort_order=i)
                    )
                await session.commit()
            
            grouped = defaultdict(list)
            for t in tariffs:
                grouped[t.get("subgroup_title")].append(t)
            
            subgroup_weights = {}
            for subgroup, tariffs_list in grouped.items():
                if subgroup:
                    total_weight = sum(t.get("sort_order", 1) for t in tariffs_list)
                    subgroup_weights[subgroup] = total_weight
            
            return {
                'tariffs': tariffs,
                'subgroup_weights': subgroup_weights
            }
        
        return tariffs
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при получении тарифов: {e}")
        return []


async def get_tariff_by_id(session: AsyncSession, tariff_id: int):
    try:
        result = await session.execute(select(Tariff).where(Tariff.id == tariff_id))
        tariff = result.scalar_one_or_none()
        return dict(tariff.__dict__) if tariff else None
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при получении тарифа по ID {tariff_id}: {e}")
        return None


async def get_tariffs_for_cluster(session: AsyncSession, cluster_name: str):
    try:
        server_row = await session.execute(
            select(Server.tariff_group).where(Server.cluster_name == cluster_name).limit(1)
        )
        row = server_row.first()

        if not row:
            server_row = await session.execute(
                select(Server.tariff_group).where(Server.server_name == cluster_name).limit(1)
            )
            row = server_row.first()

        if not row or not row[0]:
            return []

        group_code = row[0]
        result = await session.execute(
            select(Tariff)
            .where(Tariff.group_code == group_code, Tariff.is_active.is_(True))
            .order_by(Tariff.sort_order, Tariff.id)
        )
        return [dict(r.__dict__) for r in result.scalars().all()]
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при получении тарифов для кластера {cluster_name}: {e}")
        return []


async def create_tariff(session: AsyncSession, data: dict):
    try:
        data["created_at"] = datetime.utcnow()
        data["updated_at"] = datetime.utcnow()

        if "sort_order" not in data:
            group_code = data.get("group_code")
            if group_code:
                result = await session.execute(
                    select(func.max(Tariff.sort_order))
                    .where(Tariff.group_code == group_code)
                )
                max_order = result.scalar() or 0
            else:
                result = await session.execute(select(func.max(Tariff.sort_order)))
                max_order = result.scalar() or 0

            data["sort_order"] = max_order + 1

        stmt = insert(Tariff).values(**data).returning(Tariff)
        result = await session.execute(stmt)
        await session.commit()
        return result.scalar_one()
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при создании тарифа: {e}")
        await session.rollback()
        return None


async def update_tariff(session: AsyncSession, tariff_id: int, updates: dict):
    if not updates:
        return False
    try:
        updates["updated_at"] = datetime.utcnow()
        await session.execute(update(Tariff).where(Tariff.id == tariff_id).values(**updates))
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при обновлении тарифа ID={tariff_id}: {e}")
        await session.rollback()
        return False


async def delete_tariff(session: AsyncSession, tariff_id: int):
    try:
        await session.execute(delete(Tariff).where(Tariff.id == tariff_id))
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при удалении тарифа ID={tariff_id}: {e}")
        await session.rollback()
        return False


async def check_tariff_exists(session: AsyncSession, tariff_id: int):
    try:
        result = await session.execute(select(Tariff).where(Tariff.id == tariff_id, Tariff.is_active.is_(True)))
        tariff = result.scalar_one_or_none()
        if tariff:
            return True
        logger.warning(f"[TARIFF] Тариф {tariff_id} не найден в БД")
        return False
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при проверке тарифа {tariff_id}: {e}")
        return False


async def get_tariff_sort_order(session: AsyncSession, tariff_id: int) -> int:
    try:
        result = await session.execute(
            select(Tariff.sort_order).where(Tariff.id == tariff_id)
        )
        sort_order = result.scalar_one_or_none()
        
        if sort_order is None:
            await session.execute(
                update(Tariff).where(Tariff.id == tariff_id).values(sort_order=1)
            )
            await session.commit()
            return 1
        
        return sort_order
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при получении sort_order для тарифа {tariff_id}: {e}")
        return None


async def move_tariff_up(session: AsyncSession, tariff_id: int) -> bool:
    try:
        current_order = await get_tariff_sort_order(session, tariff_id)
        new_order = max(1, current_order - 1)
        
        await session.execute(
            update(Tariff).where(Tariff.id == tariff_id).values(sort_order=new_order)
        )
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при перемещении тарифа {tariff_id} вверх: {e}")
        await session.rollback()
        return False


async def move_tariff_down(session: AsyncSession, tariff_id: int) -> bool:
    try:
        current_order = await get_tariff_sort_order(session, tariff_id)
        new_order = current_order + 1
        
        await session.execute(
            update(Tariff).where(Tariff.id == tariff_id).values(sort_order=new_order)
        )
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при перемещении тарифа {tariff_id} вниз: {e}")
        await session.rollback()
        return False


async def initialize_tariff_sort_orders(session: AsyncSession, group_code: str) -> bool:
    try:
        result = await session.execute(
            select(Tariff).where(Tariff.group_code == group_code).order_by(Tariff.id)
        )
        tariffs = result.scalars().all()
        
        if not tariffs:
            return True

        for i, tariff in enumerate(tariffs):
            new_sort_order = 1 + i
            await session.execute(
                update(Tariff).where(Tariff.id == tariff.id).values(sort_order=new_sort_order)
            )
        
        await session.commit()
        return True
    except SQLAlchemyError as e:
        logger.error(f"[TARIFF] Ошибка при инициализации sort_order для группы {group_code}: {e}")
        await session.rollback()
        return False