from sqlalchemy.orm import Session
from .Mods.CompanyIndexMain.Models import CompanyIndexMain
from .Mods.CompanyBaseInfo.Models import CompanyBaseInfo
from .Mods.CompanyPlusProperty.Models import CompanyPlusProperty
from .Mods.CompanyControl.Models import CompanyControl
from . import Schemas
from Utils.SqlAlchemyUtils import query_common_core


def company_all_info_query(db: Session, query_params: Schemas.CompanyAllInfoQueryParams):
    query = db.query(CompanyIndexMain)
    model_dict = {
        "CompanyIndexMain": CompanyIndexMain,
        "CompanyBaseInfo": CompanyBaseInfo,
        "CompanyPlusProperty": CompanyPlusProperty,
        "CompanyControl": CompanyControl
    }
    params_dict = query_params.dict()
    for key in params_dict:
        if key in model_dict:
            if key == "CompanyIndexMain": continue
            query = query.join(model_dict[key], getattr(model_dict[key], 'company_id') == CompanyIndexMain.company_id)
    for key in params_dict:
        if key in model_dict:
            query = query_common_core(model_dict[key], query, getattr(query_params, key))

    count = query.count()
    page_size = min(query_params.page_size, 100)
    query = query.offset((query_params.page - 1) * page_size).limit(page_size).all()

    company_id_list = [item.company_id for item in query]
    res_dict = {}
    for key, model in model_dict.items():
        res_dict[key] = [item.to_dict() for item in
                         db.query(model).filter(getattr(model, "company_id").in_(company_id_list))]

    return count, res_dict