from __future__ import annotations

import base64
import io
import re

from docx import Document
from docx.enum.style import WD_STYLE_TYPE
from docx.enum.table import WD_CELL_VERTICAL_ALIGNMENT, WD_TABLE_ALIGNMENT
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.shared import Inches, Pt, RGBColor
from PIL import Image

from backend.services.tool_source_artifacts.image_assets import SourceImageAsset

from .document_models import DocumentSpec, TableBlock

CONTENT_WIDTH_INCHES = 6.5
MAX_IMAGE_HEIGHT_INCHES = 7.0
IMAGE_WIDTH_INCHES = {
    "small": 1.625,
    "medium": 3.25,
    "large": 4.875,
    "full": 6.5,
}
FONT_NAME = "Calibri"
BODY_SIZE = 11
BODY_AFTER = 6
BODY_LINE_SPACING = 1.10
TITLE_SIZE = 24
TITLE_COLOR = "0B2545"
SUBTITLE_SIZE = 12
HEADING_STYLES = (
    (16, "2E74B5", 16, 8),
    (13, "2E74B5", 12, 6),
    (12, "1F4D78", 8, 4),
)
LIST_INDENT = 0.5
LIST_HANGING = 0.25
LIST_AFTER = 8
TABLE_STYLE = "Table Grid"
STYLE_TITLE = "Agent Hub Document Title"
STYLE_SUBTITLE = "Agent Hub Document Subtitle"
STYLE_METADATA = "Agent Hub Document Metadata"
STYLE_CALLOUT = "Agent Hub Callout"
STYLE_QUOTE = "Agent Hub Quote"
STYLE_QUOTE_ATTRIBUTION = "Agent Hub Quote Attribution"
STYLE_TABLE_TITLE = "Agent Hub Table Title"
STYLE_TABLE_SOURCE = "Agent Hub Table Source"
STYLE_IMAGE_CAPTION = "Caption"


def render_document_to_bytes(
    spec: DocumentSpec,
    image_assets: dict[str, SourceImageAsset] | None = None,
) -> bytes:
    document = Document()
    _configure_document(document, spec)
    _render_title_block(document, spec)
    resolved_image_assets = image_assets or {}

    for block in spec.blocks:
        if block.type == "heading":
            document.add_paragraph(block.text, style=f"Heading {block.level}")
        elif block.type == "paragraph":
            _render_paragraph(document, block)
        elif block.type == "list":
            _render_list(document, block)
        elif block.type == "table":
            _render_table(document, block)
        elif block.type == "callout":
            _render_callout(document, block)
        elif block.type == "quote":
            _render_quote(document, block)
        elif block.type == "image":
            _render_image(document, block, resolved_image_assets)
        elif block.type == "page_break":
            document.add_page_break()

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


def _configure_document(document, spec: DocumentSpec) -> None:
    section = document.sections[0]
    section.page_width = Inches(8.5)
    section.page_height = Inches(11)
    section.top_margin = Inches(1)
    section.right_margin = Inches(1)
    section.bottom_margin = Inches(1)
    section.left_margin = Inches(1)
    section.header_distance = Inches(0.492)
    section.footer_distance = Inches(0.492)

    document.core_properties.title = spec.document.title
    if spec.document.subject:
        document.core_properties.subject = spec.document.subject
    if spec.document.author:
        document.core_properties.author = spec.document.author

    normal = document.styles["Normal"]
    _set_style_font(normal, BODY_SIZE, "000000")
    normal.paragraph_format.alignment = WD_ALIGN_PARAGRAPH.LEFT
    normal.paragraph_format.space_before = Pt(0)
    normal.paragraph_format.space_after = Pt(BODY_AFTER)
    normal.paragraph_format.line_spacing = BODY_LINE_SPACING
    normal.paragraph_format.widow_control = True

    for level, (size, color, before, after) in enumerate(HEADING_STYLES, start=1):
        style = document.styles[f"Heading {level}"]
        _set_style_font(style, size, color, bold=True)
        style.paragraph_format.space_before = Pt(before)
        style.paragraph_format.space_after = Pt(after)
        style.paragraph_format.keep_with_next = True
        style.paragraph_format.keep_together = True

    _configure_list_styles(document)
    _configure_custom_styles(document)


