from datetime import datetime, timedelta
import os
import io
import sys
from typing import IO, Any, List, Optional, Union, Dict
import lark_oapi as lark
from lark_oapi.api.drive.v1 import UploadAllMediaRequest, UploadAllMediaRequestBody, UploadAllMediaResponse
from lark_oapi.api.docx.v1 import Block, Callout, Image, PatchDocumentBlockRequest, UpdateBlockRequest, ReplaceImageRequest, PatchDocumentBlockResponse
from matplotlib.axes import Axes
from md_block import create_markdown_block
from lark import lark_client, LarkBlockType
from nest_block import NestBlock
import pandas as pd
import matplotlib.ticker as ticker

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from render import Render

def create_pin_block(md_str: str, children: Optional[Union[NestBlock, List[NestBlock]]]) -> Union[NestBlock, List[NestBlock]]:
    return NestBlock(
        Block.builder()
            .block_type(LarkBlockType.Callout.value)
            .callout(Callout.builder()
                .background_color(2)
                .border_color(2)
                .emoji_id('pushpin')
                .build())
            .build(),
        children
    )

def create_tips_block(md_str: str) -> NestBlock:
    return NestBlock( \
        Block.builder() \
            .block_type(LarkBlockType.Callout.value) \
            .callout(Callout.builder() \
                .background_color(4) \
                .border_color(4) \
                .emoji_id('bulb')
                .build()) \
            .build()
    )

class ImageBlock(NestBlock):
    def __init__(self, img: io.BytesIO):
        super().__init__(Block.builder()
            .block_type(LarkBlockType.Image.value)
            .image(Image.builder().token('').build())
            .build())
        self.img = img

    def callback(self):
        # upload img
        uploadRequest: UploadAllMediaRequest = UploadAllMediaRequest.builder() \
            .request_body(UploadAllMediaRequestBody.builder()
                .file_name('img')
                .parent_type('docx_image')
                .parent_node(str(self.id))
                .size(len(self.img.getvalue()))
                .file(self.img)
                .build()) \
            .build()

        uploadResponse: UploadAllMediaResponse = lark_client.drive.v1.media.upload_all(uploadRequest)

        if not uploadResponse.success():
            lark.logger.error(
                f"client.drive.v1.media.upload_all failed, code: {uploadResponse.code}, msg: {uploadResponse.msg}, log_id: {uploadResponse.get_log_id()}")
            return

        lark.logger.info(lark.JSON.marshal(uploadResponse.data, indent=4))

        # update img block
        updateRequest: PatchDocumentBlockRequest = PatchDocumentBlockRequest.builder() \
            .document_id(str(self.report_id)) \
            .block_id(str(self.id)) \
            .document_revision_id(-1) \
            .request_body(UpdateBlockRequest.builder()
                .replace_image(ReplaceImageRequest.builder()
                    .token(uploadResponse.data.file_token)
                    .build())
                .build()) \
            .build()

        updateResponse: PatchDocumentBlockResponse = lark_client.docx.v1.document_block.patch(updateRequest)

        if not updateResponse.success():
            lark.logger.error(
                f"client.docx.v1.document_block.patch failed, code: {updateResponse.code}, msg: {updateResponse.msg}, log_id: {updateResponse.get_log_id()}")
            return

        # 处理业务结果
        lark.logger.info(lark.JSON.marshal(updateResponse.data, indent=4))

def create_raw_figure_block(plot: Axes):
    f = io.BytesIO()
    fig = plot.get_figure()
    fig.savefig(f, dpi=600)
    return ImageBlock(f)

begin_query_date = (datetime.now() - timedelta(days=17)).strftime('%Y-%m-%d')
end_query_date = (datetime.now() - timedelta(days=4)).strftime('%Y-%m-%d')

def create_chart_block(sql_id: str, ext: Dict[str, str], render: Render):
    with render.context(begin_date=begin_query_date, end_date=end_query_date):
        type = ext.get('type') or 'line'
        x = ext.get('x')
        y = ext.get('y')

        df = render.dataframe(sql_id)
        df[y] = df[y].replace('[\$,]', '', regex=True).apply(pd.to_numeric)

        p = df.plot(kind=type, x=x, y=y, figsize=(16,9))
        p.set_ylim(0, 1.1 * max(df[y]))

        for k, v in df.iterrows():
            p.annotate(v[y], (v[x], v[y]))
        # p.set_xticklabels([x.strftime("%m-%d") for x in df.index], rotation=45)
        # df.apply(lambda row: p.annotate(row.name, row.values, xytext=(10,-5)), axis=1)

        return create_raw_figure_block(p)

def create_langlink_block(app_id: str, input: str) -> NestBlock:
    return NestBlock(Block.builder().build())

def create_slot_block(slot_id: str, slots: Dict[str, Union[NestBlock, List[NestBlock]]]) -> Union[NestBlock, List[NestBlock], None]:
    return slots.get(slot_id)

def create_block_by_category(cate: Dict[str, str], render: Render, slots: Dict[str, Union[NestBlock, List[NestBlock]]] = {}) -> Union[NestBlock, List[NestBlock], None]:
    t = cate.get('type')
    content = cate.get('content') or ''
    ext = cate.get('ext') or {}
    children = cate.get('children')
    childrenBlocks = None

    if children != None:
        childrenBlocks = create_block_by_category(children, render, slots)

    if t == 'pin':
        return create_pin_block(content, childrenBlocks)
    elif t == 'markdown':
        return create_markdown_block(content)
    elif t == 'tips':
        return create_tips_block(content)
    elif t == 'chart':
        return create_chart_block(content, ext, render)
    elif t == 'slot':
        return create_slot_block(content, slots)
    else:
        return None