from DBHelper.MongoHelper import MongoHelper
from Manage.company.Objects.CompanySearch import SearchCompanyApi, SearchCompanyResult


class SearchCompanyApiImpl(SearchCompanyApi):
    """企业查询请求接口实现类"""
    db = MongoHelper("tfse_v0.21")

    def conditions_search(self):
        """"""

        search_body = self.search.make_search_body()
        sort = self.sort.make_sort_body()
        page_size = 10 if self.page_size > 10 else self.page_size
        page_no = self.page_no

        total = self.db.find_all_data_with_count(
            "企业数据",
            "企业数据_更新汇总",
            search_body
        )

        records = self.db.find_data_by_page_with_sort(
            "企业数据",
            "企业数据_更新汇总",
            search_body,
            ["企业ID", "企业名称", "一级行业", "综信评价结果.信用等级", "ESG评价结果.评价等级"],
            sort,
            page_size,
            page_no
        )

        table_data = list()
        if records:
            for record in records:
                company = SearchCompanyResult()
                company.cid = record['企业ID']
                company.name = record['企业名称']
                company.industry = record['一级行业']

                # if record.__contains__('综信评价结果'):
                try:
                    company.credit_level = record['综信评价结果']['信用等级']
                # else:
                except KeyError:
                    company.credit_level = None

                # if record.__contains__('ESG评价结果'):
                try:
                    company.esg_level = record['ESG评价结果']['评价等级']
                # else:
                except KeyError:
                    company.esg_level = None

                data = company.dict_to_show(columns=["企业ID", "企业名称", "一级行业", "信用等级", "ESG等级"])

                table_data.append(data)

        result = {
            "records": table_data,
            "total": total
        }

        return result