# -*- coding: utf-8 -*-
"""
author: zengbin93
email: zeng_bin8888@163.com
create_dt: 2022/12/16 19:45
describe: 飞书电子表格接口
"""
import string
import pandas as pd
from czsc.fsa.base import FeishuApiBase, request


class SpreadSheets(FeishuApiBase):
    """
    电子表格概述: https://open.feishu.cn/document/ukTMukTMukTM/uATMzUjLwEzM14CMxMTN/overview
    """

    def __init__(self, app_id, app_secret):
        super().__init__(app_id, app_secret)

    def create(self, folder_token, title):
        """创建电子表格

        https://open.feishu.cn/document/ukTMukTMukTM/uUDN04SN0QjL1QDN/sheets-v3/spreadsheet/create

        :param folder_token: 文件夹 token
        :param title: 表格标题，长度范围：0 ～ 255 字符
        :return: 返回数据样例如下
                {
                    "code": 0,
                    "msg": "success",
                    "data": {
                        "spreadsheet": {
                            "title": "title",
                            "folder_token": "fldcnMsNb*****hIW9IjG1LVswg",
                            "url": "https://bytedance.feishu.cn/sheets/shtcnmBA*****yGehy8",
                            "spreadsheet_token": "shtcnmBA*****yGehy8"
                        }
                    }
                }
        """
        url = f"{self.host}/open-apis/sheets/v3/spreadsheets"
        payload = {"title": title, "folder_token": folder_token}
        return request("POST", url, self.get_headers(), payload)

    def check(self, token):
        """获取电子表格信息

        https://open.feishu.cn/document/ukTMukTMukTM/uUDN04SN0QjL1QDN/sheets-v3/spreadsheet/get

        :param token: 表格的token，示例值："shtxxxxxxxxxxxxxxx"
        :return: 返回数据样例如下
                {
                    "code": 0,
                    "msg": "success",
                    "data": {
                        "spreadsheet": {
                            "title": "title",
                            "owner_id": "ou_xxxxxxxxxxxx",
                            "token": "shtxxxxxxxxxxxxxx",
                            "url": "https://bytedance.feishu.cn/sheets/shtcnmBA*****yGehy8"
                        }
                    }
                }
        """
        url = f"{self.host}/open-apis/sheets/v3/spreadsheets/{token}"
        return request('GET', url, self.get_headers())

    def get_sheets(self, token):
        """获取工作表

        https://open.feishu.cn/document/ukTMukTMukTM/uUDN04SN0QjL1QDN/sheets-v3/spreadsheet-sheet/query

        :param token: 电子表格的token
        :return: 返回数据样例如下
                {'code': 0,
                 'data': {'sheets': [{'grid_properties': {'column_count': 20,
                     'frozen_column_count': 0,
                     'frozen_row_count': 0,
                     'row_count': 200},
                    'hidden': False,
                    'index': 0,
                    'resource_type': 'sheet',
                    'sheet_id': '50bafb',
                    'title': 'Sheet1'}]},
                 'msg': ''}
        """
        url = f"{self.host}/open-apis/sheets/v3/spreadsheets/{token}/sheets/query"
        return request("GET", url, self.get_headers())

    def get_sheet_meta(self, token, sheet_id):
        """查询工作表

        https://open.feishu.cn/document/ukTMukTMukTM/uUDN04SN0QjL1QDN/sheets-v3/spreadsheet-sheet/get

        :param token:
        :param sheet_id:
        :return: 返回数据样例
            {'code': 0,
             'data': {'sheet': {'grid_properties': {'column_count': 20,
                'frozen_column_count': 0,
                'frozen_row_count': 0,
                'row_count': 200},
               'hidden': False,
               'index': 0,
               'resource_type': 'sheet',
               'sheet_id': '50bafb',
               'title': 'Sheet1'}},
             'msg': ''}
        """
        url = f"{self.host}/open-apis/sheets/v3/spreadsheets/{token}/sheets/{sheet_id}"
        return request('GET', url, self.get_headers())

    def update_values(self, token, data):
        """向多个范围写入数据

        https://open.feishu.cn/document/ukTMukTMukTM/uEjMzUjLxIzM14SMyMTN

        :param token: spreadsheetToken
        :param data: 数据样例
            {
              "valueRanges": [
                {
                  "range": "range1",
                  "values": [
                    [
                      "string1", 1, "http://www.xx.com"
                    ]
                  ]
                },
                {
                  "range": "range2",
                  "values": [
                    [
                      "string2", 2, "http://www.xx.com"
                    ]
                  ]
                }
              ]
            }
        :return:
        """
        url = f"{self.host}/open-apis/sheets/v2/spreadsheets/{token}/values_batch_update"
        return request("POST", url, self.get_headers(), data)

    def update_styles(self, token, data):
        url = self.host + "/open-apis/sheets/v2/spreadsheets/" + token + "/styles_batch_update"
        return request("POST", url, self.get_headers(), data)

    def read_sheet(self, token, sheet_id):
        """
        https://open.feishu.cn/document/ukTMukTMukTM/ugTMzUjL4EzM14COxMTN

        :return:
        """
        url = f"{self.host}/open-apis/sheets/v2/spreadsheets/{token}/values/{sheet_id}"
        return request("GET", url, self.get_headers())

    def delete_values(self, token, sheet_id):
        """删除行列，清空数据

        https://open.feishu.cn/document/ukTMukTMukTM/ucjMzUjL3IzM14yNyMTN

        :param token:
        :param sheet_id:
        :return:
        """
        url = f"https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{token}/dimension_range"
        row_count = self.get_sheet_meta(token, sheet_id)['data']['sheet']['grid_properties']['row_count'] - 1
        while row_count > 1:
            data = {
                "dimension": {
                    "sheetId": sheet_id,
                    "majorDimension": "ROWS",
                    "startIndex": 1,
                    "endIndex": min(4001, row_count)
                }
            }
            request('DELETE', url, self.get_headers(), data)
            row_count = self.get_sheet_meta(token, sheet_id)['data']['sheet']['grid_properties']['row_count'] - 1

        col_count = self.get_sheet_meta(token, sheet_id)['data']['sheet']['grid_properties']['column_count'] - 1
        if col_count > 1:
            data = {
                "dimension": {
                    "sheetId": sheet_id,
                    "majorDimension": "COLUMNS",
                    "startIndex": 1,
                    "endIndex": min(4001, col_count)
                }
            }
            request('DELETE', url, self.get_headers(), data)

    def dimension_range(self, token, data):
        """增加行列

        https://open.feishu.cn/document/ukTMukTMukTM/uUjMzUjL1IzM14SNyMTN

        :param token:
        :param data:
            {
              "dimension":{
                   "sheetId": "string",
                    "majorDimension": "ROWS",
                    "length": 1
                 }
            }
        :return:
        """
        url = f"https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{token}/dimension_range"
        return request('POST', url, self.get_headers(), data)

    def update_sheets(self, token, operates):
        """增加工作表，复制工作表、删除工作表

        https://open.feishu.cn/document/ukTMukTMukTM/uYTMzUjL2EzM14iNxMTN

        :param token: spreadsheet 的 token
        :param operates: 定义工作表操作
            {
              "requests": [
                {
                  "addSheet": {
                    "properties": {
                      "title": "string",
                      "index": 0
                    }
                  }
                },
                {
                  "copySheet": {
                    "source": {
                      "sheetId": "string"
                    },
                    "destination": {
                      "title": "string"
                    }
                  }
                },
                {
                  "deleteSheet": {
                    "sheetId": "string"
                  }
                }
              ]
            }
        :return:
        """
        url = f"{self.host}/open-apis/sheets/v2/spreadsheets/{token}/sheets_batch_update"
        return request("POST", url, self.get_headers(), operates)

    def add_permissions_member(self, token, doctype, member_type, member_id, perm):
        url = self.host + "/open-apis/drive/v1/permissions/" + token + "/members?type=" \
              + doctype + "&need_notification=false"
        payload = {
            "member_type": member_type,
            "member_id": member_id,
            "perm": perm
        }
        request("POST", url, self.get_headers(), payload)

    # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    # 以下是便捷使用的封装，非官方API接口
    # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    def append(self, token, sheet_id, df: pd.DataFrame, batch_size=2000):
        """往 sheet 中追加数据

        :param token:
        :param sheet_id:
        :param df:
        :param batch_size: 批次写入行数
        :return:
        """
        cols = df.columns.tolist()
        col_range = f"{sheet_id}!A1:{string.ascii_uppercase[len(cols) - 1]}1"
        self.update_values(token, {'valueRanges': [{"range": col_range, "values": [cols]}]})
        meta = self.get_sheet_meta(token, sheet_id)
        start_index = meta['data']['sheet']['grid_properties']['row_count']

        for i in range(0, len(df), batch_size):
            dfi = df.iloc[i: i + batch_size]
            si = i + start_index + 1
            ei = si + batch_size
            vol_range = f"{sheet_id}!A{si}:{string.ascii_uppercase[len(cols) - 1]}{ei}"
            self.update_values(token, {'valueRanges': [{"range": vol_range, "values": dfi.values.tolist()}]})

    def read_table(self, token, sheet_id):
        """读取表格

        :param token:
        :param sheet_id:
        :return:
        """
        res = self.read_sheet(token, sheet_id)
        values = res['data']['valueRange']['values']
        cols = values.pop(0)
        return pd.DataFrame(values, columns=cols)
