#!/usr/bin/env python3
"""
Unified Document Parser — объединённый парсер DOCX и PDF.

Поддерживает два режима:
  1. API mode (по умолчанию) — Unstructured API server (smart_indexer / parse_pdf)
  2. Local mode — локальная библиотека unstructured (docx_pipeline)

Возвращает единый список чанков с метаданными, очищенный от шума
(шапки, футеры, листы согласования, подписи, страницы).
"""

import os
import sys
import re
import json
import hashlib
import uuid
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Optional, Any


# ============================================================
# CONFIG
# ============================================================

DEFAULT_API_URL = "http://192.168.1.103:8005/general/v0/general"
DEFAULT_LOCAL_MODE = False


def get_config() -> dict:
    """Собирает конфиг из переменных окружения или дефолтов."""
    return {
        "api_url": os.environ.get("UNSTRUCTURED_API_URL", DEFAULT_API_URL),
        "local_mode": os.environ.get("LOCAL_MODE", "0").lower() in ("1", "true", "yes"),
        "languages": os.environ.get("LANGUAGES", "rus").split(","),
        "strategy": os.environ.get("STRATEGY", "hi_res"),
        "pdf_infer_table_structure": os.environ.get("PDF_INFER_TABLE", "true").lower() == "true",
        "min_text_length": int(os.environ.get("MIN_TEXT_LENGTH", "15")),
        "min_chunk_length": int(os.environ.get("MIN_CHUNK_LENGTH", "40")),
        "chunk_max_chars": int(os.environ.get("CHUNK_MAX_CHARS", "1800")),
        "chunk_combine_under": int(os.environ.get("CHUNK_COMBINE_UNDER", "300")),
        "chunk_new_after": int(os.environ.get("CHUNK_NEW_AFTER", "6000")),
        "batch_size": int(os.environ.get("BATCH_SIZE", "32")),
    }


# ============================================================
# CLEAN / FILTER HELPERS
# ============================================================

def clean_text(text: str) -> str:
    """Убирает множественные пробелы/переносы."""
    if not text:
        return ""
    text = re.sub(r"\s+", " ", text)
    return text.strip()




# ============================================================
# SECTION DETECTION & ID GENERATION
# ============================================================

# Regex for section titles like "1. Нормативные ссылки", "2. Термины и определения"
_SECTION_RE = re.compile(r"^\s*(\d+)[\.\s]\s*(.+)$")

# Keywords that indicate section titles from Unstructured API
SECTION_KEYWORDS = [
    "Нормативные ссылки",
    "Термины",
    "Термины и определения",
    "Области применения",
    "Общие положения",
    "Порядок организации",
    "Обязанности",
    "Контроль",
    "Ответственность",
    "Заключительные положения",
    "Приложения",
]

def _detect_section_number(text: str) -> tuple:
    """Extract section number and title from text.
    
    Returns (section_number, section_title) or (None, None).
    """
    if not text:
        return None, None
    
    # Try numbered section pattern: "3. Нормативные ссылки"
    m = _SECTION_RE.match(text.strip())
    if m:
        num = int(m.group(1))
        title = m.group(2).strip()
        return num, title
    
    # Try unnumbered section pattern (fallback)
    for kw in SECTION_KEYWORDS:
        if kw in text:
            # Extract just the keyword part as title
            idx = text.index(kw)
            title = text[idx:].strip()
            return None, title
    
    return None, None


def _generate_section_id(section_number: int, section_title: str) -> str:
    """Generate a stable section_id based on section number and title."""
    if section_number is not None:
        return f"section_{section_number}"
    # Fallback: hash the title
    h = hashlib.md5(section_title.encode("utf-8")).hexdigest()[:8]
    return f"section_{h}"


def is_page_number(text: str) -> bool:
    return bool(re.match(r"^Стр\.\s*\d+\s*из\s*\d+$", text))


def is_header_noise(text: str) -> bool:
    patterns = [
        "АО «ХК «Сибцем»",
        "Тип документа:",
        "Ведущее подразделение:",
        "Дата утверждения:",
        "Редакция 1",
        "Оглавление",
    ]
    return any(p in text for p in patterns)


