#!/usr/bin/env python3
"""
Утилиты для подстановки переменных в DOCX-шаблоны.

Запуск:
    python3 docx_form_controls.py replace input.docx output.docx key1=value1 key2=value2 [--tables tables.json]
"""
import argparse
import copy
import json
import os
import re
import shutil
import sys
import zipfile
from xml.sax.saxutils import escape

from docx import Document
from docx.shared import Pt

VARIABLE_RE = re.compile(r"\{\{\s*([^}]+?)\s*\}\}")


def _replace_in_xml(xml_path: str, values: dict[str, str]):
    """Заменяет {{variable}} на значения в XML-файле DOCX."""
    with open(xml_path, "r", encoding="utf-8") as f:
        text = f.read()

    def repl(match: re.Match) -> str:
        key = match.group(1)
        value = values.get(key, match.group(0))
        return escape(value, {"'": "'", '"': """})

    text = VARIABLE_RE.sub(repl, text)

    with open(xml_path, "w", encoding="utf-8") as f:
        f.write(text)


def _replace_text_in_doc(doc: Document, values: dict[str, str]):
    """Заменяет {{variable}} на значения во всех параграфах и таблицах документа."""
    placeholder_re = re.compile(r"\{\{\s*([^}]+?)\s*\}\}")

    for paragraph in doc.paragraphs:
        text = paragraph.text
        if not placeholder_re.search(text):
            continue
        new_text = placeholder_re.sub(lambda m: values.get(m.group(1), m.group(0)), text)
        paragraph.clear()
        run = paragraph.add_run()
        run.text = new_text

    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                for paragraph in cell.paragraphs:
                    text = paragraph.text
                    if not placeholder_re.search(text):
                        continue
                    new_text = placeholder_re.sub(lambda m: values.get(m.group(1), m.group(0)), text)
                    paragraph.clear()
                    run = paragraph.add_run()
                    run.text = new_text


def _insert_tables(doc: Document, tables: dict[str, dict]):
    """Вставляет таблицы вместо плейсхолдеров {{table_code}}."""
    paragraphs = list(doc.paragraphs)
    for paragraph in paragraphs:
        text = paragraph.text.strip()
        match = VARIABLE_RE.fullmatch(text)
        if not match:
            continue
        code = match.group(1)
        table_data = tables.get(code)
        if not table_data:
            continue

        columns = table_data.get("columns", [])
        rows = table_data.get("rows", [])

        paragraph.clear()

        if not columns:
            paragraph.add_run("[пустая таблица]")
            continue

        table = doc.add_table(rows=1, cols=len(columns))
        table.style = "Table Grid"

        hdr_cells = table.rows[0].cells
        for i, col in enumerate(columns):
            hdr_cells[i].text = str(col.get("name", ""))
            for paragraph in hdr_cells[i].paragraphs:
                for run in paragraph.runs:
                    run.font.bold = True

        for row_data in rows:
            row_cells = table.add_row().cells
            for i, col in enumerate(columns):
                col_id = col.get("id")
                value = ""
                if isinstance(row_data, dict) and col_id is not None:
                    value = row_data.get(col_id, "")
                row_cells[i].text = str(value) if value is not None else ""

        # Перемещаем таблицу на место параграфа
        tbl_element = copy.deepcopy(table._tbl)
        paragraph._element.addnext(tbl_element)


def replace_variables(
    input_path: str,
    output_path: str,
    values: dict[str, str],
    tables: dict[str, dict] | None = None,
    assets: dict[str, dict] | None = None,
):
    """Подменяет переменные {{var}} в DOCX на значения."""
    if tables or assets:
        # Используем python-docx для точечной вставки таблиц и изображений
        doc = Document(input_path)
        _replace_text_in_doc(doc, values)
        if tables:
            _insert_tables(doc, tables)
        if assets:
            _insert_assets(doc, assets)
        doc.save(output_path)
        return

    tmp_dir = f"{output_path}.tmp"
    if os.path.exists(tmp_dir):
        shutil.rmtree(tmp_dir)

    with zipfile.ZipFile(input_path, "r") as zip_in:
        zip_in.extractall(tmp_dir)

    xml_path = os.path.join(tmp_dir, "word", "document.xml")
    if os.path.exists(xml_path):
        _replace_in_xml(xml_path, values)

    # Также заменяем в header/footer
    for subdir in ("word",):
        sub_path = os.path.join(tmp_dir, subdir)
        if not os.path.isdir(sub_path):
            continue
        for fname in os.listdir(sub_path):
            if fname.startswith(("header", "footer")) and fname.endswith(".xml"):
                _replace_in_xml(os.path.join(sub_path, fname), values)

    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zip_out:
        for root, _, files in os.walk(tmp_dir):
            for file in files:
                full_path = os.path.join(root, file)
                arcname = os.path.relpath(full_path, tmp_dir)
                zip_out.write(full_path, arcname)

    shutil.rmtree(tmp_dir)


def _insert_assets(doc: Document, assets: dict[str, dict]):
    """Вставляет изображения вместо плейсхолдеров {{asset_code}}."""
    from docx.shared import Cm

    paragraphs = list(doc.paragraphs)
    for paragraph in paragraphs:
        text = paragraph.text.strip()
        match = VARIABLE_RE.fullmatch(text)
        if not match:
            continue
        code = match.group(1)
        asset = assets.get(code)
        if not asset:
            continue

        paragraph.clear()
        asset_type = asset.get("type", "image")
        if asset_type != "image":
            paragraph.add_run("[неподдерживаемый актив]")
            continue

        asset_path = asset.get("path")
        if not asset_path or not os.path.exists(asset_path):
            paragraph.add_run("[изображение не найдено]")
            continue

        try:
            run = paragraph.add_run()
            picture = run.add_picture(asset_path)
            width_cm = asset.get("width")
            height_cm = asset.get("height")
            if width_cm:
                picture.width = Cm(float(width_cm))
            if height_cm:
                picture.height = Cm(float(height_cm))
            elif width_cm and picture.width:
                # Сохраняем пропорции
                aspect = picture.height / picture.width if picture.width else 1
                picture.height = int(picture.width * aspect)
        except Exception as e:
            paragraph.add_run(f"[ошибка вставки изображения: {e}]")


def main():
    parser = argparse.ArgumentParser(description="DOCX variable substitution")
    subparsers = parser.add_subparsers(dest="command", required=True)

    replace_parser = subparsers.add_parser("replace", help="replace {{var}} with values")
    replace_parser.add_argument("input", help="input DOCX")
    replace_parser.add_argument("output", help="output DOCX")
    replace_parser.add_argument("values", nargs="*", help="key=value pairs")
    replace_parser.add_argument("--tables", dest="tables_path", help="JSON file with table variables")
    replace_parser.add_argument("--assets", dest="assets_path", help="JSON file with asset variables")

    args = parser.parse_args()

    if args.command == "replace":
        values = {}
        for pair in args.values or []:
            if "=" not in pair:
                print(f"Invalid pair: {pair}", file=sys.stderr)
                sys.exit(1)
            key, value = pair.split("=", 1)
            values[key] = value

        tables = None
        if args.tables_path:
            with open(args.tables_path, "r", encoding="utf-8") as f:
                tables = json.load(f)

        assets = None
        if args.assets_path:
            with open(args.assets_path, "r", encoding="utf-8") as f:
                assets = json.load(f)

        replace_variables(args.input, args.output, values, tables, assets)
        print(f"Replaced: {args.output}")


if __name__ == "__main__":
    main()