from __future__ import annotations

import base64
import io
from typing import List, Optional

import structlog
from docx import Document
from docx.enum.style import WD_STYLE_TYPE
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.shared import Pt, RGBColor

from export.models import (
    IDExportDocument,
    IDExportPageNumber,
    IDExportParagraphStyle,
    IDExportRun,
    IDExportTable,
)
from word.units import rgb

log = structlog.get_logger()

ALIGNMENT_MAP = {
    'left':    WD_ALIGN_PARAGRAPH.LEFT,
    'center':  WD_ALIGN_PARAGRAPH.CENTER,
    'right':   WD_ALIGN_PARAGRAPH.RIGHT,
    'justify': WD_ALIGN_PARAGRAPH.JUSTIFY,
}

# InDesign paragraph style name -> Word built-in style
HEADING_STYLE_MAP = {
    'heading 1':    'Heading 1',
    'h1':           'Heading 1',
    'title':        'Heading 1',
    'heading 2':    'Heading 2',
    'h2':           'Heading 2',
    'subheading':   'Heading 2',
    'heading 3':    'Heading 3',
    'h3':           'Heading 3',
    'subhead':      'Heading 3',
    'caption':      'Caption',
    'table header': 'Table Grid',
}


class StyleFactory:
    """Registers InDesign paragraph styles as Word styles and applies run formatting."""

    @staticmethod
    def register_all(doc: Document, document: IDExportDocument) -> None:
        """Register all paragraph styles once before any content is added."""
        for style_data in document.paragraph_styles:
            StyleFactory._register_paragraph_style(doc, style_data)

    @staticmethod
    def _register_paragraph_style(doc: Document, style_data: IDExportParagraphStyle) -> None:
        builtin = HEADING_STYLE_MAP.get(style_data.name.lower().strip())
        try:
            if builtin:
                style = doc.styles[builtin]
            else:
                try:
                    style = doc.styles[style_data.name]
                except KeyError:
                    style = doc.styles.add_style(style_data.name, WD_STYLE_TYPE.PARAGRAPH)
        except Exception:
            return

        try:
            if style_data.font:
                style.font.name = style_data.font
            if style_data.size_pt:
                style.font.size = Pt(style_data.size_pt)
            if style_data.color_hex:
                style.font.color.rgb = rgb(style_data.color_hex)
            style.font.bold = style_data.bold or None
            style.font.italic = style_data.italic or None
            if style_data.space_before_pt:
                style.paragraph_format.space_before = Pt(style_data.space_before_pt)
            if style_data.space_after_pt:
                style.paragraph_format.space_after = Pt(style_data.space_after_pt)
            if style_data.leading_pt:
                style.paragraph_format.line_spacing = Pt(style_data.leading_pt)
            alignment = ALIGNMENT_MAP.get(style_data.alignment)
            if alignment is not None:
                style.paragraph_format.alignment = alignment
        except Exception as e:
            log.warning('style_property_failed', style=style_data.name, error=str(e))

    @staticmethod
    def apply_paragraph_style(para, style_name: str, doc: Document) -> None:
        """Apply a named style to a paragraph; fall back to Normal if missing."""
        builtin = HEADING_STYLE_MAP.get(style_name.lower().strip())
        target = builtin or style_name
        try:
            para.style = doc.styles[target]
        except KeyError:
            pass

    @staticmethod
    def apply_run_formatting(run, run_data: IDExportRun) -> None:
        """Apply character-level formatting to a Word run."""
        if run_data.bold:
            run.bold = True
        if run_data.italic:
            run.italic = True
        if run_data.font:
            run.font.name = run_data.font
        if run_data.size_pt:
            run.font.size = Pt(run_data.size_pt)
        if run_data.color_hex:
            try:
                run.font.color.rgb = rgb(run_data.color_hex)
            except Exception:
                pass


class ImageFactory:
    """Adds an image to the document from base64-encoded data."""

    @staticmethod
    def add_to_doc(doc: Document, image_data_b64: Optional[str],
                   width_pt: float, height_pt: float) -> None:
        if not image_data_b64:
            return
        try:
            image_bytes = base64.b64decode(image_data_b64)
            para = doc.add_paragraph()
            run = para.add_run()
            run.add_picture(io.BytesIO(image_bytes), width=Pt(width_pt), height=Pt(height_pt))
        except Exception as e:
            log.warning('image_insert_failed', error=str(e))


