Scheduler/logic/warehouse/views.py

from fastapi import APIRouter
from pydantic import BaseModel
from typing import Optional, List, Any
from starlette.requests import Request
from tortoise.queryset import QuerySet, Prefetch
import traceback
from models.warehouse import (
    Vacancy,
    Connection, 
    Schelve, 
    Warehouse,
    Vacancy_Pydantic,
    Vacancy_Pydantic_List,
    Connection_Pydantic,
    Schelve_Pydantic,
    Warehouse_Pydantic
    )

from helper import respond_to
from conf import setting
from tortoise.transactions import in_transaction

warehouse_router = APIRouter(prefix='/warehouse')


class VacancyM(BaseModel):
    number: str
    traynum: Optional[str] = None
    memo: Optional[Any] = None
    orderids: Optional[List[str]] = []

class ConnectionM(BaseModel):
    name: str
    vacancies: list[VacancyM]
    memo: Optional[Any] = None

class SchelveM(BaseModel):
    name: str
    vacancies: list[VacancyM]
    memo: Optional[Any] = None
    
class WarehouseM(BaseModel):
    name: str
    schelves: list[SchelveM]
    memo: Optional[Any] = None
    
class WarehouseMClear(BaseModel):
    name: str
    schelvenames: Optional[list[str]] = []

class ConnectionM(BaseModel):
    name: str
    vacancies: list[VacancyM]
    memo: Optional[Any] = None
    
class ConnectionMClear(BaseModel):
    name: str

@warehouse_router.get("/")
def read_root():
    return {"Hello": "World"}



@warehouse_router.post('/addcon', summary='仓库添加接驳台库存')
async def createcon(request: Request, connectionM: ConnectionM):

    # print("orderm: ", orderm)
    print("connectionM: ", connectionM.name, connectionM.memo)

    for vacancy in connectionM.vacancies:
        print("vacancy ", vacancy.number)
    
    async with in_transaction() as conn:
        try:
            connection = None
            qs1 = await Connection.filter(name = connectionM.name, is_active = True, is_valid = True)
            print("qs1: ", qs1, connectionM.name, connectionM.memo)
            if len(qs1) > 0:
                connection = qs1[0]
            else:
                connection = await Connection.create(
                    name = connectionM.name,
                    memo = connectionM.memo,
                    connection=conn
                )

            for vacancy in connectionM.vacancies:
                vacancyObj = None
                qs3 = await Vacancy.filter(number = vacancy.number, is_active = True, is_valid = True)
                if len(qs3) > 0:
                    vacancyObj = qs3[0]
                else:
                    vacancyObj = await Vacancy.create(
                        number = vacancy.number,
                        traynum = vacancy.traynum,
                        memo = vacancy.memo,
                        is_connect = True,
                        connection=conn
                    )

                #add to connection
                await connection.vacancies.add(vacancyObj)

        except Exception as e:
            # printing stack trace
            traceback.print_exc()
            msg = "Something wrong when creating connection: {}".format(e)
            print(msg) 
            return respond_to(code=404, desc="失败", data={"state": msg})
    
    return respond_to(data={"state": "success"})

@warehouse_router.post('/clearcon', summary='仓库清空接驳台库存')
async def clearcon(request: Request, connectionMClear: ConnectionMClear):

    # print("orderm: ", orderm)
    print("connectionMClear: ", connectionMClear.name)

    try:
        connection = None
        qs1 = await Connection.filter(name = connectionMClear.name, is_active = True, is_valid = True)
        print("qs1: ", qs1)
        if len(qs1) > 0:
            connection = qs1[0]
        else:
            raise Exception("Cannot find connection name {}".format(connectionMClear.name))

        await connection.fetch_related("vacancies")
        for vacancy in connection.vacancies:
            vacancy.is_valid = False
            vacancy.is_active = False
            await vacancy.save()

    except Exception as e:
        # printing stack trace
        traceback.print_exc()
        msg = "Something wrong when creating connection: {}".format(e)
        print(msg) 
        return respond_to(code=404, desc="失败", data={"state": msg})

    return respond_to(data={"state": "success"})


@warehouse_router.post('/add', summary='仓库添加仓库暂存货架库存')
async def create(request: Request, warehouseM: WarehouseM):

    # print("orderm: ", orderm)
    # print("warehouseM Add: ", warehouseM.name, warehouseM.memo)

    # for schelve in warehouseM.schelves:
    #     print("schelve ", schelve.name)
    #     for vacancy in schelve.vacancies:
    #         print("vacancy ", vacancy.number)
    
    async with in_transaction() as conn:
        try:
            warehouse = None
            qs1 = await Warehouse.filter(name = warehouseM.name, is_active = True, is_valid = True)
            print("qs1: ", qs1, warehouseM.name, warehouseM.memo)
            if len(qs1) > 0:
                warehouse = qs1[0]
            else:
                warehouse = await Warehouse.create(
                    name = warehouseM.name,
                    memo = warehouseM.memo,
                    connection=conn
                )

            for schelve in warehouseM.schelves:
                schelveObj = None
                qs2 = await Schelve.filter(name = schelve.name, is_active = True, is_valid = True)
                if len(qs2) > 0:
                    schelveObj = qs2[0]
                else:
                    schelveObj = await Schelve.create(
                        name = schelve.name,
                        memo = schelve.memo,
                        connection=conn
                    )

                for vacancy in schelve.vacancies:
                    vacancyObj = None
                    qs3 = await Vacancy.filter(number = vacancy.number, is_active = True, is_valid = True)
                    if len(qs3) > 0:
                        vacancyObj = qs3[0]
                    else:
                        vacancyObj = await Vacancy.create(
                            number = vacancy.number,
                            traynum = vacancy.traynum,
                            memo = vacancy.memo,
                            is_shelf = True,
                            connection=conn
                        )

                    
                    # print("schelveObj.vacancies: ", schelveObj)
                    #Add to schelf
                    await schelveObj.vacancies.add(vacancyObj)
                #add to warehouse
                await warehouse.schelves.add(schelveObj)

        except Exception as e:
            # printing stack trace
            traceback.print_exc()
            msg = "Something wrong when creating connection: {}".format(e)
            print(msg) 
            return respond_to(code=404, desc="失败", data={"state": msg})

    return respond_to(data={"state": "success"})



