table_prep.py

from llama_index.core.text_splitter import SentenceSplitter
from llama_index.core import Document
from config import CHUNK_SIZE, CHUNK_OVERLAP
from my_logging import log_message

def create_table_content(table_data):
    """Create formatted content from table data"""
    doc_id = table_data.get('document_id', table_data.get('document', 'Неизвестно'))
    table_num = table_data.get('table_number', 'Неизвестно')
    table_title = table_data.get('table_title', 'Неизвестно')
    section = table_data.get('section', 'Неизвестно')
    
    # Header section
    content = f"Таблица: {table_num}\n"
    content += f"Название: {table_title}\n"
    content += f"Документ: {doc_id}\n"
    content += f"Раздел: {section}\n"
    
    headers = table_data.get('headers', [])
    if headers:
        content += f"\nЗаголовки: {' | '.join(headers)}\n"
    
    # Data section
    if 'data' in table_data and isinstance(table_data['data'], list):
        content += "\nДанные таблицы:\n"
        for row_idx, row in enumerate(table_data['data'], start=1):
            if isinstance(row, dict):
                row_text = " | ".join([f"{k}: {v}" for k, v in row.items() if v])
                content += f"Строка {row_idx}: {row_text}\n"
    
    return content


def chunk_table_document(doc, chunk_size=None, chunk_overlap=None):
    """
    Smart table chunking:
    - Small tables: keep whole
    - Large tables: split by row-blocks, preserve headers in each chunk
    """
    if chunk_size is None:
        chunk_size = CHUNK_SIZE
    if chunk_overlap is None:
        chunk_overlap = CHUNK_OVERLAP
    
    table_num = doc.metadata.get('table_number', 'unknown')
    doc_id = doc.metadata.get('document_id', 'unknown')
    
    # Parse table structure
    lines = doc.text.strip().split('\n')
    
    table_header_lines = []
    data_rows = []
    in_data = False
    
    for line in lines:
        if line.startswith('Данные таблицы:'):
            in_data = True
            table_header_lines.append(line)
        elif in_data and line.startswith('Строка'):
            data_rows.append(line)
        elif not in_data:
            table_header_lines.append(line)
    
    table_header = '\n'.join(table_header_lines) + '\n'
    
    # If no data rows or small table, use standard splitting
    if not data_rows or len(doc.text) < chunk_size * 1.5:
        log_message(f"  📊 Таблица {table_num}: малая, без разбиения")
        return [doc]
    
    # Row-block chunking for large tables
    log_message(f"  📋 Таблица {table_num}: {len(data_rows)} строк → row-block chunking")
    
    header_size = len(table_header)
    available_size = chunk_size - header_size - 100  # Reserve space
    
    text_chunks = []
    current_chunk_rows = []
    current_size = 0
    
    for row in data_rows:
        row_size = len(row) + 1
        
        # Check if adding this row exceeds limit
        if current_size + row_size > available_size and current_chunk_rows:
            # Create chunk with header + rows
            chunk_text = table_header + '\n'.join(current_chunk_rows)
            text_chunks.append(chunk_text)
            
            # Overlap: keep last 2 rows for context continuity
            overlap_count = min(2, len(current_chunk_rows))
            current_chunk_rows = current_chunk_rows[-overlap_count:]
            current_size = sum(len(r) + 1 for r in current_chunk_rows)
        
        current_chunk_rows.append(row)
        current_size += row_size
    
    # Final chunk
    if current_chunk_rows:
        chunk_text = table_header + '\n'.join(current_chunk_rows)
        text_chunks.append(chunk_text)
    
    log_message(f"  ✂️ Таблица {table_num} → {len(text_chunks)} чанков")
    
    # Create Document objects
    chunked_docs = []
    for i, chunk_text in enumerate(text_chunks):
        chunk_metadata = doc.metadata.copy()
        chunk_metadata.update({
            "chunk_id": i,
            "total_chunks": len(text_chunks),
            "chunk_size": len(chunk_text),
            "is_chunked": True
        })
        
        chunked_doc = Document(
            text=chunk_text,
            metadata=chunk_metadata
        )
        chunked_docs.append(chunked_doc)
    
    return chunked_docs


def table_to_document(table_data, document_id=None):
    """Convert table data to Document, with smart chunking if needed"""
    if not isinstance(table_data, dict):
        return []
    
    doc_id = document_id or table_data.get('document_id') or table_data.get('document', 'Неизвестно')
    table_num = table_data.get('table_number', 'Неизвестно')
    table_title = table_data.get('table_title', 'Неизвестно')
    section = table_data.get('section', 'Неизвестно')
    
    table_rows = table_data.get('data', [])
    if not table_rows:
        log_message(f"⚠️ Таблица {table_num} пропущена: нет данных")
        return []
    
    content = create_table_content(table_data)
    content_size = len(content)
    
    base_doc = Document(
        text=content,
        metadata={
            "type": "table",
            "table_number": table_num,
            "table_title": table_title,
            "document_id": doc_id,
            "section": section,
            "section_id": section,
            "total_rows": len(table_rows),
            "content_size": content_size
        }
    )
    
    # Apply smart chunking if too large
    if content_size > CHUNK_SIZE:
        log_message(f"📊 CHUNKING: Таблица {table_num} | {content_size} > {CHUNK_SIZE}")
        return chunk_table_document(base_doc)
    else:
        log_message(f"✓ Таблица {table_num} добавлена целиком ({content_size} символов)")
        return [base_doc]