import asyncio
from collections import defaultdict
from enum import unique, IntEnum
from typing import List

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import joinedload

from database import Marketplace, MarketplaceProduct, Warehouse, Company
from schemas.general import StockUpdate
from updaters.factory import UpdaterFactory


class StocksUpdater:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_marketplace(self, marketplace_id: int):
        marketplace = await self.session.get(Marketplace, marketplace_id, options=[
            joinedload(Marketplace.warehouses).joinedload(Warehouse.suppliers),
            joinedload(Marketplace.warehouses).joinedload(Warehouse.company_warehouses),
            joinedload(Marketplace.company).joinedload(Company.warehouse)
        ])
        return marketplace

    async def update_marketplace(self, marketplace_id: int, updates: List[StockUpdate]):
        marketplace = await self.get_marketplace(marketplace_id)
        updater = UpdaterFactory.get_updater(self.session, marketplace)
        if not updater:
            return
        await updater.update(updates)

    async def update(self, updates: list[StockUpdate]):
        updates_dict = defaultdict(list)
        for update in updates:
            # Working with marketplaces
            stmt = (
                select(
                    MarketplaceProduct.marketplace_id.distinct()
                )
                .where(
                    MarketplaceProduct.product_id == update.product_id,
                    MarketplaceProduct.marketplace_id.in_([9, 41])
                )
            )
            stmt_result = await self.session.execute(stmt)
            marketplace_ids = stmt_result.scalars().all()
            if not marketplace_ids:
                continue
            for marketplace_id in marketplace_ids:
                updates_dict[marketplace_id].append(update)
        updates_list = list(updates_dict.items())
        updates_list = sorted(updates_list, key=lambda x: len(x[1]))

        tasks = []
        for marketplace_id, marketplace_updates in updates_list:
            tasks.append(self.update_marketplace(marketplace_id, marketplace_updates))
        await asyncio.gather(*tasks)