"""
Оркестратор двухпроходной обфускации — класс TwoPassObfuscator.

Процесс:
1. Проход 1 (сбор): извлечь текст → regex-сканирование → LLM-сканирование
2. Проход 2 (замена): применить mapping ко всем файлам
3. Сборка: ZIP + mapping.csv
"""

import io
import os
import time
import logging
from typing import Dict, List, Tuple, Callable, Optional

from . import extractor
from . import scanner
from . import replacer
from . import builder

log = logging.getLogger("drhider")


def _dedupe_file_names(
    files: List[Tuple[str, bytes, str]]
) -> List[Tuple[str, bytes, str]]:
    """Уникализировать имена файлов перед обработкой.

    Точный дубль (одинаковое имя + одинаковое содержимое) пропускается.
    При коллизии имени с ДРУГИМ содержимым добавляется суффикс _2, _3, ...

    Args:
        files: [(filename, content_bytes, content_type), ...]

    Returns:
        Новый список с уникальными именами (порядок сохранён).
    """
    seen: Dict[str, bytes] = {}
    out: List[Tuple[str, bytes, str]] = []

    for fname, content, ctype in files:
        candidate = fname
        if candidate in seen:
            if seen[candidate] == content:
                continue  # точный дубль — пропускаем
            base, ext = os.path.splitext(fname)
            n = 2
            while f"{base}_{n}{ext}" in seen:
                n += 1
            candidate = f"{base}_{n}{ext}"
        seen[candidate] = content
        out.append((candidate, content, ctype))

    return out


