from flask import Response

from DBHelper.MongoHelperInstance import DB_GUA
from Utils.ObjUtil import SpecObject
from Utils.ValidateUtil import ValidateAttr


class TFSEFile(SpecObject):
    """文件类"""

    file_id = ValidateAttr(field='file_id', type=str)
    file_bucket = ValidateAttr(field='file_bucket', in_list=["报告PDF"])
    file_body = ValidateAttr(field='file_body', type=Response, default=None)

    fields_map = {
        "file_id": "文件ID",
        "file_bucket": "文件桶",
        "file_body": "文件体"
    }


class FileImpl(TFSEFile):
    """文件类实现"""

    def dict_to_return(self, **kwargs):
        return self.file_body

    def get_image(self):
        """"""

    def get_pdf(self):

        file_stream = DB_GUA.find_file(
            "评级数据",
            self.file_bucket,
            self.file_id
        )

        if file_stream:
            self.file_body = Response(file_stream, content_type='application/pdf')
        else:
            self.file_body = None

    def get_company_excel(self):

        file_stream = DB_GUA.find_file(
            "企业数据",
            '填报模板',
            '62aae021a35c00006e0031a2'
        )

        if file_stream:
            self.file_body = Response(file_stream, content_type='application/pdf')
        else:
            self.file_body = None