# -*- coding: utf-8 -*-
"""DokoGen — Генератор документов из шаблонов DOCX (чистая версия)"""

import os
import re
import io
import shutil
import tempfile
import zipfile
from typing import Dict, List, Optional, Callable, Any
from datetime import datetime
from functools import lru_cache

try:
    from docx import Document
    from docx.oxml.ns import qn
    from docx.oxml import OxmlElement
    from lxml import etree
    DOCX_AVAILABLE = True
except ImportError:
    Document = None
    DOCX_AVAILABLE = False

# ============================================================
# 1. ПЕРЕМЕННЫЕ ШАБЛОНОВ И РЕГЕКСЫ
# ============================================================

# Паттерны для поиска переменных и маркеров циклов
VAR_PATTERN = re.compile(r'\{\{[^}]+\}\}')
OLD_VAR_PATTERN = re.compile(r'<[A-Za-z0-9_]+>')
UNMATCHED_PATTERN = re.compile(r'\{\{[^}]+\}\}|<[A-Za-z0-9_]+>')

LOOP_START = re.compile(r'(?:&lt;!--|<!--)\s*loop:(\w+)\s*(?:--&gt;|-->)')
LOOP_END = re.compile(r'(?:&lt;!--|<!--)\s*loop_end\s*(?:--&gt;|-->)')

# Маппинг старых переменных <...> → {{...}}
OLD_VAR_MAP = {
    '<CompanyName1>': '{{company_full_name}}',
    '<CompanyName2>': '{{company_short_name}}',
    '<CompanyName1Gent>': '{{company_full_name_genitive}}',
    '<CompanyName2Gent>': '{{company_short_name_genitive}}',
    '<CompanyName1Datv>': '{{company_full_name_dative}}',
    '<CompanyName2Datv>': '{{company_short_name_dative}}',
    '<CompanyName1Accs>': '{{company_full_name_accs}}',
    '<CompanyName2Accs>': '{{company_short_name_accs}}',
    '<CompanyName1Ablt>': '{{company_full_name_ablt}}',
    '<CompanyName2Ablt>': '{{company_short_name_ablt}}',
    '<CompanyName1Loct>': '{{company_full_name_loct}}',
    '<CompanyName2Loct>': '{{company_short_name_loct}}',
    '<Address>': '{{company_address}}',
    '<INN>': '{{company_inn}}',
    '<OGRN>': '{{company_ogrn}}',
    '<KPP>': '{{company_kpp}}',
    '<ChiefFio>': '{{chief_fio_initials}}',
    '<ChiefFioShort>': '{{chief_fio_short}}',
    '<ChiefPosition>': '{{chief_position}}',
    '<RspFio>': '{{responsible_fio_initials}}',
    '<RspFioShort>': '{{responsible_fio_short}}',
    '<RspPosition>': '{{responsible_position}}',
    '<CommissionList>': '{{commission_list}}',
    '<DocumentNumber>': '{{document_number}}',
    '<DocumentDate>': '{{document_date}}',
    '<Date>': '{{date}}',
}


# ============================================================
# 2. ОБРАБОТКА ЦИКЛОВ НА УРОВНЕ ZIP/XML
# ============================================================

W_NS = 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'


def _normalize_loop_markers(doc_xml: str) -> str:
    """Word разбивает маркеры <!-- loop:... --> и переменные {{item.x}}
    на несколько run'ов (например '<!-- ' в одном <w:t>, 'loop:info' в другом,
    'rmation -->' в третьем). Склеиваем текст каждого параграфа в один run,
    чтобы маркеры и переменные циклов снова стали едиными."""
    try:
        root = etree.fromstring(doc_xml.encode('utf-8'))
    except Exception:
        return doc_xml

    changed = False
    for p in root.iter(f'{{{W_NS}}}p'):
        # Собираем весь текст параграфа (по всем run'ам)
        texts = []
        for t in p.iter(f'{{{W_NS}}}t'):
            texts.append(t.text or '')
        full = ''.join(texts)

        # Интересуют только параграфы с маркерами циклов ИЛИ переменными цикла
        if 'loop:' not in full and 'loop_end' not in full and '{{item.' not in full:
            continue

        # Сохраняем форматирование первого run
        first_r = p.find(f'{{{W_NS}}}r')
        rpr = None
        if first_r is not None:
            rpr_el = first_r.find(f'{{{W_NS}}}rPr')
            if rpr_el is not None:
                rpr = etree.fromstring(etree.tostring(rpr_el))

        # Оставляем только pPr, удаляем все run'ы
        for child in list(p):
            if child.tag != f'{{{W_NS}}}pPr':
                p.remove(child)

        # Создаём один run с полным текстом параграфа
        new_r = etree.SubElement(p, f'{{{W_NS}}}r')
        if rpr is not None:
            new_r.append(rpr)
        new_t = etree.SubElement(new_r, f'{{{W_NS}}}t')
        new_t.text = full
        new_t.set('{http://www.w3.org/XML/1998/namespace}space', 'preserve')
        changed = True

    if not changed:
        return doc_xml

    return etree.tostring(root, xml_declaration=True, encoding='UTF-8', standalone=True).decode('utf-8')