@warehouse_router.post('/clear', summary='仓库清空仓库暂存货架库存')
async def clear(request: Request, WarehouseMClear: WarehouseMClear):

    # print("orderm: ", orderm)
    print("warehouseM Clear: ", WarehouseMClear.name, WarehouseMClear.schelvenames)

    # for schelvename in WarehouseMClear.schelvenames:
    #     print("schelve ", schelvename)
    
    try:
        warehouse = None
        qs1 = await Warehouse.filter(name = WarehouseMClear.name, is_valid = True, is_active = True)
        # print("qs1: ", qs1, WarehouseMClear.name)
        if len(qs1) > 0:
            warehouse = qs1[0]
        else:
            raise Exception("Cannot find warehouse name {}".format(WarehouseMClear.name))

        if len(WarehouseMClear.schelvenames) > 0:
            for schelvename in WarehouseMClear.schelvenames:
                qs2 = await Schelve.filter(name = schelvename, is_valid = True, is_active = True)
                # print("qs2: ", qs2)
                if len(qs2) > 0:
                    schelveObj = qs2[0]

                    schelveObj.is_active = False
                    schelveObj.is_valid = False
                    await schelveObj.save()

                    await schelveObj.fetch_related("vacancies")
                    for vacancy in schelveObj.vacancies:
                        vacancy.is_active = False
                        vacancy.is_valid = False
                        await vacancy.save()

                else:
                    raise Exception("Cannot find schelve name {}".format(schelvename))

        else:
            # print("Clear All")
            await warehouse.fetch_related("schelves")
            for schelve in warehouse.schelves:

                schelve.is_active = False
                schelve.is_valid = False
                await schelve.save()

                await schelve.fetch_related("vacancies")
                for vacancy in schelve.vacancies:
                    vacancy.is_active = False
                    vacancy.is_valid = False
                    await vacancy.save()

    except Exception as e:
        # printing stack trace
        traceback.print_exc()
        msg = "Something wrong when creating connection: {}".format(e)
        print(msg) 
        return respond_to(code=404, desc="失败", data={"state": msg})
    
    return respond_to(data={"state": "success"})


@warehouse_router.get('/listschelves', summary="列出所有暂存货架库存")
async def index(request: Request, page_no: int = 1, page_size: int = 20):
    """
    列出所有任务
    :param page_no: 1
    :param page_size: 20
    :return:
    """
    offset = (page_no - 1) * page_size

    query = QuerySet(Schelve).filter(is_active=True).prefetch_related(
                                                    Prefetch('vacancies', queryset=Vacancy.filter(is_active = True))
                                                )
    

    count = await query.count()
    schelves = await query.limit(page_size).offset(offset)
    # print("schelves: ", schelves)
    # # print("Dir(schelves): ", schelves)
    # print("schelves[0]: ", dir(schelves[0]), await schelves[0].vacancies.all())
    # vacancies = await schelves[0].vacancies.all()

    resData = []
    for schelve in schelves:
        vacancies = await schelve.vacancies.all()
        schelvej = schelve.as_json()
        schelvej["vacancies"] = vacancies
        resData.append(schelvej)

    #cannot response directly, manytomany field cannot be serilalized
    # return respond_to(code=200, data=dict(count=count, data=vacancies))
    return respond_to(code=200, data=dict(count=count, data=resData))


@warehouse_router.get('/listcon', summary="列出所有接驳区库存")
async def index(request: Request, page_no: int = 1, page_size: int = 20):
    """
    列出所有接驳区库存
    :param page_no: 1
    :param page_size: 20
    :return:
    """
    offset = (page_no - 1) * page_size

    query = QuerySet(Connection).filter(is_active=True).prefetch_related(
                                                    Prefetch('vacancies', queryset=Vacancy.filter(is_active = True))
                                                )
    
    count = await query.count()
    connections = await query.limit(page_size).offset(offset)

    resData = []
    for connection in connections:
        vacancies = await connection.vacancies.all()
        connectionj = connection.as_json()
        connectionj["vacancies"] = vacancies
        resData.append(connectionj)

    return respond_to(code=200, data=dict(count=count, data=resData))


# return things in shelf to warehouse location
@warehouse_router.get("/")
async def return_root():
    qs2 = await Schelve.filter(is_valid = False, is_active = True)
    # print("qs2: ", qs2)
    if len(qs2) > 0:
        schelveObj = qs2[0]

        
    return {"Hello": "World"}