class TableFactory:
    """Builds a Word table from an IDExportTable."""

    @staticmethod
    def add_to_doc(doc: Document, table_data: IDExportTable) -> None:
        if not table_data.rows:
            return

        row_count = len(table_data.rows)
        col_count = max((len(r.cells) for r in table_data.rows), default=0)
        if col_count == 0:
            return

        table = doc.add_table(rows=row_count, cols=col_count)
        try:
            table.style = 'Table Grid'
        except KeyError:
            pass

        # Set column widths
        for i, col_width in enumerate(table_data.column_widths_pt):
            if i < col_count:
                try:
                    for cell in table.columns[i].cells:
                        cell.width = Pt(col_width)
                except Exception:
                    pass

        # Fill cells
        for row_idx, row_data in enumerate(table_data.rows):
            for col_idx, cell_data in enumerate(row_data.cells):
                if col_idx >= col_count:
                    break
                cell = table.cell(row_idx, col_idx)
                cell.text = cell_data.text

                if cell_data.text_color_hex:
                    for para in cell.paragraphs:
                        for run in para.runs:
                            try:
                                run.font.color.rgb = rgb(cell_data.text_color_hex)
                            except Exception:
                                pass

                if cell_data.background_hex:
                    TableFactory._shade_cell(cell, cell_data.background_hex)

                if row_data.is_header:
                    for para in cell.paragraphs:
                        for run in para.runs:
                            run.bold = True

    @staticmethod
    def _shade_cell(cell, hex_color: str) -> None:
        try:
            tc = cell._tc
            tcPr = tc.get_or_add_tcPr()
            shd = OxmlElement('w:shd')
            shd.set(qn('w:val'), 'clear')
            shd.set(qn('w:color'), 'auto')
            shd.set(qn('w:fill'), hex_color.lstrip('#').upper())
            tcPr.append(shd)
        except Exception:
            pass


class FooterFactory:
    """Adds a native Word PAGE field footer matched to InDesign page number styling."""

    @staticmethod
    def apply(doc: Document, pn: IDExportPageNumber) -> None:
        section = doc.sections[0]
        footer = section.footer
        footer.is_linked_to_previous = False

        para = footer.paragraphs[0] if footer.paragraphs else footer.add_paragraph()
        para.clear()

        alignment = ALIGNMENT_MAP.get(pn.alignment, WD_ALIGN_PARAGRAPH.CENTER)
        para.alignment = alignment

        r1 = OxmlElement('w:r')
        r1.append(_make_rpr(pn))
        fc_begin = OxmlElement('w:fldChar')
        fc_begin.set(qn('w:fldCharType'), 'begin')
        r1.append(fc_begin)
        para._p.append(r1)

        r2 = OxmlElement('w:r')
        instr = OxmlElement('w:instrText')
        instr.text = ' PAGE '
        instr.set('{http://www.w3.org/XML/1998/namespace}space', 'preserve')
        r2.append(instr)
        para._p.append(r2)

        r3 = OxmlElement('w:r')
        r3.append(_make_rpr(pn))
        fc_end = OxmlElement('w:fldChar')
        fc_end.set(qn('w:fldCharType'), 'end')
        r3.append(fc_end)
        para._p.append(r3)


def _make_rpr(pn: IDExportPageNumber):
    rPr = OxmlElement('w:rPr')
    if pn.font:
        rFonts = OxmlElement('w:rFonts')
        rFonts.set(qn('w:ascii'), pn.font)
        rFonts.set(qn('w:hAnsi'), pn.font)
        rPr.append(rFonts)
    sz = OxmlElement('w:sz')
    sz.set(qn('w:val'), str(int(pn.size_pt * 2)))
    rPr.append(sz)
    if pn.color_hex and pn.color_hex.upper() != '000000':
        color_el = OxmlElement('w:color')
        color_el.set(qn('w:val'), pn.color_hex.lstrip('#').upper())
        rPr.append(color_el)
    return rPr