def expand_loops_in_zip(docx_path: str, loops: Dict[str, List[Dict]]) -> str:
    """Обрабатывает циклы напрямую в ZIP/DOCX на уровне XML.
    Работает до загрузки python-docx — гарантирует корректную вложенность.
    
    Args:
        docx_path: путь к шаблону .docx
        loops: словарь {имя_цикла: [{field: value}, ...]}
    
    Returns:
        путь к обработанному .docx (временный файл)
    """
    if not loops:
        return docx_path

    # Временный файл
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.docx')
    tmp.close()

    wrote = False  # была ли хоть одна запись в tmp
    max_passes = 30
    for _ in range(max_passes):
        # Читаем ВСЕ файлы из исходного ZIP
        all_files = {}
        try:
            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

        # Word мог разбить маркеры циклов на несколько run'ов — склеиваем
        doc_xml = _normalize_loop_markers(doc_xml)

        # Находим ВСЕ маркеры циклов
        start_markers = list(LOOP_START.finditer(doc_xml))
        end_markers = list(LOOP_END.finditer(doc_xml))

        if not start_markers:
            break  # нет больше циклов

        # Внешний цикл = первый start (самый ранний в документе)
        outer = start_markers[0]
        outer_key = outer.group(1)

        # Соответствующий end = первый end после outer (или последний, если вложенные)
        # Считаем глубину: каждый start +1, каждый end -1
        depth = 1
        outer_end = None
        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

        items = loops.get(outer_key, [])
        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:
                for name, data in all_files.items():
                    z.writestr(name, data)
            wrote = True
            continue

        # Извлекаем шаблон (всё между start и end маркерами)
        template = doc_xml[outer.end():outer_end.start()]

        # НЕ удаляем loop_end из шаблона — они принадлежат вложенным циклам!
        template_clean = template

        # Размножаем шаблон для каждого элемента данных
        expanded_parts = []
        for idx, item in enumerate(items, 1):
            clone = template_clean
            for field, value in item.items():
                clone = clone.replace('{{item.%s}}' % field, str(value) if value else '—')
            clone = clone.replace('{{item.number}}', str(idx))
            clone = clone.replace('{{number}}', str(idx))
            expanded_parts.append(clone)

        # Собираем новый 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:
            for name, data in all_files.items():
                z.writestr(name, data)
        wrote = True

    if not wrote:
        # Циклов в шаблоне не было — возвращаем исходный файл,
        # пустой временный удаляем (иначе python-docx упадёт с Package not found)
        try:
            os.unlink(tmp.name)
        except Exception:
            pass
        return docx_path

    return tmp.name


# ============================================================
# 3. ЗАМЕНА ПЕРЕМЕННЫХ В ЭЛЕМЕНТАХ DOCX
# ============================================================

def _replace_text_in_para(para, replacements: Dict[str, str]) -> bool:
    """Заменяет {{var}} в параграфе. Возвращает True, если были замены."""
    full = para.text
    matches = VAR_PATTERN.findall(full)
    if not matches:
        return False
    
    new_text = full
    for var in matches:
        val = replacements.get(var)
        if val is not None:
            new_text = new_text.replace(var, str(val))
    
    if new_text == full:
        return False
    
    # Сохраняем форматирование первого run
    if para.runs:
        first = para.runs[0]
        para.clear()
        run = para.add_run(new_text)
        _copy_run_style(first, run)
    else:
        para.clear()
        para.add_run(new_text)
    return True


def _copy_run_style(source, target):
    """Копирует форматирование run."""
    try:
        target.bold = source.bold
        target.italic = source.italic
        target.underline = source.underline
        if source.font:
            target.font.size = source.font.size
            target.font.name = source.font.name
    except Exception:
        pass


def _replace_in_paragraphs(paragraphs, replacements):
    for para in paragraphs:
        _replace_text_in_para(para, replacements)


def _replace_in_tables(tables, replacements):
    for table in tables:
        for row in table.rows:
            for cell in row.cells:
                _replace_in_paragraphs(cell.paragraphs, replacements)


def _replace_in_headers_footers(doc, replacements):
    for section in doc.sections:
        for hf in [section.header, section.footer,
                   section.first_page_header, section.first_page_footer]:
            if hf:
                _replace_in_paragraphs(hf.paragraphs, replacements)


def _strip_prooferr(doc):
    """Удаляет proofErr из docx (красные волнистые линии)."""
    ns = 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
    for elem in doc.element.iter():
        for child in list(elem):
            if child.tag == f'{{{ns}}}proofErr':
                elem.remove(child)