def is_header_chunk(text: str) -> bool:
    """DOCX-специфичные фильтры шапки/обложки."""
    text_lower = text.lower()
    if any(kw in text_lower for kw in [
        "утверждаю", "президент", "генеральный директор",
        "согласован", "в.в. шаповалов",
    ]):
        return True
    if any(kw in text_lower for kw in ["кемерово", "2008г.", "рег-20-1", "регистр"]):
        if len(text) < 300:
            return True
    if "лист соглас" in text_lower:
        return True
    if "содержание" in text_lower and len(text) < 500:
        return True
    if any(kw in text_lower for kw in [
        "разработчик:", "дата разработки:", "стр. из:",
    ]):
        return True
    if any(kw in text_lower for kw in [
        "список лиц", "разработавших", "финансовый директор",
        "главный инженер",
    ]):
        return True
    return False


def is_footer_chunk(text: str) -> bool:
    """DOCX-специфичные фильтры футера."""
    text_lower = text.lower()
    if any(kw in text_lower for kw in [
        "документ:", "стр. из", "дата разработки:",
        "разработчик:", "унифицированная",
    ]):
        return True
    return False


def is_junk(item: dict, text: str, cfg: dict) -> bool:
    """Общий фильтр мусора для обоих режимов."""
    if not text:
        return True
    if len(text) < cfg["min_text_length"]:
        return True
    if item.get("type") in ("Header", "Footer", "Image"):
        return True
    if is_page_number(text):
        return True
    if is_header_noise(text):
        return True
    # DOCX-специфичные фильтры
    if is_header_chunk(text):
        return True
    if is_footer_chunk(text):
        return True
    return False


# ============================================================
# HTML TABLE → TEXT
# ============================================================

def _ocr_table_fixes(text: str) -> str:
    """Применяет патч-исправления для OCR-ошибок в таблицах."""
    if not text:
        return text
    
    # Порядок важен — длинные патчи первыми, более специфичные до общих
    fixes = [
        # OCR-исправления заголовков (самые важные) — длинные сначала
        ("рредакции ы", "редакции"),
        ("рредакции :", "редакции :"),
        ("рредакции:", "редакции:"),
        ("рредакции", "редакции"),
        ("едакции ы", "редакции"),
        ("едакции :", "редакции :"),
        ("едакции:", "редакции:"),
        ("едакции", "редакции"),
        ("едаты ", "даты "),
        ("едаты:", "даты:"),
        ("едаты", "даты"),
        # OCR "р" перед словами (массовое)
        ("р внесенных :", "внесенных :"),
        ("р внесенных:", "внесенных:"),
        ("р внесенных", "внесенных"),
        ("р системное", "системное"),
        ("р системного", "системного"),
        ("р системном", "системном"),
        ("р Департамента", "Департамента"),
        ("р технологий", "технологий"),
        ("р администрирование", "администрирование"),
        ("р администрирования", "администрирования"),
        ("р информационных", "информационных"),
        ("р Группы", "Группы"),
        ("р технолог", "технолог"),
        ("р технологическим", "технологическим"),
        ("р персонализированный", "персонализированный"),
        # OCR "р" перед другими словами
        ("р ", " "),
        # OCR "Ошесв" / "Обесв" / "Отесв"
        ("Ошесвииа", "Общества"),
        ("Отесват", "Общества"),
        ("Обесват", "Общества"),
        ("Ошесват", "Общества"),
        ("Ошесв", "Общества"),
        # Дублирование слов (администрирование администрирования)
        (" администрирования администрирования", " администрирования"),
        ("администрирования администрирования", " администрирования"),
        (" администрирование администрирование", " администрирование"),
        ("администрирование администрирование", " администрирование"),
        # Убираем вертикальные разделители
        ("| Создано", "Создано"),
        ("| ", " "),
        ("|", ""),
        ("Список лиц:", "Список лиц"),
        ("Список лиц,", "Список лиц"),
        # Бауэр
        ("Бауэ ", "Бауэр "),
        ("Бауэ", "Бауэр"),
        ("Бауэ С", "Бауэр С"),
        # Общие OCR-исправления
        ("пк", "ПК"),
        ("ПК пк", "ПК"),
        ("РроюСаг", "PhotoCar"),
        ("р р", " "),
    ]
    for bad, good in fixes:
        text = text.replace(bad, good)

    # Убираем "р" как отдельный OCR-мусор — ТОЛЬКО в конце слов (не в середине)
    # НЕ используем \br\s — он ломает русские слова (Номер→Номе)
    text = re.sub(r'\sр\b', '', text)

    # Убираем двойные/тройные двоеточия
    text = re.sub(r':\s*:\s*', ':', text)
    text = re.sub(r':$', '', text)

    # Убираем "№:" из заголовков
    text = re.sub(r'№:\s*', '', text)

    # Убираем двойные буквы (Бауэрр → Бауэр и т.д.)
    text = re.sub(r'(.)\1\1+', r'\1\1', text)

    # Убираем тройные пробелы
    text = re.sub(r'  +', ' ', text)

    # Чистим лишние пробелы
    text = re.sub(r'\s+', ' ', text)
    return text.strip()


