import streamlit as st
import os
import docx
import pdfplumber
import pytesseract
from pdf2image import convert_from_path
import tempfile
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
import io
from PIL import Image

# Устанавливаем путь к исполняемому файлу Tesseract OCR
pytesseract.pytesseract.tesseract_cmd = r"/usr/bin/tesseract"  # Укажите путь к tesseract в вашей системе

def safe_extract_text(page, page_num):
    try:
        text = page.extract_text()
        if not text:
            raise ValueError("Extracted text is empty")
        return text
    except Exception as e:
        st.warning(f"Не удалось извлечь текст со страницы {page_num + 1} обычным способом. Применяем OCR.")
        try:
            # Конвертируем страницу в изображение
            img = page.to_image()
            # Применяем OCR
            text = pytesseract.image_to_string(img.original, lang='rus')
            return text
        except Exception as ocr_error:
            return f"ERROR: Не удалось извлечь текст даже с помощью OCR. {str(ocr_error)}"

def split_into_chapters(pdf_file, output_folder):
    try:
        os.makedirs(output_folder, exist_ok=True)
        chapters = []
        chapter_titles = []
        
        with pdfplumber.open(pdf_file) as pdf:
            total_pages = len(pdf.pages)
            current_chapter = []
            chapter_number = 1
            
            with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
                future_to_page = {executor.submit(safe_extract_text, pdf.pages[i], i): i for i in range(total_pages)}
                
                for future in as_completed(future_to_page):
                    page_num = future_to_page[future]
                    page_text = future.result()
                    
                    if page_text.startswith("ERROR:"):
                        st.warning(f"Страница {page_num + 1}: {page_text}")
                        continue
                    
                    # Проверяем, является ли текст заголовком главы
                    if re.search(r'^(глава|раздел)\s+\d+', page_text.lower(), re.IGNORECASE):
                        if current_chapter:
                            chapter_file = os.path.join(output_folder, f"chapter_{chapter_number}.txt")
                            with open(chapter_file, 'w', encoding='utf-8') as f:
                                f.write('\n'.join(current_chapter))
                            chapters.append(chapter_file)
                            chapter_titles.append(f"Глава {chapter_number}")
                            chapter_number += 1
                            current_chapter = []
                    
                    current_chapter.append(page_text)
                    
                    # Обновляем прогресс
                    progress = (page_num + 1) / total_pages
                    st.progress(progress)
                
                # Сохраняем последнюю главу
                if current_chapter:
                    chapter_file = os.path.join(output_folder, f"chapter_{chapter_number}.txt")
                    with open(chapter_file, 'w', encoding='utf-8') as f:
                        f.write('\n'.join(current_chapter))
                    chapters.append(chapter_file)
                    chapter_titles.append(f"Глава {chapter_number}")

        return chapters, chapter_titles
    except Exception as e:
        st.error(f"Ошибка при разделении на главы: {e}")
        return [], []

# Обновленный интерфейс загрузки и обработки файлов
st.header("2. Загрузка и обработка файлов")
uploaded_file = st.file_uploader("Выберите PDF файл для загрузки", type=["pdf"])
if uploaded_file is not None:
    st.write(f"Файл **{uploaded_file.name}** успешно загружен")
    if st.button("Обработать файл"):
        progress_bar = st.progress(0)
        status_text = st.empty()

        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
            tmp_file.write(uploaded_file.getvalue())
            tmp_file_path = tmp_file.name

        output_folder = os.path.join(os.getcwd(), os.path.splitext(uploaded_file.name)[0])

        status_text.text("Разделение на главы...")
        with st.spinner("Разделение на главы..."):
            chapters, chapter_titles = split_into_chapters(tmp_file_path, output_folder)

        if chapters:
            st.success(f"Документ разделен на главы. Сохранено в: {output_folder}")
            st.write("Обзор глав:")

            st.session_state.chapters = chapters
            st.session_state.chapter_titles = chapter_titles

            for i, chapter_file in enumerate(chapters):
                st.markdown(f"### {chapter_titles[i]}")
                with open(chapter_file, 'r', encoding='utf-8') as f:
                    chapter_content = f.read()
                    display_text = (chapter_content[:500] + "...") if len(chapter_content) > 500 else chapter_content
                    st.write(display_text)
        else:
            st.error("Не удалось разделить документ на главы.")

        os.unlink(tmp_file_path)

# ... (остальной код остается без изменений)