# -*- coding: utf-8 -*-
"""
template_scanner.py：扫描 docx 模板中的 Jinja2 占位符。

只读模板，不渲染；返回模板里要求外部提供的数据字段清单。
"""
import re
import zipfile
import xml.etree.ElementTree as ET

from jinja2 import Environment, meta
from jinja2 import nodes as jinja_nodes


def _iter_docx_text(docx_path):
    """遍历 docx 中所有 XML 文本节点，产出原始字符串片段。"""
    with zipfile.ZipFile(docx_path, "r") as z:
        for name in z.namelist():
            # 只关心 word 主文档、页眉、页脚
            if not (name.startswith("word/document") or name.startswith("word/header") or name.startswith("word/footer")):
                continue
            data = z.read(name)
            try:
                root = ET.fromstring(data)
            except ET.ParseError:
                continue
            # w:t 节点存放文本
            for elem in root.iter():
                if elem.tag.endswith("}t"):
                    if elem.text:
                        yield elem.text


def _extract_source(docx_path):
    """把 docx 所有文本片段拼成一段连续的源文本，便于 Jinja2 解析。"""
    return "\n".join(_iter_docx_text(docx_path))


def _normalize_docxtpl_tags(source):
    """把 docxtpl 段落/行/单元格级标签 {%p ... %} {%tr ... %} {%tc ... %} 还原为 {% ... %}。

    docxtpl 用 {%p、{%tr、{%tc 控制块作用于段落、表格行、表格单元格；
    扫描占位符时不需要这些粒度信息，统一成标准 Jinja2 标签即可解析。
    """
    return re.sub(r"{%\s*(?:p|tr|tc)\s+", "{% ", source)


def _walk_blocks(node, result=None):
    """遍历 Jinja2 AST，收集 For / If 块信息。"""
    if result is None:
        result = []
    if node is None:
        return result

    if isinstance(node, jinja_nodes.For):
        iter_name = _expression_name(node.iter)
        target_name = _expression_name(node.target)
        result.append({
            "type": "for",
            "iterator": target_name,
            "variable": iter_name,
        })
        for child in node.body:
            _walk_blocks(child, result)
        for child in node.else_ or []:
            _walk_blocks(child, result)
        return result

    if isinstance(node, jinja_nodes.If):
        test_name = _expression_name(node.test)
        result.append({
            "type": "if",
            "condition": test_name,
        })
        for child in node.body:
            _walk_blocks(child, result)
        for child in node.else_ or []:
            _walk_blocks(child, result)
        return result

    if hasattr(node, "body"):
        for child in node.body:
            _walk_blocks(child, result)

    if hasattr(node, "else_") and node.else_:
        for child in node.else_:
            _walk_blocks(child, result)

    return result


def _expression_name(expr):
    """把 Jinja2 表达式尽量还原为可读的字符串。"""
    if expr is None:
        return None
    if isinstance(expr, jinja_nodes.Name):
        return expr.name
    if isinstance(expr, jinja_nodes.Const):
        return str(expr.value)
    if isinstance(expr, jinja_nodes.Getattr):
        return f"{_expression_name(expr.node)}.{expr.attr}"
    if isinstance(expr, jinja_nodes.Getitem):
        return f"{_expression_name(expr.node)}[{_expression_name(expr.arg)}]"
    return str(expr)


def scan_template(template_path):
    """扫描 docx 模板，返回占位符信息 JSON。

    返回:
        {
            "placeholders": ["project_title", "equipments", ...],  # 需要外部提供的最顶层变量
            "blocks": [
                {"type": "for", "iterator": "eq", "variable": "equipments"},
                {"type": "if", "condition": "show_layout"},
                ...
            ]
        }
    """
    source = _normalize_docxtpl_tags(_extract_source(template_path))
    if not source.strip():
        return {"placeholders": [], "blocks": []}

    env = Environment()
    ast = env.parse(source)

    variables = sorted(meta.find_undeclared_variables(ast))
    blocks = _walk_blocks(ast)

    return {
        "placeholders": variables,
        "blocks": blocks,
    }