def _highlight_unmatched(doc):
    """Подсвечивает незаменённые переменные красным."""
    ns = 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'

    def _highlight_run(run):
        rpr = run._element.find(f'{{{ns}}}rPr')
        if rpr is None:
            rpr = OxmlElement('w:rPr')
            run._element.insert(0, rpr)
        color = OxmlElement('w:color')
        color.set(f'{{{ns}}}val', 'FF0000')
        rpr.append(color)

    for para in doc.paragraphs:
        for run in para.runs:
            if UNMATCHED_PATTERN.search(run.text or ''):
                _highlight_run(run)

    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                for para in cell.paragraphs:
                    for run in para.runs:
                        if UNMATCHED_PATTERN.search(run.text or ''):
                            _highlight_run(run)

    for section in doc.sections:
        for hf in [section.header, section.footer]:
            if hf:
                for para in hf.paragraphs:
                    for run in para.runs:
                        if UNMATCHED_PATTERN.search(run.text or ''):
                            _highlight_run(run)


def _log_unmatched(doc, replacements=None, log_func=None):
    """Логирует незаменённые переменные."""
    if not log_func:
        return
    unmatched = []
    seen = set()
    for para in doc.paragraphs:
        for m in UNMATCHED_PATTERN.finditer(para.text):
            var = m.group(0)
            if var not in seen:
                seen.add(var)
                unmatched.append(var)
    if unmatched:
        log_func(f"⚠️ Незаменённые переменные ({len(unmatched)}): "
                 f"{', '.join(unmatched[:20])}"
                 f"{' ...' if len(unmatched) > 20 else ''}")
    else:
        log_func("✅ Все переменные успешно заменены")


# ============================================================
# 4. МИГРАЦИЯ СТАРЫХ ПЕРЕМЕННЫХ
# ============================================================

def _migrate_old_vars(doc):
    """Заменяет <OldVar> на {{new_var}}."""
    def _migrate_text(text):
        return OLD_VAR_PATTERN.sub(lambda m: OLD_VAR_MAP.get(m.group(0), m.group(0)), text)

    for para in doc.paragraphs:
        new_text = _migrate_text(para.text)
        if new_text != para.text:
            if para.runs:
                first = para.runs[0]
                para.clear()
                run = para.add_run(new_text)
                _copy_run_style(first, run)
            else:
                para.clear()
                para.add_run(new_text)

    for table in doc.tables:
        for row in table.rows:
            for cell in row.cells:
                for para in cell.paragraphs:
                    new_text = _migrate_text(para.text)
                    if new_text != para.text:
                        if para.runs:
                            first = para.runs[0]
                            para.clear()
                            run = para.add_run(new_text)
                            _copy_run_style(first, run)
                        else:
                            para.clear()
                            para.add_run(new_text)

    for section in doc.sections:
        for hf in [section.header, section.footer,
                   section.first_page_header, section.first_page_footer]:
            if hf:
                for para in hf.paragraphs:
                    new_text = _migrate_text(para.text)
                    if new_text != para.text:
                        if para.runs:
                            first = para.runs[0]
                            para.clear()
                            run = para.add_run(new_text)
                            _copy_run_style(first, run)
                        else:
                            para.clear()
                            para.add_run(new_text)


# ============================================================
# 5. ОСНОВНАЯ ФУНКЦИЯ ГЕНЕРАЦИИ
# ============================================================

def process_template(
    template_path: str,
    replacements: Dict[str, str],
    output_path: Optional[str] = None,
    log_func: Optional[Callable] = None,
) -> str:
    """Обрабатывает шаблон DOCX: подстановка переменных, циклы, сохранение.
    
    Args:
        template_path: путь к файлу шаблона .docx
        replacements: словарь {переменная: значение}
        output_path: путь для сохранения результата
        log_func: функция логирования (str → None)
    
    Returns:
        путь к сгенерированному файлу
    """
    if not DOCX_AVAILABLE:
        raise ImportError("python-docx не установлен")

    if not os.path.exists(template_path):
        raise FileNotFoundError(f"Шаблон не найден: {template_path}")

    if output_path is None:
        base, ext = os.path.splitext(template_path)
        output_path = f"{base}_result{ext}"

    # ШАГ 1: Обрабатываем циклы на уровне ZIP/XML (до загрузки python-docx)
    loops = replacements.pop('_loops', {})
    working_path = template_path
    if loops:
        working_path = expand_loops_in_zip(template_path, loops)
        replacements['_loops'] = loops  # возвращаем на всякий случай

    # ШАГ 2: Загружаем через python-docx
    doc = Document(working_path)

    # ШАГ 3: Миграция старых переменных <...> → {{...}}
    _migrate_old_vars(doc)

    # ШАГ 4: Убираем proofErr
    _strip_prooferr(doc)

    # ШАГ 5: Замена переменных (3 прохода)
    for _ in range(3):
        _replace_in_paragraphs(doc.paragraphs, replacements)
        _replace_in_tables(doc.tables, replacements)
        _replace_in_headers_footers(doc, replacements)

    # ШАГ 6: Логирование незаменённых
    _log_unmatched(doc, replacements, log_func)

    # ШАГ 7: Подсветка незаполненных
    _highlight_unmatched(doc)

    # ШАГ 8: Сохранение
    doc.save(output_path)

    # ШАГ 9: Если использовали временный файл — удаляем
    if working_path != template_path:
        try:
            os.unlink(working_path)
        except Exception:
            pass

    if log_func:
        log_func(f"✅ {os.path.basename(output_path)}")
    return output_path