app/ocr.py

"""
OCR / Document Scanner Module — Enhanced
Supports PDF, Images (Tesseract), DOCX, and XLSX.
"""
import os
import logging
import pdfplumber

logger = logging.getLogger("OCR")


class DocumentScanner:
    def __init__(self):
        pass

    def scan_pdf(self, filepath: str) -> str:
        """Extract text from an uploaded PDF file."""
        logger.info(f"Scanning PDF: {filepath}")
        extracted_text = []

        try:
            with pdfplumber.open(filepath) as pdf:
                for page in pdf.pages:
                    text = page.extract_text()
                    if text:
                        extracted_text.append(text)
                    # Also extract tables
                    tables = page.extract_tables()
                    for table in tables:
                        if table:
                            table_text = "\n".join(
                                " | ".join(str(cell or "") for cell in row)
                                for row in table
                            )
                            extracted_text.append(f"\n[Tabela]:\n{table_text}\n")

            return "\n".join(extracted_text)
        except Exception as e:
            logger.error(f"Failed to scan PDF {filepath}: {e}")
            return f"Gabim gjatë leximit të PDF: {str(e)}"

    def scan_image(self, filepath: str) -> str:
        """Scan image via Tesseract OCR."""
        try:
            import pytesseract
            from PIL import Image
            logger.info(f"Scanning Image via OCR: {filepath}")

            img = Image.open(filepath)
            # Pre-processing: convert to grayscale for better OCR
            img = img.convert('L')
            text = pytesseract.image_to_string(img, lang="sqi+eng")
            return text.strip()
        except ImportError:
            logger.error("pytesseract or PIL not installed for image OCR.")
            return "Leximi i imazheve (OCR) nuk është konfiguruar (mungon pytesseract)."
        except Exception as e:
            logger.error(f"Image OCR failed: {e}")
            return ""

    def scan_docx(self, filepath: str) -> str:
        """Extract text from DOCX file."""
        try:
            from docx import Document
            logger.info(f"Scanning DOCX: {filepath}")

            doc = Document(filepath)
            text_parts = []

            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_parts.append(paragraph.text)

            # Also extract text from tables
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join(cell.text.strip() for cell in row.cells if cell.text.strip())
                    if row_text:
                        text_parts.append(row_text)

            return "\n".join(text_parts)
        except ImportError:
            logger.error("python-docx not installed.")
            return "Leximi i DOCX nuk është konfiguruar (mungon python-docx)."
        except Exception as e:
            logger.error(f"DOCX scan failed: {e}")
            return f"Gabim gjatë leximit të DOCX: {str(e)}"

    def scan_xlsx(self, filepath: str) -> str:
        """Extract text from XLSX file."""
        try:
            from openpyxl import load_workbook
            logger.info(f"Scanning XLSX: {filepath}")

            wb = load_workbook(filepath, data_only=True)
            text_parts = []

            for sheet_name in wb.sheetnames:
                ws = wb[sheet_name]
                text_parts.append(f"\n[Fleta: {sheet_name}]")
                for row in ws.iter_rows(values_only=True):
                    row_text = " | ".join(str(cell) if cell is not None else "" for cell in row)
                    if row_text.strip(" |"):
                        text_parts.append(row_text)

            return "\n".join(text_parts)
        except ImportError:
            logger.error("openpyxl not installed.")
            return "Leximi i XLSX nuk është konfiguruar (mungon openpyxl)."
        except Exception as e:
            logger.error(f"XLSX scan failed: {e}")
            return f"Gabim gjatë leximit të XLSX: {str(e)}"

    def scan_document(self, filepath: str) -> str:
        """Determine file type and route to appropriate scanner."""
        if not filepath or not os.path.exists(filepath):
            return ""

        ext = filepath.lower().split('.')[-1]

        if ext == 'pdf':
            return self.scan_pdf(filepath)
        elif ext in ['png', 'jpg', 'jpeg']:
            return self.scan_image(filepath)
        elif ext == 'docx':
            return self.scan_docx(filepath)
        elif ext == 'xlsx':
            return self.scan_xlsx(filepath)
        else:
            return "Kjo formë dokumenti nuk mbështetet. Formate: PDF, PNG, JPG, DOCX, XLSX."