processor.py

"""
Main OCR Processor
Orchestrates the complete OCR processing pipeline
"""

import time
from datetime import datetime
import tempfile
import os
import fitz  # PyMuPDF
from pathlib import Path
from typing import Optional, Callable
from dataclasses import dataclass
import threading

from logger import ProcessingLogger
from ocr_engine import OCREngine
from text_processor import ContentFormatter
from utils import validate_page_ranges

@dataclass
class ProcessingResult:
    """Container for processing results."""
    content: str
    output_file: Optional[str]
    status: str
    logs: str
    success: bool
    vision_calls_used: int = 0
    processing_time: float = 0.0
    pages_processed: int = 0

class DocumentProcessor:
    """Main document processor orchestrating the OCR pipeline."""

    def __init__(self):
        self.logger = ProcessingLogger()
        self.ocr_engine = OCREngine(self.logger)
        self.content_formatter = ContentFormatter(self.logger)
        # --- Abort state (instance-scoped) ---
        self._abort_flag = threading.Event()

    # ---------- Abort control ----------
    def abort_processing(self) -> None:
        """Signal that processing should be aborted."""
        self._abort_flag.set()
        self.logger.log_section("Processing Aborted")
        self.logger.log_metric("Status", "Aborted by user")

    def clear_abort(self) -> None:
        """Clear the abort flag (start fresh for a new run)."""
        self._abort_flag.clear()

    def is_abort_requested(self) -> bool:
        """Check whether the user has requested an abort."""
        return self._abort_flag.is_set()
    # -----------------------------------

    # Logs as a single string (Textbox-safe)
    def _logs_text(self) -> str:
        logs = self.logger.get_logs()
        if isinstance(logs, list):
            return "\n".join(str(x) for x in logs)
        return str(logs or "")

    def save_output(self, pdf_path: Path, content: str) -> Optional[str]:
        """Save processed content to file."""
        try:
            base_name = pdf_path.stem
            now = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{base_name}_{now}.md"
            temp_dir = tempfile.gettempdir()
            output_path = os.path.join(temp_dir, filename)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(content)
            self.logger.log_success(f"Saved: {filename}")
            return output_path
        except Exception as e:
            self.logger.log_error(f"Save failed: {e}")
            return None
    
    def process_document(self, uploaded_file, page_ranges_str: Optional[str] = None, progress_callback: Optional[Callable] = None) -> ProcessingResult:
        """
        Process a PDF or Markdown document through the processing pipeline.
        """
        start_time = time.time()

        # Fresh run => ensure abort flag is clear
        self.clear_abort()
        
        try:
            if not uploaded_file:
                return ProcessingResult(
                    content="Please upload a PDF, Markdown, or TXT file.",
                    output_file=None,
                    status="No file",
                    logs=self._logs_text(),
                    success=False
                )
            
            file_path = Path(uploaded_file.name)
            self.logger.log_section(f"Processing: {file_path.name}")
            
            if progress_callback:
                progress_callback(f"📄 Processing: {file_path.name}")
            
            # Markdown / Text
            if file_path.suffix.lower() in ['.md', '.markdown', '.txt']:
                if file_path.suffix.lower() == '.txt':
                    return self._process_txt_file(file_path, progress_callback, start_time)
                else:
                    return self._process_markdown_file(file_path, progress_callback, start_time)
            
            # PDF path
            with fitz.open(file_path) as doc:
                total_pages = len(doc)
                self.logger.log_metric("Total pages in document", total_pages)
                
                # Page ranges
                if page_ranges_str and page_ranges_str.strip():
                    is_valid, error_msg, pages_to_process = validate_page_ranges(page_ranges_str, total_pages)
                    if not is_valid:
                        return ProcessingResult(
                            content=f"Invalid page ranges: {error_msg}",
                            output_file=None,
                            status="Invalid page ranges",
                            logs=self._logs_text(),
                            success=False
                        )
                    page_numbers = [p + 1 for p in pages_to_process]  # back to 1-indexed
                    self.logger.log_metric("Pages to process", f"{len(page_numbers)} pages: {page_ranges_str}")
                else:
                    page_numbers = list(range(1, total_pages + 1))
                    self.logger.log_metric("Pages to process", f"All {total_pages} pages")
                
                # Extract text
                page_texts = {}
                for i, page_no in enumerate(page_numbers):
                    if self.is_abort_requested():
                        return ProcessingResult(
                            content="Processing was aborted by user",
                            output_file=None,
                            status="Aborted",
                            logs=self._logs_text(),
                            success=False,
                            vision_calls_used=self.ocr_engine.get_vision_calls_used(),
                            processing_time=time.time() - start_time,
                            pages_processed=i
                        )
                    
                    if progress_callback:
                        progress_callback(f"📖 Processing page {page_no} ({i+1}/{len(page_numbers)})")
                    
                    page = doc[page_no - 1]  # fitz is 0-indexed
                    text = self.ocr_engine.extract_page_text(page, page_no)
                    page_texts[page_no] = text
            
            # Formatting
            self.logger.log_section("Content Formatting")
            if progress_callback:
                progress_callback("📝 Formatting content...")
            
            formatted_pages = []
            document_title = self._extract_document_title(file_path.name)
            
            for page_no in sorted(page_texts.keys()):
                if self.is_abort_requested():
                    return ProcessingResult(
                        content="Processing was aborted by user",
                        output_file=None,
                        status="Aborted",
                        logs=self._logs_text(),
                        success=False,
                        vision_calls_used=self.ocr_engine.get_vision_calls_used(),
                        processing_time=time.time() - start_time,
                        pages_processed=len(page_texts)
                    )
                
                if len(page_texts[page_no].strip()) >= 10:
                    formatted = self.content_formatter.format_content(
                        page_texts[page_no], 
                        page_no, 
                        document_title
                    )
                    formatted_pages.append(formatted)
            
            # Assemble
            self.logger.log_section("Document Assembly")
            header = self.content_formatter.build_document_header(document_title)
            if page_ranges_str and page_ranges_str.strip():
                header += f"\n\n**Pages Processed:** {page_ranges_str}"
            final_content = f"{header}\n\n---\n\n" + "\n\n---\n\n".join(formatted_pages)
            
            # Save
            output_file = self.save_output(file_path, final_content)
            
            # Metrics
            processing_time = time.time() - start_time
            vision_calls = self.ocr_engine.get_vision_calls_used()
            self.logger.log_section("Processing Complete")
            self.logger.log_metric("Processing time", f"{processing_time:.1f}s")
            self.logger.log_metric("Vision calls used", vision_calls)
            self.logger.log_metric("Pages processed", len(formatted_pages))
            self.logger.log_metric("Total words", len(final_content.split()))
            if progress_callback:
                progress_callback("✅ Complete!")
            
            return ProcessingResult(
                content=final_content,
                output_file=output_file,
                status="Complete",
                logs=self._logs_text(),
                success=True,
                vision_calls_used=vision_calls,
                processing_time=processing_time,
                pages_processed=len(formatted_pages)
            )
            
        except Exception as e:
            processing_time = time.time() - start_time
            error_msg = f"Processing error: {str(e)}"
            self.logger.log_error(error_msg)
            return ProcessingResult(
                content=error_msg,
                output_file=None,
                status="Error",
                logs=self._logs_text(),
                success=False,
                vision_calls_used=self.ocr_engine.get_vision_calls_used(),
                processing_time=processing_time,
                pages_processed=0
            )
    
    def _process_markdown_file(self, file_path: Path, progress_callback: Optional[Callable], start_time: float) -> ProcessingResult:
        """Process a markdown file - apply formatting only, no OCR needed."""
        try:
            self.logger.log_section("Markdown Processing")
            if progress_callback:
                progress_callback("📝 Reading markdown file...")
            with open(file_path, 'r', encoding='utf-8') as f:
                markdown_content = f.read()
            self.logger.log_metric("File size", f"{len(markdown_content)} characters")
            self.logger.log_metric("File type", "Markdown")
            document_title = self._extract_document_title(file_path.name)
            self.logger.log_section("Content Formatting")
            MAX_CHARS_PER_CHUNK = 100000  # ~25k tokens
            chunks_processed = 1
            if len(markdown_content) <= MAX_CHARS_PER_CHUNK:
                if progress_callback:
                    progress_callback("🎨 Formatting content...")
                formatted_content = self.content_formatter.format_content(
                    markdown_content, 1, document_title
                )
            else:
                chunks = self._split_text_into_chunks(markdown_content, MAX_CHARS_PER_CHUNK)
                chunks_processed = len(chunks)
                self.logger.log_metric("Chunks to process", chunks_processed)
                formatted_chunks = []
                for i, chunk in enumerate(chunks, 1):
                    if progress_callback:
                        progress_callback(f"🎨 Formatting chunk {i}/{chunks_processed}...")
                    formatted_chunks.append(
                        self.content_formatter.format_content(chunk, i, document_title)
                    )
                formatted_content = "\n\n---\n\n".join(formatted_chunks)
            header = self.content_formatter.build_document_header(document_title)
            final_content = f"{header}\n\n---\n\n{formatted_content}"
            output_file = self.save_output(file_path, final_content)
            processing_time = time.time() - start_time
            self.logger.log_section("Processing Complete")
            self.logger.log_metric("Processing time", f"{processing_time:.1f}s")
            self.logger.log_metric("Total characters", len(final_content))
            if progress_callback:
                progress_callback("✅ Complete!")
            return ProcessingResult(
                content=final_content,
                output_file=output_file,
                status="Complete",
                logs=self._logs_text(),
                success=True,
                vision_calls_used=0,
                processing_time=processing_time,
                pages_processed=chunks_processed
            )
        except Exception as e:
            return ProcessingResult(
                content=f"Markdown processing error: {str(e)}",
                output_file=None,
                status="Error",
                logs=self._logs_text(),
                success=False,
                vision_calls_used=0,
                processing_time=time.time() - start_time,
                pages_processed=0
            )

    def _process_txt_file(self, file_path: Path, progress_callback: Optional[Callable], start_time: float) -> ProcessingResult:
        """Process a text file - apply formatting only, no OCR needed."""
        try:
            self.logger.log_section("Text Processing")
            if progress_callback:
                progress_callback("📝 Reading text file...")
            with open(file_path, 'r', encoding='utf-8') as f:
                text_content = f.read()
            self.logger.log_metric("File size", f"{len(text_content)} characters")
            self.logger.log_metric("File type", "Text")
            document_title = self._extract_document_title(file_path.name)
            self.logger.log_section("Content Formatting")
            MAX_CHARS_PER_CHUNK = 100000
            chunks_processed = 1
            if len(text_content) <= MAX_CHARS_PER_CHUNK:
                if progress_callback:
                    progress_callback("🎨 Formatting content...")
                formatted_content = self.content_formatter.format_content(
                    text_content, 1, document_title
                )
            else:
                chunks = self._split_text_into_chunks(text_content, MAX_CHARS_PER_CHUNK)
                chunks_processed = len(chunks)
                self.logger.log_metric("Chunks to process", chunks_processed)
                formatted_chunks = []
                for i, chunk in enumerate(chunks, 1):
                    if progress_callback:
                        progress_callback(f"🎨 Formatting chunk {i}/{chunks_processed}...")
                    formatted_chunks.append(
                        self.content_formatter.format_content(chunk, i, document_title)
                    )
                formatted_content = "\n\n---\n\n".join(formatted_chunks)
            header = self.content_formatter.build_document_header(document_title)
            final_content = f"{header}\n\n---\n\n{formatted_content}"
            output_file = self.save_output(file_path, final_content)
            processing_time = time.time() - start_time
            self.logger.log_section("Processing Complete")
            self.logger.log_metric("Processing time", f"{processing_time:.1f}s")
            self.logger.log_metric("Total characters", len(final_content))
            if progress_callback:
                progress_callback("✅ Complete!")
            return ProcessingResult(
                content=final_content,
                output_file=output_file,
                status="Complete",
                logs=self._logs_text(),
                success=True,
                vision_calls_used=0,
                processing_time=processing_time,
                pages_processed=chunks_processed
            )
        except Exception as e:
            return ProcessingResult(
                content=f"Text processing error: {str(e)}",
                output_file=None,
                status="Error",
                logs=self._logs_text(),
                success=False,
                vision_calls_used=0,
                processing_time=time.time() - start_time,
                pages_processed=0
            )
    
    def _split_text_into_chunks(self, text: str, max_chunk_size: int) -> list:
        """Split text into chunks at logical boundaries."""
        chunks = []
        paragraphs = text.split('\n\n')
        current_chunk = ""
        for para in paragraphs:
            if len(current_chunk) + len(para) + 2 > max_chunk_size:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                    current_chunk = para
                else:
                    # Break a huge paragraph by sentences
                    sentences = para.replace('. ', '.\n').split('\n')
                    for sent in sentences:
                        if len(current_chunk) + len(sent) + 1 > max_chunk_size:
                            if current_chunk:
                                chunks.append(current_chunk.strip())
                            current_chunk = sent
                        else:
                            current_chunk += (" " + sent) if current_chunk else sent
            else:
                current_chunk += ("\n\n" + para) if current_chunk else para
        if current_chunk:
            chunks.append(current_chunk.strip())
        return chunks
    
    def _extract_document_title(self, filename: str) -> str:
        """Extract a clean document title from filename."""
        title = os.path.splitext(os.path.basename(filename))[0]
        title = title.replace('_', ' ').replace('-', ' ')
        title = ' '.join(word.capitalize() for word in title.split())
        return title if title else "Document"
    
    def add_log_callback(self, callback: Callable[[str], None]) -> None:
        """Add a callback for real-time log updates."""
        self.logger.add_callback(callback)
    
    def clear_logs(self) -> None:
        """Clear all logs and reset counters."""
        self.logger.clear()
        self.ocr_engine.reset_vision_counter()