def _is_ocr_table(html: str) -> bool:
    """Проверяет, является ли HTML таблицы сильно повреждённым OCR."""
    if not html:
        return False
    # Если HTML содержит много OCR-мусора
    ocr_markers = ["едации", "едаты", "Ошесв", "Обесв", "рредакции", "р систем", "р Департ"]
    ocr_count = sum(1 for m in ocr_markers if m in html)
    return ocr_count >= 2


def _extract_table_text_fallback(table_text: str) -> str:
    """Извлекает структуру таблицы из уже обработанного text поля API."""
    if not table_text:
        return ""
    lines = []
    current_record = None
    for line in table_text.split('\n'):
        line = clean_text(line)
        if not line:
            continue
        # Проверяем "Запись N:"
        m = re.match(r'^Запись\s+(\d+):\s*(.*)', line)
        if m:
            current_record = m.group(1)
            rest = m.group(2)
            if rest:
                lines.append(f"Запись {current_record}: {clean_text(rest)}")
            continue
        lines.append(line)

    return "\n".join(lines)


def html_table_to_text(html: str, table_text: str = "") -> str:
    """Конвертирует HTML таблицу в читаемый текст.

    Сначала чистит HTML от OCR-ошибок, затем парсит.
    Если OCR слишком сильный — fallback из table_text.
    """
    if not html and not table_text:
        return ""

    # Определяем, стоит ли использовать HTML-парсинг
    # Если в HTML слишком много OCR-мусора — используем fallback
    html_clean = _ocr_table_fixes(html)
    
    # Проверяем качество HTML после очистки
    ocr_markers = ["едации", "едаты", "Ошесв", "Обесв", "рредакции", "р систем", "р Департ"]
    ocr_count = sum(1 for m in ocr_markers if m in html_clean)
    
    # Если после очистки всё ещё есть OCR-ошибки — используем fallback
    if ocr_count >= 2:
        return _extract_table_text_fallback(table_text) if table_text else ""
    
    # Парсим HTML
    from html.parser import HTMLParser

    class TableParser(HTMLParser):
        def __init__(self):
            super().__init__()
            self.rows = []
            self.current_row = []
            self.in_table = False

        def handle_starttag(self, tag, attrs):
            if tag == "table":
                self.in_table = True
            elif tag == "tr" and self.in_table:
                self.current_row = []
            elif tag in ("td", "th") and self.in_table:
                attrs_dict = dict(attrs)
                rowspan = int(attrs_dict.get("rowspan", 1))
                colspan = int(attrs_dict.get("colspan", 1))
                self.current_row.append({
                    "text": "",
                    "rowspan": rowspan,
                    "colspan": colspan,
                })

        def handle_data(self, data):
            if self.current_row:
                self.current_row[-1]["text"] += data

        def handle_endtag(self, tag):
            if tag == "tr" and self.in_table and self.current_row:
                self.rows.append(self.current_row)
                self.current_row = []

    parser = TableParser()
    parser.feed(html_clean)

    rows = parser.rows
    if not rows:
        return _extract_table_text_fallback(table_text) if table_text else ""

    # Убираем пустые строки
    rows = [r for r in rows if any(cell["text"].strip() for cell in r)]

    if not rows:
        return _extract_table_text_fallback(table_text) if table_text else ""

    # Определяем ширину таблицы с учётом rowspan
    max_cols = max(len(row) for row in rows)
    table_width = max(max_cols, 1)

    # Строим виртуальную сетку с rowspan
    grid = [[None] * table_width for _ in range(len(rows))]
    for row_idx, row in enumerate(rows):
        col_idx = 0
        for cell in row:
            # Пропускаем занятые ячейки (rowspan)
            while col_idx < table_width and grid[row_idx][col_idx] is not None:
                col_idx += 1
            if col_idx >= table_width:
                break
            grid[row_idx][col_idx] = {
                "text": cell["text"],
                "rowspan": cell["rowspan"],
                "colspan": cell["colspan"],
            }
            # Занимаем ячейки для colspan
            for c in range(1, cell["colspan"]):
                next_col = col_idx + c
                if next_col < table_width:
                    grid[row_idx][next_col] = {"text": "", "rowspan": 1, "colspan": 1}

    # Определяем, является ли первая строка заголовком
    def is_header_row(row_strings):
        """row_strings = list of str. Headers are short, no numbers."""
        texts = [s.strip() for s in row_strings if s.strip()]
        if not texts:
            return False
        short = all(len(t) < 50 for t in texts)
        has_number = any(re.match(r"^\d+\.?\s*$", t) for t in texts)
        return short and not has_number

    # Парсим ячейки в плоские строки
    def parse_row_cells(row_cells):
        result = []
        for cell in row_cells:
            if cell is None:
                result.append("")
            else:
                result.append(_ocr_table_fixes(cell["text"].strip()))
        return result

    parsed_rows = []
    for row in grid:
        parsed_rows.append(parse_row_cells(row))

    # Фильтруем пустые строки
    parsed_rows = [r for r in parsed_rows if any(c.strip() for c in r)]

    if not parsed_rows:
        return _extract_table_text_fallback(table_text) if table_text else ""

    # Определяем заголовок
    header_idx = 0
    if is_header_row(parsed_rows[0]):
        header_idx = 0
    else:
        for i, row in enumerate(parsed_rows):
            if is_header_row(row):
                header_idx = i
                break

    headers = parsed_rows[header_idx]

    # Формируем вывод
    lines = []
    for row_idx, row in enumerate(parsed_rows[header_idx + 1:], header_idx + 1):
        lines.append(f"Запись {row_idx}:")
        for col_idx, cell in enumerate(row):
            if not cell:
                continue
            header = headers[col_idx] if col_idx < len(headers) else f"Колонка {col_idx+1}"
            header = _ocr_table_fixes(header)
            lines.append(f"{header}: {cell}")
        lines.append("")

    result = "\n".join(lines).strip()
    return result if result else _extract_table_text_fallback(table_text) if table_text else ""


