"""
Проход 1: обнаружение сущностей в тексте документов.

Два метода сканирования:
1. scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
2. scan_llm_ner — универсальное LLM-обнаружение (любые приватные данные)

Замена: гибридные токены — читаемый шаблон + уникальный номер.
Пример: Иванов_0001, ООО_Технология_0034, ул_Ленина_0015
"""

import re
import json
import random
import time
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Optional

from .config import (
    ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN,
    RU_SURNAMES, COMPANY_NOUNS, CITY_POOL, STREET_POOL,
)

log = logging.getLogger("drhider")


# ═══════════════════════════════════════════════════════════════════════════
# Гибридные токены: тип → (пул_шаблонов, префикс_для_счётчика)
# ═══════════════════════════════════════════════════════════════════════════

TYPE_POOLS = {
    "person_name":  (RU_SURNAMES, "person"),
    "person":       (RU_SURNAMES, "person"),
    "company":      (COMPANY_NOUNS, "company"),
    "phone":        (["+7_000_000"], "phone"),
    "email":        (["email"], "email"),
    "address":      (CITY_POOL, "address"),
    "inn":          (["ИНН"], "inn"),
    "inn_ul":       (["ИНН"], "inn"),
    "inn_fl":       (["ИНН"], "inn"),
    "kpp":          (["КПП"], "kpp"),
    "ogrn":         (["ОГРН"], "ogrn"),
    "bik":          (["БИК"], "bik"),
    "passport":     (["Паспорт"], "passport"),
    "bank_account": (["Счёт"], "bank"),
    "rs":           (["РС"], "bank"),
    "ks":           (["КС"], "bank"),
    "contract_number": (["Договор"], "contract"),
    "other_id":     (["ID"], "other"),
}
_FALLBACK_POOL = (["Данные"], "data")


def _next_token(entity_type: str, counters: Dict[str, int]) -> str:
    """Сгенерировать токен замены: случайный шаблон из пула + номер.

    Одна и та же сущность в mapping получает один и тот же токен
    (mapping dict гарантирует согласованность).
    Разные сущности одного типа получают РАЗНЫЕ шаблоны из пула.

    Args:
        entity_type: Тип сущности (person_name, company, phone...)
        counters: Глобальные счётчики {prefix_key: count}

    Returns:
        Строка вида "Иванов_0001", "ООО_Технология_0034", "[Адрес_0015]"
    """
    pool, prefix_key = TYPE_POOLS.get(entity_type, _FALLBACK_POOL)
    template = random.choice(pool)
    key = f"{prefix_key}_{template}"
    counters[key] = counters.get(key, 0) + 1
    return f"{template}_{counters[key]:04d}"

log = logging.getLogger("drhider")


def scan_regex(text: str, mapping: Dict[str, str], counters: Dict[str, int]) -> None:
    """Сканировать текст regex-паттернами, заполнить словарь замен.

    Замена: гибридные токены (шаблон + номер).

    Args:
        text: Текст документа для сканирования
        mapping: Словарь замен (мутабельный, пополняется)
        counters: Глобальные счётчики токенов (мутабельный)
    """
    # ── Сущности по regex-паттернам ──
    for entity_type, pattern in ENTITY_PATTERNS.items():
        for match in re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE):
            original = match.group(0).strip()
            if not original or original in mapping:
                continue
            mapping[original] = _next_token(entity_type, counters)

    # ── Названия компаний ──
    for match in COMPANY_PATTERN.finditer(text):
        original = match.group(0).strip()
        if not original or original in mapping:
            continue
        mapping[original] = _next_token("company", counters)

    # ── ФИО ──
    for match in PERSON_PATTERN.finditer(text):
        original = match.group(0).strip()
        if not original or original in mapping:
            continue
        mapping[original] = _next_token("person_name", counters)


# ═══════════════════════════════════════════════════════════════════════════
# LLM-NER: пофайловый чанкинг (Sonnet-схема)
# ═══════════════════════════════════════════════════════════════════════════

class CancelRequested(Exception):
    """Обработка прервана пользователем (мягкая остановка между файлами LLM)."""

    def __init__(self, llm_failures: int = 0):
        self.llm_failures = llm_failures