def _configure_list_styles(document) -> None:
    for base_name in ("List Bullet", "List Number"):
        for level in range(3):
            style_name = base_name if level == 0 else f"{base_name} {level + 1}"
            try:
                style = document.styles[style_name]
            except KeyError:
                continue
            _set_style_font(style, BODY_SIZE, "000000")
            style.paragraph_format.left_indent = Inches(LIST_INDENT + level * 0.25)
            style.paragraph_format.first_line_indent = Inches(-LIST_HANGING)
            style.paragraph_format.space_after = Pt(LIST_AFTER)
            style.paragraph_format.line_spacing = BODY_LINE_SPACING
            style.paragraph_format.widow_control = True


def _configure_custom_styles(document) -> None:
    title = _get_or_add_paragraph_style(document, STYLE_TITLE)
    _set_style_font(title, TITLE_SIZE, TITLE_COLOR, bold=True)
    title.paragraph_format.space_before = Pt(0)
    title.paragraph_format.space_after = Pt(5)
    title.paragraph_format.keep_with_next = True

    subtitle = _get_or_add_paragraph_style(document, STYLE_SUBTITLE)
    _set_style_font(subtitle, SUBTITLE_SIZE, "4F5965")
    subtitle.paragraph_format.space_before = Pt(0)
    subtitle.paragraph_format.space_after = Pt(12)
    subtitle.paragraph_format.keep_with_next = True

    metadata = _get_or_add_paragraph_style(document, STYLE_METADATA)
    _set_style_font(metadata, 9.5, "606A76")
    metadata.paragraph_format.space_before = Pt(0)
    metadata.paragraph_format.space_after = Pt(3)

    callout = _get_or_add_paragraph_style(document, STYLE_CALLOUT)
    _set_style_font(callout, BODY_SIZE, "1F3A5F")
    callout.paragraph_format.left_indent = Inches(0.25)
    callout.paragraph_format.right_indent = Inches(0.1)
    callout.paragraph_format.space_before = Pt(6)
    callout.paragraph_format.space_after = Pt(8)
    callout.paragraph_format.keep_together = True

    quote = _get_or_add_paragraph_style(document, STYLE_QUOTE)
    _set_style_font(quote, 11.5, "3F4750", italic=True)
    quote.paragraph_format.left_indent = Inches(0.4)
    quote.paragraph_format.right_indent = Inches(0.4)
    quote.paragraph_format.space_before = Pt(8)
    quote.paragraph_format.space_after = Pt(3)
    quote.paragraph_format.keep_together = True

    attribution = _get_or_add_paragraph_style(document, STYLE_QUOTE_ATTRIBUTION)
    _set_style_font(attribution, 9.5, "606A76")
    attribution.paragraph_format.left_indent = Inches(0.4)
    attribution.paragraph_format.space_after = Pt(8)

    table_title = _get_or_add_paragraph_style(document, STYLE_TABLE_TITLE)
    _set_style_font(table_title, 10, "1F4D78", bold=True)
    table_title.paragraph_format.space_before = Pt(6)
    table_title.paragraph_format.space_after = Pt(4)
    table_title.paragraph_format.keep_with_next = True

    table_source = _get_or_add_paragraph_style(document, STYLE_TABLE_SOURCE)
    _set_style_font(table_source, 8.5, "606A76", italic=True)
    table_source.paragraph_format.space_before = Pt(4)
    table_source.paragraph_format.space_after = Pt(4)

    image_caption = _get_or_add_paragraph_style(document, STYLE_IMAGE_CAPTION)
    _set_style_font(image_caption, 9, "606A76", italic=True)
    image_caption.paragraph_format.space_before = Pt(2)
    image_caption.paragraph_format.space_after = Pt(8)
    image_caption.paragraph_format.keep_together = True