# ============================================================
# MERGE FRAGMENTS
# ============================================================

def merge_fragments(elements: list) -> list:
    """Объединяет короткие чанки в буфер, таблицы — отдельно.
    
    Улучшена логика слияния:
    - Учитывает номер страницы (объединяет только с соседних)
    - Учитывает section_title (не объединяет через разделы)
    - Более высокий порог для разделения (300 вместо 80)
    """
    merged = []
    buffer = []
    buffer_page = None
    buffer_section = None

    def flush_buffer():
        nonlocal buffer
        if not buffer:
            return
        text = clean_text(" ".join(buffer))
        if text:
            merged.append({
                "type": "MergedText",
                "text": text,
                "page_number": buffer_page,
                "section_title": buffer_section,
            })
        buffer = []

    for el in elements:
        text = clean_text(el.get("text", ""))
        if not text:
            continue
        el_type = el.get("type")
        page = el.get("metadata", {}).get("page_number", 0)
        section = el.get("section_title", "")

        if el_type == "Table":
            flush_buffer()
            merged.append({
                "type": "Table",
                "text": text,
                "html": el.get("metadata", {}).get("text_as_html", ""),
                "page_number": page,
                "section_title": section,
            })
            continue

        if el_type in ("NarrativeText", "ListItem", "UncategorizedText"):
            # Check if we should flush the buffer (different section or page jump)
            if buffer_section and section and buffer_section != section:
                flush_buffer()
                buffer = []
                buffer_section = section
                buffer_page = page
            
            if len(text) < 300:
                buffer.append(text)
                buffer_page = page
                buffer_section = section
            else:
                flush_buffer()
                merged.append({
                    "type": "Text",
                    "text": text,
                    "page_number": page,
                    "section_title": section,
                })
        else:
            flush_buffer()

    flush_buffer()
    return merged


