extractors.cpp

#include "engine.h"
#include <iostream>
#include <sstream>
#include <regex>
#include <fstream>
#include <cstdio>
#include <memory>
#include <vector>
#include <cctype>
#include <thread>
#include <algorithm>

#include <poppler/cpp/poppler-document.h>
#include <poppler/cpp/poppler-page.h>
#include <poppler/GlobalParams.h>
#include <poppler/PDFDoc.h>
#include <poppler/Outline.h>
#include <poppler/Link.h>
#include <poppler/goo/GooString.h>
#include <poppler/Catalog.h>

void flatten_outline( const std::vector<OutlineItem*>* items, std::vector<OutlineItem*>& flat_list ) {
    if ( !items ) return;
    for ( OutlineItem *item : *items ) {
        flat_list.push_back( item );
        if ( item->hasKids() ) {
            item->open(); 
            if ( item->getKids() ) {
                flatten_outline( item->getKids(), flat_list );
            }
        }
    }
}

bool get_metadata_chapter_bounds( const char* file_path, int target_chapter, int* start_page, int* end_page ) {
    if ( !globalParams ) globalParams = std::make_unique<GlobalParams>();
    
    auto goo_file = std::make_unique<GooString>( file_path );
    auto doc = std::make_unique<PDFDoc>( std::move( goo_file ) );
    
    if ( !doc->isOk() || !doc->getOutline() ) return false;
    
    const auto *items = doc->getOutline()->getItems();
    if ( !items ) return false;
    
    std::vector<OutlineItem*> all_items;
    flatten_outline( items, all_items );
    
    int current_chapter = 0;
    int found_start = -1;
    int found_end = doc->getNumPages();
    std::regex chapter_regex( R"(^\s*(chapter\s+[0-9]+|chapter\s+[ivxlcdm]+|[ivxlcdm]+)\s*$)", std::regex_constants::icase );

    for ( OutlineItem *item : all_items ) {
        const std::vector<Unicode>& title_uni = item->getTitle();
        std::string title_str;
        for ( Unicode u : title_uni ) title_str += (char)( u & 0xFF );
        
        if ( std::regex_match( title_str, chapter_regex ) ) {
            current_chapter++;
            int page_num = -1;
            const LinkAction *action = item->getAction();
            if ( action && action->getKind() == actionGoTo ) {
                const LinkGoTo *goto_action = static_cast<const LinkGoTo*>( action );
                const LinkDest *dest = goto_action->getDest();
                
                std::unique_ptr<LinkDest> resolved_dest;
                if ( !dest ) {
                    const GooString *named = goto_action->getNamedDest();
                    if ( named ) {
                        resolved_dest = doc->getCatalog()->findDest( named );
                        dest = resolved_dest.get();
                    }
                }
                if ( dest ) page_num = dest->isPageRef() ? doc->findPage( dest->getPageRef() ) : dest->getPageNum();
            }
            
            if ( current_chapter == target_chapter && page_num != -1 ) {
                found_start = page_num;
            } else if ( current_chapter == target_chapter + 1 && page_num != -1 ) {
                found_end = page_num - 1; 
                break; 
            }
        }
    }
    
    if ( found_start != -1 ) {
        *start_page = found_start;
        *end_page = found_end;
        return true;
    }
    return false;
}

std::vector<std::string> extract_epub_html_files( const char* file_path ) {
    std::vector<std::string> chapters;
    std::string cmd_list = std::string( "unzip -Z1 \"" ) + file_path + "\" \"*.html\" \"*.xhtml\" \"*.htm\" 2>/dev/null";
    FILE* pipe_list = popen( cmd_list.c_str(), "r" );
    if ( !pipe_list ) return chapters;

    char buf[ 256 ];
    std::vector<std::string> filenames;
    while ( fgets( buf, sizeof( buf ), pipe_list ) ) {
        std::string name = buf;
        name.erase( name.find_last_not_of( " \n\r\t" ) + 1 );
        filenames.push_back( name );
    }
    pclose( pipe_list );

    for ( const auto& fname : filenames ) {
        std::string cmd_extract = std::string( "unzip -p \"" ) + file_path + "\" \"" + fname + "\" 2>/dev/null";
        FILE* pipe_ext = popen( cmd_extract.c_str(), "r" );
        if ( !pipe_ext ) continue;

        char buf2[ 2048 ];
        std::string html;
        while ( fgets( buf2, sizeof( buf2 ), pipe_ext ) != nullptr ) html += buf2;
        pclose( pipe_ext );
        chapters.push_back( html );
    }
    return chapters;
}