def _render_title_block(document, spec: DocumentSpec) -> None:
    document.add_paragraph(spec.document.title, style=STYLE_TITLE)
    if spec.document.subtitle:
        document.add_paragraph(spec.document.subtitle, style=STYLE_SUBTITLE)
    metadata_parts = [part for part in (spec.document.author, spec.document.date) if part]
    if metadata_parts:
        paragraph = document.add_paragraph(style=STYLE_METADATA)
        paragraph.add_run(" | ".join(metadata_parts))
        paragraph.paragraph_format.space_after = Pt(14)
    elif not spec.document.subtitle:
        spacer = document.add_paragraph(style=STYLE_METADATA)
        spacer.paragraph_format.space_after = Pt(8)


def _render_paragraph(document, block) -> None:
    paragraph = document.add_paragraph()
    if block.text is not None:
        paragraph.add_run(block.text)
        return
    for run_spec in block.runs:
        run = paragraph.add_run(run_spec.text)
        run.bold = run_spec.bold
        run.italic = run_spec.italic
        run.underline = run_spec.underline


def _render_list(document, block) -> None:
    base_style = "List Number" if block.style == "numbered" else "List Bullet"
    for item in block.items:
        style_name = base_style if item.level == 0 else f"{base_style} {item.level + 1}"
        try:
            document.styles[style_name]
        except KeyError:
            style_name = base_style
        document.add_paragraph(item.text, style=style_name)


def _render_table(document, block: TableBlock) -> None:
    if block.title:
        document.add_paragraph(block.title, style=STYLE_TABLE_TITLE)

    table = document.add_table(rows=1, cols=len(block.columns))
    table.style = TABLE_STYLE
    table.alignment = WD_TABLE_ALIGNMENT.LEFT
    table.autofit = False
    widths = _resolve_column_widths(block)

    header_cells = table.rows[0].cells
    for index, (cell, label) in enumerate(zip(header_cells, block.columns)):
        cell.text = label
        cell.width = widths[index]
        cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
        paragraph = cell.paragraphs[0]
        paragraph.alignment = (
            WD_ALIGN_PARAGRAPH.CENTER if _is_short_value_column(block, index) else WD_ALIGN_PARAGRAPH.LEFT
        )
        _format_cell_paragraph(paragraph, bold=True)

    for values in block.rows:
        cells = table.add_row().cells
        for index, (cell, value) in enumerate(zip(cells, values)):
            cell.text = value
            cell.width = widths[index]
            cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
            paragraph = cell.paragraphs[0]
            paragraph.alignment = (
                WD_ALIGN_PARAGRAPH.CENTER if _is_short_value_column(block, index) else WD_ALIGN_PARAGRAPH.LEFT
            )
            _format_cell_paragraph(paragraph, bold=False)

    for index, column in enumerate(table.columns):
        column.width = widths[index]

    if block.source:
        document.add_paragraph(f"Source: {block.source}", style=STYLE_TABLE_SOURCE)
    else:
        trailing = document.add_paragraph()
        trailing.paragraph_format.space_before = Pt(0)
        trailing.paragraph_format.space_after = Pt(2)


def _render_callout(document, block) -> None:
    paragraph = document.add_paragraph(style=STYLE_CALLOUT)
    if block.label:
        label_run = paragraph.add_run(f"{block.label}: ")
        label_run.bold = True
    paragraph.add_run(block.text)


def _render_quote(document, block) -> None:
    document.add_paragraph(block.text, style=STYLE_QUOTE)
    if block.attribution:
        document.add_paragraph(f"— {block.attribution}", style=STYLE_QUOTE_ATTRIBUTION)