# ============================================================
# API MODE — parse via Unstructured API
# ============================================================

def parse_document_api(file_path: str, cfg: dict) -> list:
    """Парсинг через Unstructured API (работает и с PDF, и с DOCX)."""
    import requests

    filename = os.path.basename(file_path)
    print(f"\n📄 API Парсинг: {filename}")

    payload_data = {
        "strategy": cfg["strategy"],
        "languages": cfg["languages"],
        "pdf_infer_table_structure": str(cfg["pdf_infer_table_structure"]),
    }

    ext = os.path.splitext(filename)[1].lower()
    mime_type = "application/pdf" if ext == ".pdf" else (
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    )

    with open(file_path, "rb") as f:
        files = {"files": (filename, f, mime_type)}
        timeout = cfg.get("timeout", 900)
        response = requests.post(
            cfg["api_url"],
            files=files,
            data=payload_data,
            timeout=timeout,
        )
        response.raise_for_status()
        return response.json()


# ============================================================
# LOCAL MODE — parse DOCX via python unstructured
# ============================================================

def parse_document_local(file_path: str, cfg: dict) -> list:
    """Локальный парсинг DOCX через unstructured.partition.docx."""
    from unstructured.partition.docx import partition_docx
    from unstructured.chunking.title import chunk_by_title

    print(f"\n📄 Local DOCX Парсинг: {os.path.basename(file_path)}")

    elements = partition_docx(
        file_path,
        strategy=cfg["strategy"],
        infer_table_structure=True,
    )
    print(f"  Partitioned: {len(elements)} elements")

    chunks = chunk_by_title(
        elements,
        max_characters=cfg["chunk_max_chars"],
        combine_text_under_n_chars=cfg["chunk_combine_under"],
        new_after_n_chars=cfg["chunk_new_after"],
    )
    print(f"  Chunked: {len(chunks)} chunks")

    # Конвертируем в единый формат с API
    result = []
    for c in chunks:
        item = {
            "type": c.category,
            "text": c.text,
            "metadata": {
                "page_number": getattr(c.metadata, "page_number", 0) if hasattr(c, "metadata") and c.metadata else 0,
                "text_as_html": getattr(c.metadata, "text_as_html", "") if hasattr(c, "metadata") and c.metadata else "",
            },
        }
        # section_title из metadata
        if hasattr(c, "metadata") and c.metadata:
            if hasattr(c.metadata, "page_label") and c.metadata.page_label:
                item["metadata"]["page_number"] = c.metadata.page_label
        result.append(item)

    return result


# ============================================================
# CLEAN + MERGE (единый пайплайн)
# ============================================================

def clean_and_merge(data: list, cfg: dict) -> list:
    """Очистка + фильтрация + слияние — общий для обоих режимов."""
    cleaned = []
    current_title = "Общая информация"
    current_section_num = None
    
    for item in data:
        raw_text = item.get("text", "")
        text = clean_text(raw_text)
        if is_junk(item, text, cfg):
            continue
        
        # Try to detect section titles from content (not just Title type)
        section_num, section_title = _detect_section_number(text)
        if section_num is not None or section_title != text:
            if section_num is not None:
                current_section_num = section_num
            current_title = section_title if section_title else text
            # Don't add the title itself as content
            continue
        
        item["section_title"] = current_title
        item["section_number"] = current_section_num
        cleaned.append(item)

    merged = merge_fragments(cleaned)
    return merged


# ============================================================
# BUILD CHUNKS + CONTENT
# ============================================================