extern "C" {

bool load_pdf_session( const char* file_path, int start_page, int end_page ) {
    cleanup_session();
    
    poppler::document* doc = poppler::document::load_from_file( file_path );
    if ( !doc ) return false;

    int total_pages = doc->pages();
    if ( start_page < 1 ) start_page = 1;
    if ( end_page > total_pages || end_page < start_page ) end_page = total_pages;

    for ( int i = start_page - 1; i < end_page; i++ ) {
        poppler::page* p = doc->create_page( i );
        if ( p ) {
            poppler::ustring ustr = p->text();
            std::vector<char> utf8_bytes = ustr.to_utf8();
            std::string text( utf8_bytes.begin(), utf8_bytes.end() );
            
            std::istringstream iss( text );
            std::string word;
            while ( iss >> word ) push_processed_word( word );
            delete p;
        }
    }
    delete doc;
    return true;
}

bool load_pdf_chapter( const char* file_path, int target_chapter ) {
    int start_page = -1;
    int end_page = -1;
    if ( get_metadata_chapter_bounds( file_path, target_chapter, &start_page, &end_page ) ) {
        return load_pdf_session( file_path, start_page, end_page );
    }
    return false;
}

bool load_txt_session( const char* file_path ) {
    cleanup_session();
    std::ifstream file( file_path );
    if ( !file.is_open() ) return false;

    std::string word;
    while ( file >> word ) push_processed_word( word );
    return session_words.size() > 0;
}

bool load_docx_session( const char* file_path ) {
    cleanup_session();
    std::string cmd = std::string( "unzip -p \"" ) + file_path + "\" word/document.xml 2>/dev/null";
    FILE* pipe = popen( cmd.c_str(), "r" );
    if ( !pipe ) return false;

    char buffer[ 1024 ];
    std::string xml_data;
    while ( fgets( buffer, sizeof( buffer ), pipe ) != nullptr ) xml_data += buffer;
    pclose( pipe );

    if ( xml_data.empty() ) return false;

    std::regex xml_tags( "<[^>]+>" );
    std::string plain_text = std::regex_replace( xml_data, xml_tags, " " );

    std::regex html_entities( "&[a-zA-Z0-9#]{1,10};" );
    plain_text = std::regex_replace( plain_text, html_entities, " " );

    std::istringstream iss( plain_text );
    std::string word;
    while ( iss >> word ) push_processed_word( word );
    return session_words.size() > 0;
}

bool load_epub_session( const char* file_path ) {
    cleanup_session();
    std::string cmd = std::string( "unzip -p \"" ) + file_path + "\" \"*.html\" \"*.xhtml\" \"*.htm\" 2>/dev/null";
    FILE* pipe = popen( cmd.c_str(), "r" );
    if ( !pipe ) return false;

    char buffer[ 2048 ];
    std::string html_data;
    while ( fgets( buffer, sizeof( buffer ), pipe ) != nullptr ) html_data += buffer;
    pclose( pipe );

    if ( html_data.empty() ) return false;

    std::regex html_tags( "<[^>]+>" );
    std::string plain_text = std::regex_replace( html_data, html_tags, " " );

    std::regex html_entities( "&[a-zA-Z0-9#]{1,10};" );
    plain_text = std::regex_replace( plain_text, html_entities, " " );

    std::istringstream iss( plain_text );
    std::string word;
    while ( iss >> word ) push_processed_word( word );

    return session_words.size() > 0;
}

int get_pdf_chapter_count( const char* file_path ) {
    if ( !globalParams ) globalParams = std::make_unique<GlobalParams>();
    
    auto goo_file = std::make_unique<GooString>( file_path );
    auto doc = std::make_unique<PDFDoc>( std::move( goo_file ) );
    
    if ( !doc->isOk() || !doc->getOutline() ) return 0;
    
    const auto *items = doc->getOutline()->getItems();
    if ( !items ) return 0;
    
    std::vector<OutlineItem*> all_items;
    flatten_outline( items, all_items );
    
    int count = 0;
    std::regex chapter_regex( R"(^\s*(chapter\s+[0-9]+|chapter\s+[ivxlcdm]+|[ivxlcdm]+)\s*$)", std::regex_constants::icase );

    for ( OutlineItem *item : all_items ) {
        const std::vector<Unicode>& title_uni = item->getTitle();
        std::string title_str;
        for ( Unicode u : title_uni ) title_str += (char)( u & 0xFF );
        
        if ( std::regex_match( title_str, chapter_regex ) ) count++;
    }
    return count;
}

bool compile_txt_concurrent( const char* file_path, const char* out_path ) {
    std::ifstream file( file_path, std::ios::in | std::ios::binary );
    if ( !file ) return false;

    std::string raw_text( ( std::istreambuf_iterator<char>( file ) ), std::istreambuf_iterator<char>() );
    file.close();

    size_t num_threads = std::thread::hardware_concurrency();
    if ( num_threads == 0 ) num_threads = 4;

    std::vector< std::string > chunks;
    size_t chunk_size = raw_text.length() / num_threads;
    size_t start = 0;

    for ( size_t i = 0; i < num_threads; ++i ) {
        if ( start >= raw_text.length() ) break;

        size_t end = start + chunk_size;
        if ( end < raw_text.length() ) {
            while ( end < raw_text.length() && !std::isspace( raw_text[ end ] ) ) {
                end++;
            }
        } else {
            end = raw_text.length();
        }

        chunks.push_back( raw_text.substr( start, end - start ) );
        start = end;
    }

    std::vector< std::string > processed_chunks( chunks.size() );

    {
        ThreadPool pool( chunks.size() );

        for ( size_t i = 0; i < chunks.size(); ++i ) {
            std::string chunk_data = chunks[ i ]; 

            pool.enqueue( [ i, chunk_data, &processed_chunks ] () mutable {
                
                std::regex html_entities( "&[a-zA-Z0-9#]{1,10};" );
                chunk_data = std::regex_replace( chunk_data, html_entities, " " );

                std::replace( chunk_data.begin(), chunk_data.end(), '-', ' ' );

                std::stringstream ss( chunk_data );
                std::string final_chunk;
                std::string temp_word;
                while ( ss >> temp_word ) {
                    final_chunk += temp_word;
                    final_chunk += '\0';
                }

                processed_chunks[ i ] = std::move( final_chunk );
            } );
        }
    }

    std::ofstream out_file( out_path, std::ios::out | std::ios::binary );
    if ( !out_file ) return false;

    uint32_t total_words = 0;
    for ( const std::string& clean_chunk : processed_chunks ) {
        total_words += std::count( clean_chunk.begin(), clean_chunk.end(), '\0' );
    }

    SRHeader header = { 0 };
    header.magic[0] = 'S'; header.magic[1] = 'R'; header.magic[2] = '0'; header.magic[3] = '1';
    header.saved_index = 0;
    header.total_words = total_words;
    header.chapter_count = 0;
    out_file.write( reinterpret_cast<char*>( &header ), sizeof( SRHeader ) );

    for ( const std::string& clean_chunk : processed_chunks ) {
        out_file.write( clean_chunk.c_str(), clean_chunk.length() );
    }

    out_file.close();
    return true;
}

bool compile_epub_concurrent( const char* file_path, const char* out_path ) {
    std::vector<std::string> raw_chapters = extract_epub_html_files( file_path );
    if ( raw_chapters.empty() ) return false;

    std::vector<std::string> processed_chapters( raw_chapters.size() );

    {
        size_t num_threads = std::thread::hardware_concurrency();
        if ( num_threads == 0 ) num_threads = 4;
        ThreadPool pool( num_threads );

        for ( size_t i = 0; i < raw_chapters.size(); ++i ) {
            std::string chapter_html = raw_chapters[ i ];

            pool.enqueue( [ i, chapter_html, &processed_chapters ] () mutable {
                std::regex html_tags( "<[^>]+>" );
                chapter_html = std::regex_replace( chapter_html, html_tags, " " );

                std::regex html_entities( "&[a-zA-Z0-9#]{1,10};" );
                chapter_html = std::regex_replace( chapter_html, html_entities, " " );

                std::replace( chapter_html.begin(), chapter_html.end(), '-', ' ' );

                std::stringstream ss( chapter_html );
                std::string final_chunk;
                std::string temp_word;
                while ( ss >> temp_word ) {
                    final_chunk += temp_word;
                    final_chunk += '\0';
                }
                
                processed_chapters[ i ] = std::move( final_chunk );
            } );
        }
    }

    std::ofstream out_file( out_path, std::ios::out | std::ios::binary );
    if ( !out_file ) return false;

    uint32_t total_words = 0;
    for ( const std::string& clean_chapter : processed_chapters ) {
        total_words += std::count( clean_chapter.begin(), clean_chapter.end(), '\0' );
    }

    SRHeader header = { 0 };
    header.magic[0] = 'S'; header.magic[1] = 'R'; header.magic[2] = '0'; header.magic[3] = '1';
    header.saved_index = 0;
    header.total_words = total_words;
    header.chapter_count = 0;
    out_file.write( reinterpret_cast<char*>( &header ), sizeof( SRHeader ) );

    for ( const std::string& clean_chapter : processed_chapters ) {
        out_file.write( clean_chapter.c_str(), clean_chapter.length() );
    }

    out_file.close();
    return true;
}

} // extern "C"