def _render_image(
    document,
    block,
    image_assets: dict[str, SourceImageAsset],
) -> None:
    image_asset = image_assets.get(block.asset_id)
    if image_asset is None:
        raise ValueError(f"Unknown image asset '{block.asset_id}'")

    try:
        image_bytes = base64.b64decode(image_asset.data_base64)
        image_buffer, image_width_px, image_height_px = _prepare_word_picture(image_bytes)
    except Exception as exc:
        raise ValueError(f"Image asset '{image_asset.asset_id}' could not be decoded as a valid image") from exc
    if image_width_px <= 0 or image_height_px <= 0:
        raise ValueError(f"Image asset '{image_asset.asset_id}' has invalid dimensions")

    max_width = IMAGE_WIDTH_INCHES[block.width]
    scale = min(max_width / image_width_px, MAX_IMAGE_HEIGHT_INCHES / image_height_px)
    rendered_width = image_width_px * scale
    rendered_height = image_height_px * scale
    alignment = {
        "left": WD_ALIGN_PARAGRAPH.LEFT,
        "center": WD_ALIGN_PARAGRAPH.CENTER,
        "right": WD_ALIGN_PARAGRAPH.RIGHT,
    }[block.alignment]

    image_paragraph = document.add_paragraph()
    image_paragraph.alignment = alignment
    image_paragraph.paragraph_format.space_before = Pt(6)
    image_paragraph.paragraph_format.space_after = Pt(4 if block.caption else 8)
    image_paragraph.paragraph_format.keep_together = True
    image_paragraph.paragraph_format.keep_with_next = bool(block.caption)
    image_paragraph.add_run().add_picture(
        image_buffer,
        width=Inches(rendered_width),
        height=Inches(rendered_height),
    )

    if block.caption:
        caption = document.add_paragraph(block.caption, style=STYLE_IMAGE_CAPTION)
        caption.alignment = alignment


def _prepare_word_picture(image_bytes: bytes) -> tuple[io.BytesIO, int, int]:
    """Normalize readable image formats that Word cannot embed reliably to PNG."""

    with Image.open(io.BytesIO(image_bytes)) as image:
        image_width_px, image_height_px = image.size
        image_format = (image.format or "").upper()
        if image_format not in {"BMP", "GIF", "JPEG", "PNG", "TIFF", "WMF"}:
            converted_buffer = io.BytesIO()
            converted_image = image.convert("RGBA" if image.mode in {"RGBA", "LA", "P"} else "RGB")
            converted_image.save(converted_buffer, format="PNG")
            converted_buffer.seek(0)
            return converted_buffer, image_width_px, image_height_px

    passthrough_buffer = io.BytesIO(image_bytes)
    passthrough_buffer.seek(0)
    return passthrough_buffer, image_width_px, image_height_px


def _resolve_column_widths(block: TableBlock):
    weights = block.column_widths or _infer_column_weights(block)
    total = sum(weights)
    return [Inches(CONTENT_WIDTH_INCHES * weight / total) for weight in weights]


def _infer_column_weights(block: TableBlock) -> list[float]:
    weights: list[float] = []
    for index, heading in enumerate(block.columns):
        lengths = [len(heading), *[len(row[index]) for row in block.rows]]
        representative = max(8, min(max(lengths), 60))
        weights.append(float(representative))
    return weights


def _is_short_value_column(block: TableBlock, column_index: int) -> bool:
    values = [block.columns[column_index], *[row[column_index] for row in block.rows]]
    if any("\n" in value for value in values):
        return False
    average_length = sum(len(value.strip()) for value in values) / len(values)
    numeric_like = sum(bool(re.fullmatch(r"[\s$€£¥+\-()\d.,%/:]+", value)) for value in values[1:])
    return average_length <= 14 or (bool(values[1:]) and numeric_like * 2 >= len(values[1:]))


def _format_cell_paragraph(paragraph, *, bold: bool) -> None:
    paragraph.paragraph_format.space_before = Pt(2)
    paragraph.paragraph_format.space_after = Pt(2)
    paragraph.paragraph_format.line_spacing = 1.0
    paragraph.paragraph_format.widow_control = True
    for run in paragraph.runs:
        run.font.name = FONT_NAME
        run.font.size = Pt(9.5)
        run.bold = bold


def _get_or_add_paragraph_style(document, name: str):
    try:
        return document.styles[name]
    except KeyError:
        return document.styles.add_style(name, WD_STYLE_TYPE.PARAGRAPH)


def _set_style_font(style, size: float, color: str, *, bold: bool = False, italic: bool = False) -> None:
    style.font.name = FONT_NAME
    style.font.size = Pt(size)
    style.font.color.rgb = RGBColor.from_string(color)
    style.font.bold = bold
    style.font.italic = italic
