lib/parser/ntr_fileparser/data_classes/parsed_table.py

"""
Модуль содержит классы для представления таблиц в документе.
"""

import warnings
from dataclasses import asdict, dataclass, field
from typing import Any, Callable, Optional

from .parsed_structure import DocumentElement
from .parsed_text_block import TextStyle


@dataclass
class TableTag:
    """
    Тег для классификации таблицы.
    """

    name: str = ""
    value: str = ""

    def to_dict(self) -> dict[str, Any]:
        """
        Преобразует тег таблицы в словарь.

        Returns:
            dict[str, Any]: Словарное представление тега таблицы.
        """
        return asdict(self)


@dataclass
class ParsedRow(DocumentElement):
    """
    Строка таблицы.
    """

    index: int | str | None = None
    cells: list[str] = field(default_factory=list)
    style: TextStyle = field(default_factory=TextStyle)
    anchors: list[str] = field(default_factory=list)
    links: list[str] = field(default_factory=list)
    is_header: bool = False

    def to_string(
        self,
        header: Optional['ParsedRow'] = None,
        note: Optional[str] = None,
    ) -> str:
        """
        Преобразует строку таблицы в строковое представление в виде маркированного списка.

        Args:
            header (Optional[ParsedRow]): Заголовок столбцов для форматирования.
            note (Optional[str]): Примечание к строке. Не будет использовано, если строка не содержит *,
            а в примечании есть *

        Returns:
            str: Строковое представление строки таблицы.
        """
        if not self.cells:
            return ""

        # Если у нас есть хедер, то форматируем как "ключ: значение" с маркерами
        if header:
            if len(header.cells) != len(self.cells):
                raise ValueError("Количество ячеек в строке и хедере не совпадает")
            result = '\n'.join(
                f"- {header.cells[i]}: {self.cells[i]}" for i in range(len(self.cells))
            )

        # Если у нас есть только две колонки, форматируем как "ключ: значение" с маркером
        elif len(self.cells) == 2:
            result = f"- {self.cells[0].strip()}: {self.cells[1].strip()}"

        # Иначе просто форматируем все ячейки через разделитель
        else:
            result = '\n'.join(f"- {cell.strip()}" for cell in self.cells)

        if note:

            if ('*' in result) != ('*' in note):
                return result
            else:
                return f"{result}\nПримечание: {note}"

        return result

    def apply(self, func: Callable[[str], str]) -> None:
        """
        Применяет функцию ко всем ячейкам строки.

        Args:
            func (Callable[[str], str]): Функция для применения к текстовым элементам.
        """
        self.cells = [func(cell) for cell in self.cells]
        self.anchors = [func(anchor) for anchor in self.anchors]
        self.links = [func(link) for link in self.links]

    def to_dict(self) -> dict[str, Any]:
        """
        Преобразует строку таблицы в словарь.

        Returns:
            dict[str, Any]: Словарное представление строки таблицы.
        """
        return {
            'index': self.index,
            'cells': self.cells,
            'style': self.style.to_dict(),
            'anchors': self.anchors,
            'links': self.links,
            'is_header': self.is_header,
            'page_number': self.page_number,
            'index_in_document': self.index_in_document,
        }


@dataclass
class ParsedSubtable(DocumentElement):
    """
    Подтаблица внутри основной таблицы.
    """

    title: str | None = None
    header: ParsedRow | None = None
    rows: list[ParsedRow] = field(default_factory=list)

    def to_string(
        self,
        header: Optional['ParsedRow'] = None,
        note: Optional[str] = None,
    ) -> str:
        """
        Преобразует подтаблицу в строковое представление.

        Returns:
            str: Строковое представление подтаблицы.
        """
        if self.header:
            header = self.header

        result = []
        if self.title:
            result.append(f"## {self.title}")

        if len(self.rows) == 0:
            if header:
                result.append(header.to_string(note=note))
            if note:
                result.append(f"Примечание: {note}")

        # Обрабатываем каждую строку таблицы
        for i, row in enumerate(self.rows, start=1):
            # Добавляем номер строки (начиная с 1)
            result.append(f"### Строка {i}")
            result.append(row.to_string(header=header, note=note))

        return "\n".join(result)

    def apply(self, func: Callable[[str], str]) -> None:
        """
        Применяет функцию ко всем элементам подтаблицы.

        Args:
            func (Callable[[str], str]): Функция для применения к текстовым элементам.
        """
        if self.title:
            self.title = func(self.title)
        if self.header:
            self.header.apply(func)
        for row in self.rows:
            row.apply(func)

    def to_dict(self) -> dict[str, Any]:
        """
        Преобразует подтаблицу в словарь.

        Returns:
            dict[str, Any]: Словарное представление подтаблицы.
        """
        result = {'title': self.title, 'rows': [row.to_dict() for row in self.rows]}
        if self.header:
            result['header'] = self.header.to_dict()

        # Добавляем поля из DocumentElement
        result['page_number'] = self.page_number
        result['index_in_document'] = self.index_in_document

        return result

    def has_merged_cells(self) -> bool:
        """
        Проверяет наличие объединенных ячеек в подтаблице.

        Returns:
            bool: True, если в подтаблице есть строки с разным количеством ячеек.
        """
        if not self.rows:
            return False

        # Получаем количество ячеек в строках
        cell_counts = [len(row.cells) for row in self.rows]
        if len(set(cell_counts)) > 1:
            return True

        return False


