from io import BytesIO

from docx import Document
from docx.document import Document as DocumentObject
from docx.oxml import OxmlElement
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table, _Cell
from docx.text.paragraph import Paragraph


def _iter_block_items(parent: DocumentObject | _Cell):
    parent_elm = parent.element.body if isinstance(parent, DocumentObject) else parent._tc
    for child in parent_elm.iterchildren():
        if isinstance(child, CT_P):
            yield Paragraph(child, parent)
        elif isinstance(child, CT_Tbl):
            yield Table(child, parent)


def _is_heading(paragraph: Paragraph) -> bool:
    style_name = paragraph.style.name if paragraph.style is not None else ""
    normalized = style_name.lower().replace(" ", "")
    return normalized.startswith("heading")


def _delete_block(block):
    element = block._element
    parent = element.getparent()
    if parent is not None:
        parent.remove(element)


def _clear_paragraph(paragraph: Paragraph):
    element = paragraph._element
    for child in list(element):
        if child.tag.endswith("}r"):
            element.remove(child)


def _append_paragraph_after(paragraph: Paragraph, text: str, style_name: str | None = None) -> Paragraph:
    new_p = OxmlElement("w:p")
    paragraph._element.addnext(new_p)
    new_para = Paragraph(new_p, paragraph._parent)
    if style_name:
        try:
            new_para.style = style_name
        except Exception:
            pass
    if text:
        new_para.add_run(text)
    return new_para


def _append_table_after(paragraph: Paragraph, rows: list[list[str]], headers: list[str] | None = None):
    container = paragraph._parent
    table = container.add_table(rows=1, cols=max(len(headers or []), len(rows[0]) if rows else 1))
    if headers:
        header_cells = table.rows[0].cells
        for index, value in enumerate(headers):
            header_cells[index].text = value
    else:
        if rows:
            first = rows.pop(0)
            for index, value in enumerate(first):
                table.rows[0].cells[index].text = value
    for row in rows:
        new_row = table.add_row().cells
        for index, value in enumerate(row):
            new_row[index].text = value

    tbl = table._tbl
    tbl.getparent().remove(tbl)
    paragraph._element.addnext(tbl)
    return Table(tbl, container)


def _append_empty_paragraph_after_table(table: Table, style_name: str | None = None) -> Paragraph:
    new_p = OxmlElement("w:p")
    table._tbl.addnext(new_p)
    new_para = Paragraph(new_p, table._parent)
    if style_name:
        try:
            new_para.style = style_name
        except Exception:
            pass
    return new_para


def _find_heading_paragraph(document: DocumentObject, heading_text: str) -> Paragraph | None:
    for block in _iter_block_items(document):
        if isinstance(block, Paragraph) and _is_heading(block) and block.text.strip() == heading_text.strip():
            return block
    return None


def _replace_section_content(document: DocumentObject, heading_title: str, content: dict):
    heading = _find_heading_paragraph(document, heading_title)
    if heading is None:
        return

    first_body_style = None
    current = heading._element.getnext()
    blocks_to_remove = []
    while current is not None:
        if isinstance(current, CT_P):
            current_paragraph = Paragraph(current, heading._parent)
            if _is_heading(current_paragraph):
                break
            if first_body_style is None and current_paragraph.style is not None:
                first_body_style = current_paragraph.style.name
            blocks_to_remove.append(current_paragraph)
        elif isinstance(current, CT_Tbl):
            blocks_to_remove.append(Table(current, heading._parent))
        current = current.getnext()

    for block in blocks_to_remove:
        _delete_block(block)

    insert_after = heading
    content_blocks = content.get("content", [])
    for block in content_blocks:
        block_type = block.get("type")
        if block_type == "table":
            rows = [list(row) for row in block.get("rows", [])]
            headers = block.get("headers") or []
            table = _append_table_after(insert_after, rows, headers)
            insert_after = _append_empty_paragraph_after_table(table, first_body_style)
        else:
            text = block.get("text", "")
            insert_after = _append_paragraph_after(insert_after, text, first_body_style)


def export_document_bytes(template_bytes: bytes, logs: list[dict]) -> bytes:
    document = Document(BytesIO(template_bytes))
    for item in logs:
        _replace_section_content(document, item["title"], item["content"])

    output = BytesIO()
    document.save(output)
    return output.getvalue()