# Приоритет границ для чанков (от предпочтительных к жёстким)
_CHUNK_BOUNDARIES = ['\n\n', '\n', '. ', '? ', '! ', '; ', ', ', ' ']
_CHUNK_SIZE = 6000          # символов на чанк (≈1500-2000 токенов RU — NER-качество)
_CHUNK_OVERLAP = 500        # перекрытие чанков (сущности на стыке)
_LLM_CONCURRENCY = 1        # параллельность убрана — всё последовательно (liberta/LLM не тянет конкурентность)


def split_into_chunks(text: str, size: int = _CHUNK_SIZE,
                      overlap: int = _CHUNK_OVERLAP) -> List[str]:
    """Разбить текст на чанки ~size символов с перекрытием overlap.

    Режем по границам (\n\n, \n, ..., пробел) — чтобы не разрывать слова/сущности.
    """
    chunks: List[str] = []
    start = 0
    tl = len(text)
    while start < tl:
        end = min(start + size, tl)
        if end == tl:
            chunks.append(text[start:])
            break
        cut = None
        for boundary in _CHUNK_BOUNDARIES:
            pos = text.rfind(boundary, start + size // 2, end)
            if pos != -1:
                cut = pos + len(boundary)
                break
        if cut is None:
            cut = end  # жёсткий разрез — нет ни одного пробела/границы
        chunks.append(text[start:cut])
        # overlap: отступить назад, найти начало слова
        ov_start = max(start, cut - overlap)
        space = text.find(' ', ov_start)
        start = (space + 1) if (space != -1 and space < cut) else ov_start
    return chunks


def normalize_entity(s: str) -> str:
    """Нормализовать строку сущности для надёжного дедупа (регистр/пробелы/тире)."""
    s = s.strip()
    s = re.sub(r'\s+', ' ', s)
    s = s.lower()
    s = re.sub(r'[«»\u201c\u201d\u2018\u2019"\u0027]', '"', s)
    s = re.sub(r'[\u2014\u2013\u2212-]', '-', s)
    s = s.replace('\u00ad', '')   # мягкий перенос
    s = s.replace('ё', 'е')       # Ё→Е (OCR/PDF)
    return s


def _verify_entity(value: str, full_text: str) -> bool:
    """Верификация: сущность реально есть в полном тексте файла (отсекает галлюцинации)."""
    if value in full_text:
        return True
    return normalize_entity(value) in normalize_entity(full_text)


def _build_llm_prompt(text: str, already_found: List[str]) -> str:
    """Собрать промпт для одного текстового фрагмента (чанк/файл)."""
    af = "\n".join(f"- {k}" for k in already_found) if already_found else "(none)"
    return (
        "You are a PII (Personally Identifiable Information) detector for Russian legal documents.\n\n"
        "Find ALL sensitive or private information that is NOT already captured below.\n\n"

        "Already captured by regex (skip these exact strings):\n"
        f"{af}\n\n"

        "Return a JSON array. Each item must have:\n"
        '  - "type": MUST be one of: person_name, company, phone, email, address,\n'
        '    inn, kpp, ogrn, bik, passport, bank_account, contract_number, other_id\n'
        '    Use "other_id" only if nothing else fits. If unsure — do NOT include.\n'
        '  - "value": the EXACT string from the text (copy verbatim)\n\n'

        "NOT PII — do NOT include these (legal terms, not private data):\n"
        "- Roles: Исполнитель, Заказчик, Стороны, Сторона, Покупатель, Поставщик\n"
        "- Titles: Генеральный директор, Директор, действующего на основании\n"
        "- Legal terms: Услуги, Договор, Приложение, Спецификация, НДС, Устава\n\n"

        "Rules:\n"
        "- Copy value VERBATIM — same spaces, punctuation, case as in text\n"
        "- Same value → include only once\n"
        "- Return ONLY the JSON array, no other text, no markdown wrapping\n\n"

        f"Text:\n{text}"
    )


def _call_llm(text: str, llm_client) -> List[dict]:
    """Один вызов LLM по фрагменту → проверенный список сущностей."""
    raw = llm_client.complete(text).strip()
    if raw.startswith("```"):
        lines = raw.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        raw = "\n".join(lines).strip()

    try:
        entities = json.loads(raw)
    except json.JSONDecodeError as original_error:
        # Модель иногда добавляет текст вокруг массива; raw_decode корректно
        # учитывает вложенные массивы и квадратные скобки внутри JSON-строк.
        decoder = json.JSONDecoder()
        entities = None
        for index, char in enumerate(raw):
            if char != "[":
                continue
            try:
                candidate, _ = decoder.raw_decode(raw, index)
            except json.JSONDecodeError:
                continue
            if isinstance(candidate, list):
                entities = candidate
                break
        if entities is None:
            raise ValueError("LLM response does not contain a JSON array") from original_error

    if not isinstance(entities, list):
        raise ValueError("LLM response must be a JSON array")
    for entity in entities:
        if not isinstance(entity, dict):
            raise ValueError("LLM entity must be a JSON object")
        value = entity.get("value")
        if not isinstance(value, str) or not value.strip():
            raise ValueError("LLM entity value must be a non-empty string")
        entity_type = entity.get("type")
        if entity_type is not None and not isinstance(entity_type, str):
            raise ValueError("LLM entity type must be a string")
    return entities


def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
                  llm_client, counters: Dict[str, int],
                  cancel_event=None, file_progress=None) -> int:
    """Универсальное LLM-обнаружение приватных данных (пофайлово, целиком).

    Каждый файл обрабатывается ПОЛНОСТЬЮ: текст разбивается на чанки
    (_CHUNK_SIZE=6000, overlap=_CHUNK_OVERLAP), чанки обрабатываются последовательно
    (_LLM_CONCURRENCY), сущности дедуплицируются и верифицируются против полного
    текста файла (отсечка галлюцинаций). regex-найденное не дублируется.

    Args:
        all_texts: {filename: text_content} — тексты всех файлов
        mapping: Словарь замен (мутабельный, пополняется)
        llm_client: Объект с методом .complete(prompt) -> str
        counters: Глобальные счётчики токенов (мутабельный)
        cancel_event: threading.Event — мягкая остановка между файлами (текущий добирается до конца)
        file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам
    """
    already_found = list(mapping.keys())

    # Порядок файлов: крупные вперёд (приоритет на долгую обработку)
    file_items = sorted(
        all_texts.items(), key=lambda kv: len(kv[1]), reverse=True
    )
    llm_failures = 0

    for fname, full_text in file_items:
        # Мягкая остановка ТОЛЬКО между файлами — текущий файл добирается до конца
        if cancel_event is not None and cancel_event.is_set():
            raise CancelRequested(llm_failures)

        if not full_text or full_text.startswith("[DOC binary"):
            continue
        chunks = split_into_chunks(full_text)

        if file_progress:
            file_progress("file_start", fname, chars=len(full_text), chunks=len(chunks))

        t0 = time.time()

        # Параллельные вызовы LLM по чанкам этого файла
        if len(chunks) > 1 and _LLM_CONCURRENCY > 1:
            with ThreadPoolExecutor(max_workers=_LLM_CONCURRENCY) as _ex:
                futures = [_ex.submit(_call_llm, _build_llm_prompt(c, mapping.keys()), llm_client)
                           for c in chunks]
                all_ent = []
                for fut in futures:
                    try:
                        all_ent.extend(fut.result())
                    except Exception as e:
                        llm_failures += 1
                        log.warning("LLM NER chunk failed in %r: %s", fname, e)
        else:
            all_ent = []
            for k, c in enumerate(chunks):
                try:
                    all_ent.extend(_call_llm(_build_llm_prompt(c, mapping.keys()), llm_client))
                except Exception as e:
                    llm_failures += 1
                    log.warning("LLM NER chunk failed in %r: %s", fname, e)
                if file_progress:
                    file_progress("file_chunk", fname,
                                  chunks_done=k + 1, chunks_total=len(chunks))

        # Дедуп + верификация + добавление в mapping
        seen: set = set()
        for ent in all_ent:
            val = ent.get("value", "").strip()
            if not val or val in mapping:
                continue
            nk = normalize_entity(val)
            if nk in seen:
                continue
            seen.add(nk)
            if not _verify_entity(val, full_text):
                log.debug("LLM hallucination dropped: %r", val)
                continue
            raw_type = ent.get("type")
            ent_type = raw_type.strip().lower().replace(" ", "_") if raw_type else ""
            ent_type = ent_type or "other_id"
            mapping[val] = _next_token(ent_type, counters)

        if file_progress:
            file_progress("file_done", fname, elapsed=round(time.time() - t0, 2))

    return llm_failures