@dataclass
class ParsedTable(DocumentElement):
    """
    Таблица из документа.
    """

    title: str | None = None
    note: str | None = None
    classified_tags: list[TableTag] = field(default_factory=list)
    index: list[str] = field(default_factory=list)
    headers: list[ParsedRow] = field(default_factory=list)
    subtables: list[ParsedSubtable] = field(default_factory=list)
    table_style: dict[str, Any] = field(default_factory=dict)
    title_index_in_paragraphs: int | None = None

    def to_string(self) -> str:
        """
        Преобразует таблицу в строковое представление.

        Returns:
            str: Строковое представление таблицы.
        """
        # Формируем заголовок таблицы
        table_header = ""
        if self.title:
            table_header = f"# {self.title}"

        final_result = []

        common_header = None
        if self.headers:
            common_header = ParsedRow(
                cells=[
                    '/'.join(header.cells[i] for header in self.headers)
                    for i in range(len(self.headers[0].cells))
                ]
            )
            
        if len(self.subtables) == 0:
            if common_header:
                final_result.append(common_header.to_string(note=self.note))
            else:
                final_result.append(self.note)

        # Обрабатываем каждую подтаблицу
        for subtable in self.subtables:
            # Получаем строковое представление подтаблицы
            subtable_lines = subtable.to_string(common_header, self.note).split('\n')

            # Для каждой линии в подтаблице
            current_block = []
            for line in subtable_lines:
                # Если это начало новой строки (заголовок строки)
                if line.startswith('### Строка'):
                    # Если у нас уже есть блок данных, добавляем его с дополнительным переносом
                    if current_block:
                        final_result.append('\n'.join(current_block))
                        final_result.append("")  # Дополнительный перенос между строками

                    # Начинаем новый блок с заголовка таблицы
                    current_block = []
                    if table_header:
                        current_block.append(table_header)

                    # Если у подтаблицы есть заголовок, добавляем его
                    if subtable.title:
                        current_block.append(f"## {subtable.title}")

                    # Добавляем заголовок строки
                    current_block.append(line)
                else:
                    # Добавляем данные строки
                    current_block.append(line)

            # Добавляем последний блок, если он есть
            if current_block:
                final_result.append('\n'.join(current_block))
                final_result.append("")  # Дополнительный перенос между блоками

        return '\n'.join(final_result)

    def apply(self, func: Callable[[str], str]) -> None:
        """
        Применяет функцию ко всем элементам таблицы.

        Args:
            func (Callable[[str], str]): Функция для применения к текстовым элементам.
        """
        if self.title:
            self.title = func(self.title)
        self.note = func(self.note)
        self.index = [func(idx) for idx in self.index]
        for tag in self.classified_tags:
            tag.name = func(tag.name)
            tag.value = func(tag.value)
        for header in self.headers:
            header.apply(func)
        for subtable in self.subtables:
            subtable.apply(func)

    def to_dict(self) -> dict[str, Any]:
        """
        Преобразует таблицу в словарь.

        Returns:
            dict[str, Any]: Словарное представление таблицы.
        """
        result = {
            'title': self.title,
            'note': self.note,
            'classified_tags': [tag.to_dict() for tag in self.classified_tags],
            'index': self.index,
            'headers': [header.to_dict() for header in self.headers],
            'subtables': [subtable.to_dict() for subtable in self.subtables],
            'table_style': self.table_style,
            'page_number': self.page_number,
            'index_in_document': self.index_in_document,
            'title_index_in_paragraphs': self.title_index_in_paragraphs,
        }
        return result

    def has_merged_cells(self) -> bool:
        """
        Проверяет наличие объединенных ячеек в таблице.

        Returns:
            bool: True, если в таблице есть строки с разным количеством ячеек.
        """
        # Проверяем заголовки
        if self.headers:
            header_cell_counts = [len(header.cells) for header in self.headers]
            if len(set(header_cell_counts)) > 1:
                return True

            expected_cell_count = header_cell_counts[0] if header_cell_counts else 0
        else:
            expected_cell_count = 0

        # Проверяем подтаблицы
        for subtable in self.subtables:
            if subtable.has_merged_cells():
                return True

            # Проверяем соответствие количества ячеек заголовку
            if subtable.rows and expected_cell_count > 0:
                for row in subtable.rows:
                    if len(row.cells) != expected_cell_count:
                        return True

        return False

    def to_pandas(self, merged_ok: bool = False) -> Optional['pandas.DataFrame']:  # type: ignore
        """
        Преобразует таблицу в pandas DataFrame.

        Args:
            merged_ok (bool): Флаг, указывающий, допустимы ли объединенные ячейки.
                Если False и обнаружены объединенные ячейки, будет выдано предупреждение.

        Returns:
            pandas.DataFrame: DataFrame, представляющий таблицу.

        Примечание:
            Этот метод требует установленного пакета pandas.
        """
        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "Для использования to_pandas требуется установить pandas."
            )

        # Проверка объединенных ячеек
        if not merged_ok and self.has_merged_cells():
            warnings.warn(
                "Таблица содержит объединенные ячейки, что может привести к некорректному "
                "отображению в DataFrame. Установите параметр merged_ok=True, чтобы скрыть это предупреждение."
            )

        # Собираем данные для DataFrame
        data = []

        # Заголовки столбцов
        columns = []
        if self.headers:
            # Объединяем многострочные заголовки, используя разделитель '->'
            if len(self.headers) > 1:
                # Собираем все строки заголовков
                header_cells = []
                for i in range(len(self.headers[0].cells)):
                    header_values = [
                        header.cells[i] if i < len(header.cells) else ""
                        for header in self.headers
                    ]
                    header_cells.append(" -> ".join(filter(None, header_values)))
                columns = header_cells
            else:
                columns = self.headers[0].cells

        # Собираем данные из подтаблиц
        for subtable in self.subtables:
            # Если есть заголовок подтаблицы, добавляем его как строку с пустыми значениями
            if subtable.title:
                row_data = (
                    [subtable.title] + [""] * (len(columns) - 1)
                    if columns
                    else [subtable.title]
                )
                data.append(row_data)

            # Добавляем данные из строк подтаблицы
            for row in subtable.rows:
                row_data = row.cells

                # Если количество ячеек не совпадает с количеством столбцов, заполняем пустыми
                if columns and len(row_data) < len(columns):
                    row_data.extend([""] * (len(columns) - len(row_data)))

                data.append(row_data)

        # Создаем DataFrame
        if not columns:
            # Если нет заголовков, определяем максимальное количество столбцов
            max_cols = max([len(row) for row in data]) if data else 0
            df = pd.DataFrame(data)
        else:
            df = pd.DataFrame(data, columns=columns)

        # Добавляем название таблицы как атрибут
        if self.title:
            df.attrs['title'] = self.title

        # Добавляем примечание как атрибут
        if self.note:
            df.attrs['note'] = self.note

        return df

    def to_markdown(self, merged_ok: bool = False) -> str:
        """
        Преобразует таблицу в формат Markdown.

        Args:
            merged_ok (bool): Флаг, указывающий, допустимы ли объединенные ячейки.
                Если False и обнаружены объединенные ячейки, будет выдано предупреждение.

        Returns:
            str: Markdown представление таблицы.
        """
        # Проверка объединенных ячеек
        if not merged_ok and self.has_merged_cells():
            warnings.warn(
                "Таблица содержит объединенные ячейки, что может привести к некорректному "
                "отображению в Markdown. Установите параметр merged_ok=True, чтобы скрыть это предупреждение."
            )

        lines = []

        # Добавляем заголовок таблицы, если он есть
        if self.title:
            lines.append(f"**{self.title}**\n")

        # Если есть заголовок таблицы, используем его
        if self.headers:
            # Берем первую строку заголовка
            header_cells = self.headers[0].cells

            # Формируем строку заголовка
            header_line = "| " + " | ".join(header_cells) + " |"
            lines.append(header_line)

            # Формируем разделительную строку
            separator_line = "| " + " | ".join(["---"] * len(header_cells)) + " |"
            lines.append(separator_line)

            # Если есть дополнительные строки заголовка, добавляем их
            for i in range(1, len(self.headers)):
                subheader_cells = self.headers[i].cells
                if len(subheader_cells) < len(header_cells):
                    subheader_cells.extend(
                        [""] * (len(header_cells) - len(subheader_cells))
                    )
                subheader_line = (
                    "| " + " | ".join(subheader_cells[: len(header_cells)]) + " |"
                )
                lines.append(subheader_line)

            # Обходим подтаблицы
            for subtable in self.subtables:
                # Если есть заголовок подтаблицы, добавляем его как строку
                if subtable.title:
                    lines.append(
                        f"| **{subtable.title}** | "
                        + " | ".join([""] * (len(header_cells) - 1))
                        + " |"
                    )

                # Добавляем строки подтаблицы
                for row in subtable.rows:
                    row_cells = row.cells

                    # Если количество ячеек не совпадает с количеством заголовков, добавляем пустые
                    if len(row_cells) < len(header_cells):
                        row_cells.extend([""] * (len(header_cells) - len(row_cells)))

                    row_line = "| " + " | ".join(row_cells[: len(header_cells)]) + " |"
                    lines.append(row_line)
        else:
            # Если заголовка нет, просто выводим строки как текст
            for subtable in self.subtables:
                if subtable.title:
                    lines.append(f"**{subtable.title}**")

                for row in subtable.rows:
                    lines.append(row.to_string())

        # Добавляем примечание, если оно есть
        if self.note:
            lines.append(f"\n*Примечание: {self.note}*")

        return "\n".join(lines)