"""Content injection for MCP documents.

Follows the same patterns as WordDocumentEmitter._add_sample_content().
Uses XML templates from xml-structures/ as single source of truth.
"""

from __future__ import annotations

import zipfile
from io import BytesIO
from typing import Any, Dict, List, Literal, Optional

from lxml import etree

from tokenmoulds.emitters import TemplateResolver
from tokenmoulds.mcp.tools.token_utils import escape_xml as _escape_xml

# Word XML namespace
NSMAP = {"w": "http://schemas.openxmlformats.org/wordprocessingml/2006/main"}
W = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"

ContentBlockType = Literal[
    "heading1",
    "heading2",
    "heading3",
    "heading4",
    "paragraph",
    "bullet_list",
    "numbered_list",
    "table",
    "quote",
]


def inject_content(
    template_bytes: bytes,
    content_blocks: List[Dict[str, Any]],
    style_context: Optional[Dict[str, Any]] = None,
) -> bytes:
    """Inject content blocks into a Word template.

    Follows the same pattern as WordDocumentEmitter._customize_document():
    1. Unzip template
    2. Parse document.xml
    3. Add content elements using XML templates
    4. Convert content type from template to document
    5. Repackage

    Args:
        template_bytes: The .dotx template as bytes
        content_blocks: List of content block dicts from MCP input
        style_context: Style configuration from tokens (table style, etc.)

    Returns:
        Modified document as bytes (.docx)
    """
    if not content_blocks:
        return template_bytes

    # Load templates (format_key "word" maps to wordprocessingml)
    templates = TemplateResolver("word")

    # Default style context
    if style_context is None:
        style_context = {}

    # Read and modify the package
    input_zip = BytesIO(template_bytes)
    output_zip = BytesIO()

    with zipfile.ZipFile(input_zip, "r") as zin:
        with zipfile.ZipFile(
            output_zip, "w", zipfile.ZIP_DEFLATED, compresslevel=9
        ) as zout:
            for item in zin.namelist():
                data = zin.read(item)

                if item == "word/document.xml":
                    data = _inject_into_document(
                        data, content_blocks, templates, style_context
                    )
                elif item == "[Content_Types].xml":
                    # Convert from template to document content type
                    data = _convert_template_to_document_content_type(data)

                zout.writestr(item, data)

    return output_zip.getvalue()


def _convert_template_to_document_content_type(content_types_xml: bytes) -> bytes:
    """Convert Content_Types.xml from template to document format.

    Changes the main document content type from:
        wordprocessingml.template.main+xml
    To:
        wordprocessingml.document.main+xml
    """
    xml_str = content_types_xml.decode("utf-8")
    xml_str = xml_str.replace(
        "wordprocessingml.template.main+xml", "wordprocessingml.document.main+xml"
    )
    return xml_str.encode("utf-8")


def _inject_into_document(
    doc_xml: bytes,
    content_blocks: List[Dict[str, Any]],
    templates: TemplateResolver,
    style_context: Dict[str, Any],
) -> bytes:
    """Inject content blocks into document.xml.

    Follows same pattern as WordDocumentEmitter._add_sample_content().
    """
    root = etree.fromstring(doc_xml)
    body = root.find(".//w:body", NSMAP)

    if body is None:
        return doc_xml

    # Preserve sectPr (section properties) - must stay at end
    sect_pr = body.find("w:sectPr", NSMAP)

    # Remove existing paragraphs (sample content)
    for p in body.findall("w:p", NSMAP):
        body.remove(p)

    # Add content blocks
    for block in content_blocks:
        elements = _render_content_block(block, templates, style_context)
        for el in elements:
            body.append(el)

    # Ensure sectPr is at the end
    if sect_pr is not None:
        body.remove(sect_pr)
        body.append(sect_pr)

    return etree.tostring(root, xml_declaration=True, encoding="UTF-8")


def _render_content_block(
    block: Dict[str, Any],
    templates: TemplateResolver,
    style_context: Dict[str, Any],
) -> List[etree._Element]:
    """Render a content block to XML elements.

    Uses XML templates from xml-structures/ following emitter pattern.
    """
    block_type = block.get("type", "paragraph")
    text = _escape_xml(block.get("text", ""))
    items = block.get("items", [])
    rows = block.get("rows", [])

    if block_type in ("heading1", "heading2", "heading3", "heading4"):
        # Use heading template with style ID
        level = block_type[-1]  # "1", "2", etc.
        return [_render_heading(templates, f"Heading{level}", text)]

    elif block_type == "paragraph":
        return [_render_paragraph(templates, text)]

    elif block_type == "quote":
        # Quote uses heading template with Quote style
        return [_render_heading(templates, "Quote", text)]

    elif block_type == "bullet_list":
        return _render_list(templates, items, numbered=False)

    elif block_type == "numbered_list":
        return _render_list(templates, items, numbered=True)

    elif block_type == "table":
        table_style = style_context.get("table_style", "LightGrid-Accent1")
        return [_render_table(templates, rows, table_style)]

    return []


def _render_heading(
    templates: TemplateResolver, style_id: str, text: str
) -> etree._Element:
    """Render a heading/styled paragraph using template."""
    xml = templates.read_text("paragraph_heading.xml")
    xml = xml.replace("${STYLE_ID}", style_id)
    xml = xml.replace("${TEXT}", text)
    return etree.fromstring(xml.encode("utf-8"))


def _render_paragraph(templates: TemplateResolver, text: str) -> etree._Element:
    """Render a body paragraph using template."""
    xml = templates.read_text("paragraph_body.xml")
    xml = xml.replace("${TEXT}", text)
    return etree.fromstring(xml.encode("utf-8"))


def _render_list(
    templates: TemplateResolver,
    items: List[str],
    numbered: bool = False,
) -> List[etree._Element]:
    """Render list items using template."""
    elements = []
    # numId 1 = bullets, numId 2 = numbered (standard Word convention)
    num_id = "2" if numbered else "1"

    for item in items:
        xml = templates.read_text("paragraph_list_item.xml")
        xml = xml.replace("${LEVEL}", "0")
        xml = xml.replace("${NUM_ID}", num_id)
        xml = xml.replace("${TEXT}", _escape_xml(item))
        elements.append(etree.fromstring(xml.encode("utf-8")))

    return elements


def _render_table(
    templates: TemplateResolver,
    rows: List[List[str]],
    table_style: str,
) -> etree._Element:
    """Render a table using templates.

    Args:
        templates: Template resolver for XML templates
        rows: Table rows (first row is header)
        table_style: Word table style ID from tokens (e.g., "LightGrid-Accent1")
    """
    if not rows:
        # Return empty paragraph if no rows
        return _render_paragraph(templates, "")

    # Build rows
    row_elements = []
    for row in rows:
        cells_xml = ""
        for cell_text in row:
            cell_xml = templates.read_text("table_cell.xml")
            cell_xml = cell_xml.replace("${TEXT}", _escape_xml(cell_text))
            cells_xml += cell_xml

        row_xml = templates.read_text("table_row.xml")
        row_xml = row_xml.replace("${CELLS}", cells_xml)
        row_elements.append(row_xml)

    # Build table with style from tokens
    table_xml = templates.read_text("table.xml")
    table_xml = table_xml.replace("${TABLE_STYLE}", table_style)
    table_xml = table_xml.replace("${ROWS}", "".join(row_elements))

    return etree.fromstring(table_xml.encode("utf-8"))
