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

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

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

import re
import json
import random
import logging
from typing import Dict

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)


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

    LLM САМА определяет типы. Замена — гибридные токены.

    Args:
        all_texts: {filename: text_content} — тексты всех файлов
        mapping: Словарь замен (мутабельный, пополняется)
        llm_client: Объект с методом .complete(prompt) -> str
        counters: Глобальные счётчики токенов (мутабельный)
    """
    # Формируем комбинированный текст: имя файла + первые 3000 символов
    combined = "\n\n---FILE---\n\n".join(
        f"FILE: {fname}\n{t[:3000]}" for fname, t in all_texts.items()
    )

    # ── Универсальный промпт: STEP 1 (Исполнитель) + STEP 2 (ПДн) ──
    already_found = "\n".join(f"- {k}" for k in mapping.keys()) if mapping else "(none)"

    prompt = (
        "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"{already_found}\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{combined[:8000]}"
    )

    try:
        # Отправляем запрос в LLM
        raw = llm_client.complete(prompt)

        # Чистим markdown-обёртку (если LLM вернула ```json ... ```)
        raw = raw.strip()
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[1]
            if raw.endswith("```"):
                raw = raw[:-3]

        entities = json.loads(raw)

        # Добавляем найденные сущности в mapping
        for ent in entities:
            val = ent.get("value", "").strip()
            if not val or val in mapping:
                continue

            # Тип — любой (LLM сама придумала)
            ent_type = ent.get("type", "").strip().lower().replace(" ", "_")
            mapping[val] = _next_token(ent_type, counters)

    except Exception as e:
        log.warning("LLM NER failed: %s", e)