class TwoPassObfuscator:
    """Двухпроходный обфускатор документов.

    Проход 1: собрать все сущности из всех файлов → глобальный словарь замен.
    Проход 2: применить замены ко всем файлам → ZIP с результатом.

    Согласованность: одна и та же сущность во всех файлах получает
    одно и то же фиктивное значение.

    Attributes:
        _mapping: {оригинал: замена} — глобальный словарь
        _sorted_keys: ключи mapping, отсортированные по длине (убывание)
        _llm_client: опциональный LLM-клиент для NER
    """

    def __init__(self, llm_client=None):
        """Инициализировать обфускатор.

        Args:
            llm_client: Объект с методом .complete(prompt) -> str.
                        Если None — LLM-сканирование не выполняется.
        """
        self._mapping: Dict[str, str] = {}
        self._sorted_keys: List[str] = []
        self._compiled_re = None  # кэш скомпилированного regex (A2)
        self._counters: Dict[str, int] = {}  # Глобальные счётчики токенов
        self._llm_client = llm_client

    # ═══════════════════════════════════════════════════════════════════
    # Главная точка входа
    # ═══════════════════════════════════════════════════════════════════

    def obfuscate(
        self, files: List[Tuple[str, bytes, str]],
        progress_cb: Optional[Callable[[str, int, int, str], None]] = None,
        cancel_event=None, file_progress=None
    ) -> Tuple[bytes, str, Optional[dict]]:
        """Обфусцировать список файлов.

        Все форматы → Markdown → замена → результат.

        Args:
            files: [(filename, content_bytes, content_type), ...]
            progress_cb: Опциональный коллбек (phase, idx, total, fname, elapsed),
                         где phase ∈ {"start", "done"}, idx — 0-based индекс,
                         elapsed — время обработки конкретного файла (сек).
            cancel_event: threading.Event — мягкая остановка. Проверяется между
                          файлами LLM и между файлами замены; текущий файл добирается.
            file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам.
                           События: extract_done{total_chars}, file_start{chars,chunks,idx},
                           file_chunk{chunks_done,chunks_total,idx}, file_done{elapsed,idx}.

        Returns:
            (zip_bytes, csv_string, meta):
                zip_bytes — ZIP-архив с обфусцированными .md файлами (без mapping.csv)
                csv_string — содержимое mapping.csv как строка
                meta — None при штатном завершении; {"cancelled": True, "processed", "total"}
                       при прерывании (частичный результат)
        """
        # ── Предобработка: распаковать ZIP + уникализировать имена ──
        files = extractor.expand_zips(files)
        files = _dedupe_file_names(files)

        # имя файла (внутр. ключ all_texts) -> idx в files (для событий file_progress)
        name_to_idx = {fname: i for i, (fname, _, _) in enumerate(files)}
        llm_done: set = set()  # файлы, чей LLM-анализ ПОЛНОСТЬЮ завершён

        def _fp(event, fname, **fields):
            if event == "file_done":
                llm_done.add(fname)
            if file_progress:
                fields["idx"] = name_to_idx.get(fname)
                file_progress(event, fname, **fields)

        try:
            # ── Проход 1: сбор сущностей ──
            # Извлекаем Markdown из каждого файла
            all_texts: Dict[str, str] = {}
            total = len(files)
            file_times: List[float] = [0.0] * total
            skipped: set = set()  # имена файлов, которые не удалось обработать (битые)

            # A1: .doc конвертируются через HTTP-сервис liberta ПОСЛЕДОВАТЕЛЬНО.
            # Без параллельности: liberta не выдерживает конкурентные вызовы
            # (libreoffice падает с javaldx).

            cancelled_extract = False
            for i, (fname, content, ctype) in enumerate(files):
                # Мягкая остановка в фазе извлечения (между файлами) — не ждём все .doc
                if cancel_event is not None and cancel_event.is_set():
                    log.info("obfuscate: cancelled during extract, processed=%d/%d", i, total)
                    cancelled_extract = True
                    break
                display_name = os.path.basename(fname) or fname
                if progress_cb:
                    progress_cb("start", i, total, display_name, 0.0)

                t0 = time.time()
                try:
                    text = extractor.extract_text(fname, content, ctype)
                except Exception as e:
                    log.warning("obfuscate: skip file %r: %r", fname, e)
                    skipped.add(fname)
                    if progress_cb:
                        progress_cb("done", i, total, display_name, 0.0)
                    continue
                all_texts[fname] = text

                # Regex-сканирование (быстрое, локальное)
                if text and not text.startswith("[DOC binary"):
                    scanner.scan_regex(text, self._mapping, self._counters)
                file_times[i] += time.time() - t0

            # Объём извлечённого текста — для оценки времени (глобальной ETA)
            total_chars = sum(len(t) for t in all_texts.values())
            if file_progress:
                file_progress("extract_done", None, total_chars=total_chars,
                              per_file={fname: len(t) for fname, t in all_texts.items()})

            # LLM-сканирование (получает уже найденное regex'ом чтобы не дублировать)
            cancelled_llm = cancelled_extract
            if not cancelled_llm and self._llm_client:
                try:
                    scanner.scan_llm_ner(
                        all_texts, self._mapping, self._llm_client, self._counters,
                        cancel_event=cancel_event, file_progress=_fp,
                    )
                except scanner.CancelRequested:
                    log.info("obfuscate: LLM cancelled, llm_done=%d", len(llm_done))
                    cancelled_llm = True

            # Предсортировать ключи один раз (по убыванию длины)
            self._sorted_keys = sorted(
                self._mapping.keys(), key=len, reverse=True
            )
            # A2: кэш скомпилированного regex — один на все файлы прохода 2
            self._compiled_re = replacer._build_combined_re(self._sorted_keys)

            # ── Проход 2: замена сущностей ──
            cancelled = cancelled_llm
            results: List[Tuple[str, bytes]] = []

            for i, (in_fname, content, ctype) in enumerate(files):
                fname = in_fname
                display_name = os.path.basename(in_fname) or in_fname

                # Битый файл (не удалось извлечь текст) — пропускаем
                if fname in skipped:
                    if progress_cb:
                        progress_cb("done", i, total, display_name, 0.0)
                    continue

                # При прерывании в LLM — в результат идут ТОЛЬКО файлы с завершённым LLM
                if cancelled_llm and fname not in llm_done:
                    continue

                # Мягкая остановка в фазе замены (между файлами).
                # НЕ проверяем, если LLM уже прерван: файлы из llm_done добираем (сохраняем всё готовое).
                if not cancelled_llm and cancel_event is not None and cancel_event.is_set():
                    log.info("obfuscate: cancelled in replace phase, processed=%d", len(results))
                    cancelled = True
                    break

                t0 = time.time()
                obf_content = content  # По умолчанию — без изменений

                if fname.endswith('.doc'):
                    # .doc — конвертируется через liberta, обрабатывается как .docx
                    md_text = all_texts.get(fname, '')
                    replaced = replacer.apply_replacements(
                        md_text, self._mapping, self._sorted_keys, self._compiled_re
                    )
                    md_name = os.path.splitext(fname)[0] + '.md'
                    obf_content = replaced.encode('utf-8')
                    fname = md_name
                else:
                    # Все форматы → Markdown → замена в тексте
                    md_text = all_texts.get(fname, '')
                    replaced = replacer.apply_replacements(
                        md_text, self._mapping, self._sorted_keys, self._compiled_re
                    )
                    # Имя файла: заменить расширение на .md
                    md_name = os.path.splitext(fname)[0] + '.md'
                    obf_content = replaced.encode('utf-8')
                    fname = md_name

                file_times[i] += time.time() - t0
                results.append((fname, obf_content))
                if progress_cb:
                    progress_cb("done", i, total, display_name, round(file_times[i], 2))

            # ── Сборка результата ──
            csv_str = builder.build_mapping_csv(self._mapping)
            zip_data = builder.build_zip(results, csv_str)

            meta = None
            if cancelled:
                expected_total = total - len(skipped)
                meta = {"cancelled": True, "processed": len(results), "total": expected_total}
            return zip_data, csv_str, meta

        finally:
            # Очистка состояния (обфускатор может использоваться повторно)
            self._mapping.clear()
            self._sorted_keys.clear()
            self._compiled_re = None
            self._counters.clear()


# ═══════════════════════════════════════════════════════════════════════
# Удобная функция для быстрого вызова
# ═══════════════════════════════════════════════════════════════════════

def obfuscate_files(
    files: List[Tuple[str, bytes, str]], llm_client=None,
    progress_cb: Optional[Callable[[str, int, int, str], None]] = None,
    cancel_event=None, file_progress=None
) -> Tuple[bytes, str, Optional[dict]]:
    """Обфусцировать список файлов — удобная функция.

    Создаёт экземпляр TwoPassObfuscator и вызывает .obfuscate().

    Args:
        files: [(filename, content_bytes, content_type), ...]
        llm_client: Опциональный LLM-клиент
        progress_cb: Опциональный коллбек (phase, idx, total, fname)
        cancel_event: threading.Event — мягкая остановка (см. TwoPassObfuscator.obfuscate)
        file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам

    Returns:
        (zip_bytes, csv_string, meta)
    """
    obf = TwoPassObfuscator(llm_client=llm_client)
    return obf.obfuscate(files, progress_cb=progress_cb,
                         cancel_event=cancel_event, file_progress=file_progress)