from typing import Union
from lark_oapi.api.docx.v1 import Block, List, Optional, Text, TextElement, TextRun
from lark import LarkBlockType
from nest_block import NestBlock
from markdown_it import MarkdownIt

class MdBlock:
    def __init__(self, content: str, children: Optional[List["MdBlock"]]) -> None:
        self.content = content
        self.children = children

    def to_nest_block(self) -> Union[NestBlock, None]:
        return None

class TextBlock(MdBlock):
    def __init__(self, content: str, children: Optional[List[MdBlock]]) -> None:
        super().__init__(content, children)

    def to_nest_block(self) -> Union[NestBlock, None]:
        return create_text_block(self.content)

def create_text_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Text.value) \
            .text(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

class HeadingBlock(MdBlock):
    def __init__(self, level: Union[str, int, float], content: str, children: Optional[List[MdBlock]]) -> None:
        super().__init__(content, children)
        self.level = level

    def to_nest_block(self) -> Union[NestBlock, None]:
        return globals().get(f'create_heading{self.level}_block', create_text_block)(self.content)

def create_heading1_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Heading1.value) \
            .heading1(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

def create_heading2_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Heading2.value) \
            .heading2(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

def create_heading3_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Heading3.value) \
            .heading3(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

def create_heading4_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Heading4.value) \
            .heading4(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

def create_heading5_block(content: str) -> NestBlock:
    return NestBlock(
        Block.builder() \
            .block_type(LarkBlockType.Heading5.value) \
            .heading5(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build()
    )

class BulletBlock(MdBlock):
    def __init__(self, content: str, children: Optional[List[MdBlock]]) -> None:
        super().__init__(content, children)

    def to_nest_block(self) -> Union[NestBlock, None]:
        children = convert_to_nest_blocks(self.children or [])
        return create_bullet_block(self.content, children)

def create_bullet_block(content: str, next: Optional[Union[NestBlock, List[NestBlock]]] = None) -> NestBlock:
    return NestBlock( \
        Block.builder() \
            .block_type(LarkBlockType.Bullet.value) \
            .bullet(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build(), \
        next
    )

class OrderBlock(MdBlock):
    def __init__(self, content: str, children: Optional[List[MdBlock]]) -> None:
        super().__init__(content, children)

    def to_nest_block(self) -> Union[NestBlock, None]:
        return create_ordered_block(self.content)

def create_ordered_block(content: str, next: Optional[Union[NestBlock, List[NestBlock]]] = None) -> NestBlock:
    return NestBlock( \
        Block.builder() \
            .block_type(LarkBlockType.Ordered.value) \
            .ordered(Text.builder() \
                .elements([TextElement.builder()
                    .text_run(TextRun.builder()
                        .content(content)
                        .build())
                        .build()
                    ])
                .build()) \
            .build(), \
        next
    )

def create_markdown_block(md: str) -> List[NestBlock]:
    md_parser = MarkdownIt()
    tokens = md_parser.parse(md)
    md_blocks = []
    stack: List[MdBlock] = []
    bullet_parent_stack: List[MdBlock] = []
    bullet_parent_level = 0

    for token in tokens:
        token_type = token.type

        if token_type == 'heading_open':
            level = len(token.markup)
            stack.append(HeadingBlock(level, '', None))

        elif token_type == 'heading_close':
            if stack:
                heading_block = stack.pop()
                md_blocks.append(heading_block)

        elif token_type == 'bullet_list_open':
            bullet_parent_level += 1
            if bullet_parent_level > 1:
                bullet_parent_stack.append(stack[-1])

        elif token_type == 'bullet_list_close':
            if bullet_parent_level > 1:
                bullet_parent_stack.pop()
            bullet_parent_level -= 1

        elif token_type == 'list_item_open':
            b = BulletBlock('', None)
            stack.append(b)
            if bullet_parent_stack:
                if bullet_parent_stack[-1].children == None:
                    bullet_parent_stack[-1].children = []
                bullet_parent_stack[-1].children.append(b)

        elif token_type == 'list_item_close':
            b = stack.pop()
            if bullet_parent_level == 1:
                md_blocks.append(b)

        elif token_type == 'inline':
            c = token.content
            if stack:
                outter_block = stack[-1]
                outter_block.content += c
            else:
                md_blocks.append(TextBlock(c, None))

    return convert_to_nest_blocks(md_blocks)

def convert_to_nest_blocks(md_blocks: List[MdBlock]) -> List[NestBlock]:
    nest_blocks = []
    for md_block in md_blocks:
        nest_block = md_block.to_nest_block()
        nest_blocks.append(nest_block)
    return nest_blocks

# markdown_str = """
# - B1
#   - B2
# - B3
# - B4
# """

# blocks = create_markdown_block(markdown_str)
# for block in blocks:
#     print(block.block)