def build_chunks(merged: list, filename: str, cfg: dict) -> list:
    """Формирует финальные чанки с content для индексации.
    
    Добавлен section_id, улучшена фильтрация мусорных чанков.
    """
    chunks = []
    for idx, chunk in enumerate(merged):
        chunk_type = chunk["type"]
        section_title = chunk["section_title"]
        page_number = chunk["page_number"]
        section_num = chunk.get("section_number", None)
        
        # Generate section_id
        section_id = _generate_section_id(section_num, section_title)
        
        if chunk_type == "Table":
            table_text = html_table_to_text(chunk.get("html", ""), chunk.get("text", ""))
            content = f"Раздел: {section_title}\nТип: Таблица\n\n{table_text}"
        else:
            content = f"Раздел: {section_title}\n\n{chunk['text']}"

        content = clean_text(content)
        if len(content) < cfg["min_chunk_length"]:
            continue
        
        # Filter out garbage chunks (too short content after header)
        if chunk_type == "Table" and len(table_text) < 50 if "table_text" in locals() else len(content) < 100:
            continue
        
        # Filter out chunks that are just headers with no real content
        if chunk["type"] in ("MergedText", "Text") and content.count("\n") <= 1 and len(content) < 100:
            continue

        chunk_hash = hashlib.sha256(
            (filename + str(idx) + content).encode("utf-8")
        ).hexdigest()

        chunks.append({
            "id": str(uuid.uuid5(uuid.NAMESPACE_DNS, chunk_hash)),
            "content": content,
            "metadata": {
                "filename": filename,
                "chunk_index": idx,
                "section_title": section_title,
                "section_id": section_id,
                "page_number": page_number,
                "content_type": chunk_type,
                "chunk_hash": chunk_hash,
            },
        })

    return chunks





# ============================================================
# MAIN ENTRY POINT
# ============================================================

def parse_document(file_path: str, cfg: Optional[dict] = None) -> dict:
    """
    Универсальный парсер документов (PDF + DOCX).

    Возвращает:
        {
            "filename": str,
            "chunks": list[dict],  # каждый: {id, content, metadata}
            "stats": {
                "total_elements": int,
                "total_chunks_before_filter": int,
                "total_chunks_after_filter": int,
                "total_final_chunks": int,
            }
        }
    """
    if cfg is None:
        cfg = get_config()

    filename = os.path.basename(file_path)

    # --- PARSE ---
    if cfg["local_mode"]:
        raw_data = parse_document_local(file_path, cfg)
    else:
        raw_data = parse_document_api(file_path, cfg)

    total_elements = len(raw_data)

    # --- CLEAN + MERGE ---
    merged = clean_and_merge(raw_data, cfg)
    total_after_filter = len(merged)

    # --- BUILD CHUNKS ---
    chunks = build_chunks(merged, filename, cfg)
    total_final = len(chunks)

    stats = {
        "total_elements": total_elements,
        "total_chunks_before_filter": total_elements,
        "total_chunks_after_filter": total_after_filter,
        "total_final_chunks": total_final,
    }

    return {
        "filename": filename,
        "chunks": chunks,
        "stats": stats,
    }


# ============================================================
# CLI
# ============================================================

def main():
    """CLI: python document_parser.py <file> [--api|--local] [--json]"""
    import argparse

    parser = argparse.ArgumentParser(description="Unified DOCX/PDF Document Parser")
    parser.add_argument("file", help="Path to PDF or DOCX file")
    parser.add_argument("--api", action="store_true", default=None,
                        help="Use API mode (default: auto-detect from env)")
    parser.add_argument("--local", action="store_true", default=None,
                        help="Use local unstructured library")
    parser.add_argument("--json", action="store_true", default=False,
                        help="Output parsed JSON to stdout")
    parser.add_argument("--output", "-o", default=None,
                        help="Output JSON file path")
    args = parser.parse_args()

    file_path = os.path.abspath(args.file)
    if not os.path.exists(file_path):
        print(f"❌ File not found: {file_path}")
        sys.exit(1)

    # Override mode from CLI
    if args.local:
        cfg = get_config()
        cfg["local_mode"] = True
    elif args.api:
        cfg = get_config()
        cfg["local_mode"] = False
    else:
        cfg = get_config()

    result = parse_document(file_path, cfg)

    print(f"\n✅ {result['filename']}")
    print(f"   Elements: {result['stats']['total_elements']}")
    print(f"   After filter: {result['stats']['total_chunks_after_filter']}")
    print(f"   Final chunks: {result['stats']['total_final_chunks']}")

    # Output
    output_json = json.dumps(result, ensure_ascii=False, indent=2)

    if args.json or args.output:
        if args.output:
            with open(args.output, "w", encoding="utf-8") as f:
                f.write(output_json)
            print(f"\n📄 Saved to: {args.output}")
        else:
            print(output_json)

    return result


if __name__ == "__main__":
    main()