From c5466f2992c23f7cc308834a6add2f35a1a428b1 Mon Sep 17 00:00:00 2001
From: DokoGen <prog1764@labnethub.ru>
Date: Fri, 31 Jul 2026 00:29:12 +0400
Subject: [PATCH] fix: nested loops work - preserve inner loop_end markers

---
 dokogen/generator.py | 55 +++++++++++++++++++++++++++++---------------
 test_generator.py    |  4 ++--
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/dokogen/generator.py b/dokogen/generator.py
index 7e79cc9..0ac3e27 100644
--- a/dokogen/generator.py
+++ b/dokogen/generator.py
@@ -82,20 +82,25 @@ def expand_loops_in_zip(docx_path: str, loops: Dict[str, List[Dict]]) -> str:
     if not loops:
         return docx_path
 
-    # Временная копия
+    # Временный файл
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.docx')
     tmp.close()
-    shutil.copy2(docx_path, tmp.name)
 
     max_passes = 30
     for _ in range(max_passes):
-        # Читаем document.xml
+        # Читаем ВСЕ файлы из исходного ZIP
+        all_files = {}
         try:
-            with zipfile.ZipFile(tmp.name, 'r') as z:
-                doc_xml = z.read('word/document.xml').decode('utf-8')
+            with zipfile.ZipFile(docx_path if _ == 0 else tmp.name, 'r') as z:
+                for name in z.namelist():
+                    all_files[name] = z.read(name)
+                doc_xml = all_files.get('word/document.xml', b'').decode('utf-8')
         except Exception:
             break
 
+        if not doc_xml:
+            break
+
         # Находим ВСЕ маркеры циклов
         start_markers = list(LOOP_START.finditer(doc_xml))
         end_markers = list(LOOP_END.finditer(doc_xml))
@@ -107,12 +112,26 @@ def expand_loops_in_zip(docx_path: str, loops: Dict[str, List[Dict]]) -> str:
         outer = start_markers[0]
         outer_key = outer.group(1)
 
-        # Соответствующий end = последний end после outer
+        # Соответствующий end = первый end после outer (или последний, если вложенные)
+        # Считаем глубину: каждый start +1, каждый end -1
+        depth = 1
         outer_end = None
-        for m in reversed(end_markers):
-            if m.start() > outer.start():
-                outer_end = m
-                break
+        all_markers = sorted(
+            [('start', m) for m in start_markers] +
+            [('end', m) for m in end_markers],
+            key=lambda x: x[1].start()
+        )
+        found_start = False
+        for kind, m in all_markers:
+            if m.start() <= outer.start():
+                continue
+            if kind == 'start':
+                depth += 1
+            elif kind == 'end':
+                depth -= 1
+                if depth == 0:
+                    outer_end = m
+                    break
 
         if outer_end is None:
             break
@@ -121,19 +140,17 @@ def expand_loops_in_zip(docx_path: str, loops: Dict[str, List[Dict]]) -> str:
         if not items:
             # Удаляем блок цикла полностью
             doc_xml = doc_xml[:outer.start()] + doc_xml[outer_end.end():]
+            all_files['word/document.xml'] = doc_xml.encode('utf-8')
             with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as z:
-                z.writestr('word/document.xml', doc_xml.encode('utf-8'))
+                for name, data in all_files.items():
+                    z.writestr(name, data)
             continue
 
         # Извлекаем шаблон (всё между start и end маркерами)
         template = doc_xml[outer.end():outer_end.start()]
 
-        # Удаляем только последний loop_end (принадлежит внешнему циклу)
+        # НЕ удаляем loop_end из шаблона — они принадлежат вложенным циклам!
         template_clean = template
-        last_end_markers = list(LOOP_END.finditer(template_clean))
-        if last_end_markers:
-            last_end = last_end_markers[-1]
-            template_clean = template_clean[:last_end.start()] + template_clean[last_end.end():]
 
         # Размножаем шаблон для каждого элемента данных
         expanded_parts = []
@@ -148,9 +165,11 @@ def expand_loops_in_zip(docx_path: str, loops: Dict[str, List[Dict]]) -> str:
         # Собираем новый XML
         new_xml = doc_xml[:outer.start()] + ''.join(expanded_parts) + doc_xml[outer_end.end():]
 
-        # Сохраняем
+        # Сохраняем ВСЕ файлы ZIP, обновляя только document.xml
+        all_files['word/document.xml'] = new_xml.encode('utf-8')
         with zipfile.ZipFile(tmp.name, 'w', zipfile.ZIP_DEFLATED) as z:
-            z.writestr('word/document.xml', new_xml.encode('utf-8'))
+            for name, data in all_files.items():
+                z.writestr(name, data)
 
     return tmp.name
 
diff --git a/test_generator.py b/test_generator.py
index 3064586..91ae39d 100644
--- a/test_generator.py
+++ b/test_generator.py
@@ -9,7 +9,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from dokogen.models import Company, CommissionMember, InformationSystem
 from dokogen.variables import build_replacements
-from dokogen.generator import process_template, _expand_loops_in_zip
+from dokogen.generator import process_template, expand_loops_in_zip
 
 def create_test_template():
     """Creates a test .docx with nested loops for verification."""
@@ -108,7 +108,7 @@ def test_zip_loop_expansion():
     
     # Process with ZIP-level loop expansion
     print('\nProcessing loops via ZIP...')
-    result_path = _expand_loops_in_zip(template_path, loops)
+    result_path = expand_loops_in_zip(template_path, loops)
     
     # Verify result
     with zipfile.ZipFile(result_path, 'r') as z: