backend/app/services/document_processor.py

"""
document_processor.py

Multi-format document processor for DocWeave.

Migrated from PilotMaster/DocPilot/backend/app/services/ingestion.py.

Removed:
    - process_document()     RAG pipeline entry point
    - add_chunks()           FAISS embedding call
    - TracePilot HTTP callbacks
    - All pilotcore imports

Preserved:
    - All format extractors  (PDF, DOCX, PPTX, TXT, CSV, XLSX, images, code)
    - PyMuPDF → Docling → OCR fallback cascade
    - TextSection dataclass
    - clean_text()
    - detect_section_title()
    - SECTION_TYPES mapping
    - extract_text_sections()   primary public API
    - extract_text()            convenience flat-text API
"""

from dataclasses import dataclass, field
from statistics import median
import logging
import mimetypes
import os
import re
import shutil

import pandas as pd
import pytesseract

# Locate tesseract binary for Linux/Docker and Windows
_tesseract = shutil.which("tesseract")
if not _tesseract and os.name == "nt":
    _tesseract = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
if _tesseract:
    pytesseract.pytesseract.tesseract_cmd = _tesseract

from docx import Document as DocxDocument
from pdf2image import convert_from_path
from PIL import Image
from pypdf import PdfReader

try:
    from docling.document_converter import DocumentConverter
except ImportError:
    DocumentConverter = None

try:
    from pptx import Presentation
except ImportError:
    Presentation = None

try:
    import fitz
except ImportError:
    fitz = None

logger = logging.getLogger(__name__)

SECTION_TYPES = {
    "abstract": "abstract",
    "introduction": "introduction",
    "background": "background",
    "related work": "related_work",
    "literature review": "related_work",
    "methods": "methods",
    "methodology": "methods",
    "experimental setup": "methods",
    "experiments": "experiments",
    "evaluation": "evaluation",
    "results": "results",
    "discussion": "discussion",
    "conclusion": "conclusion",
    "future work": "future_work",
    "limitations": "limitations",
    "references": "references",
    "bibliography": "references",
    "appendix": "appendix",
}


class TextExtractionError(Exception):
    pass


@dataclass
class TextSection:
    text: str
    metadata: dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Text utilities
# ---------------------------------------------------------------------------

def clean_text(text: str) -> str:
    text = (text or "").replace("\x00", "")
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"(\w+)-\s*\n\s*(\w+)", r"\1\2", text)  # PDF hyphenation
    text = re.sub(r"[|]{3,}", "", text)                     # OCR garbage
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

def extraction_quality(sections: list[TextSection]) -> float:
    """
    Estimate whether extracted PDF text is substantial enough
    to be considered usable.

    Returns a score from 0.0 to 1.0.
    """

    if not sections:
        return 0.0

    total_chars = sum(len(section.text) for section in sections)

    if total_chars == 0:
        return 0.0

    text = "\n".join(section.text for section in sections)

    words = re.findall(r"\b\w+\b", text)
    if not words:
        return 0.0

    alphanumeric_chars = sum(
        char.isalnum()
        for char in text
    )

    alphanumeric_ratio = alphanumeric_chars / max(len(text), 1)

    word_count = len(words)

    score = 0.0

    # Amount of actual text
    if total_chars >= 2000:
        score += 0.4
    elif total_chars >= 1000:
        score += 0.25
    elif total_chars >= 500:
        score += 0.1

    # Number of words
    if word_count >= 300:
        score += 0.3
    elif word_count >= 150:
        score += 0.2
    elif word_count >= 75:
        score += 0.1

    # Mostly actual text rather than symbols/noise
    if alphanumeric_ratio >= 0.75:
        score += 0.3
    elif alphanumeric_ratio >= 0.60:
        score += 0.2
    elif alphanumeric_ratio >= 0.45:
        score += 0.1

    return min(score, 1.0)


def detect_type(file_path: str, mime_type: str = None):
    extension = os.path.splitext(file_path)[1].lower()
    detected_mime = mime_type or mimetypes.guess_type(file_path)[0] or ""
    return extension, detected_mime


# ---------------------------------------------------------------------------
# Section title detection (PyMuPDF page dict)
# ---------------------------------------------------------------------------

def detect_section_title(page_dict: dict):
    """
    Return the highest-confidence heading on a page, or None.
    Uses font size ratio, bold flags, text length, and position scoring.
    """
    font_sizes = []

    for block in page_dict.get("blocks", []):
        if block.get("type") != 0:
            continue
        for line in block.get("lines", []):
            for span in line.get("spans", []):
                size = span.get("size")
                if size:
                    font_sizes.append(size)

    if not font_sizes:
        return None

    median_size = median(font_sizes)
    page_height = page_dict.get("height", 1)
    candidates = []

    for block in page_dict.get("blocks", []):
        if block.get("type") != 0:
            continue

        lines = block.get("lines", [])
        if not lines:
            continue

        text_parts = []
        max_score = 0

        for line in lines:
            for span in line.get("spans", []):
                text = (span.get("text") or "").strip()
                if not text:
                    continue

                score = 0
                size = span.get("size", 0)
                flags = span.get("flags", 0)
                font = span.get("font", "")

                ratio = size / median_size if median_size else 1
                if ratio >= 1.4:
                    score += 4
                elif ratio >= 1.1:
                    score += 2

                if flags & 16:
                    score += 3

                if any(t in font.lower() for t in ["bold", "medi"]):
                    score += 2

                if len(text) <= 60:
                    score += 2
                elif len(text) <= 120:
                    score += 1
                else:
                    score -= 2

                if len(text) > 2 and text.isupper():
                    score += 2

                if len(text.split()) > 1 and text.istitle():
                    score += 1

                text_parts.append(text)
                max_score = max(max_score, score)

        candidate_text = " ".join(text_parts).strip()

        if not candidate_text or len(candidate_text.split()) > 12:
            continue

        if len(lines) == 1:
            max_score += 1

        y0 = block.get("bbox", [0, 0, 0, 0])[1]
        if y0 < page_height * 0.15:
            max_score += 1

        if re.fullmatch(r"[\d.]+", candidate_text):
            continue

        candidates.append({"text": candidate_text, "score": max_score})

    if not candidates:
        return None

    candidates.sort(key=lambda x: x["score"], reverse=True)
    filtered = [c for c in candidates if len(c["text"].split()) <= 8]
    best = filtered[0] if filtered else candidates[0]

    return best["text"] if best["score"] >= 8 else None


# ---------------------------------------------------------------------------
# PDF extractors
# ---------------------------------------------------------------------------

def extract_pdf_text_pymupdf(file_path: str) -> list[TextSection]:
    sections = []

    with fitz.open(file_path) as doc:
        for page_index, page in enumerate(doc, start=1):
            page_dict = page.get_text("dict")
            text = clean_text(page.get_text("text"))
            section_title = detect_section_title(page_dict)

            metadata = {"page": page_index}

            if section_title:
                metadata["section_title"] = section_title
                lower = section_title.lower()
                for key, value in SECTION_TYPES.items():
                    if key in lower:
                        metadata["section_type"] = value
                        break

            if text:
                sections.append(TextSection(text=text, metadata={"element_type": "paragraph", **metadata}))

        logger.info("PyMuPDF processed %s pages", doc.page_count)

    return sections


def extract_pdf_text_pypdf(file_path: str) -> list[TextSection]:
    reader = PdfReader(file_path)
    sections = []

    for page_number, page in enumerate(reader.pages, start=1):
        text = clean_text(page.extract_text() or "")
        if text:
            sections.append(TextSection(
                text=text,
                metadata={"page": page_number, "element_type": "paragraph"},
            ))

    logger.info("pypdf processed %s pages", len(reader.pages))
    return sections


def extract_pdf_docling(file_path: str) -> list[TextSection]:
    if DocumentConverter is None:
        return []

    try:
        result = DocumentConverter().convert(file_path)
        text = clean_text(result.document.export_to_markdown())

        if not text:
            return []

        return [TextSection(text=text, metadata={"extractor": "docling", "element_type": "document"})]

    except Exception as e:
        logger.exception("Docling extraction failed: %s", e)
        return []


def extract_pdf_ocr(file_path: str) -> list[TextSection]:
    try:
        images = convert_from_path(
            file_path,
            poppler_path=os.getenv("POPPLER_PATH"),
        )
        sections = []

        for page_number, image in enumerate(images, start=1):
            text = clean_text(pytesseract.image_to_string(image))
            if text:
                sections.append(TextSection(
                    text=text,
                    metadata={"page": page_number, "ocr": True, "element_type": "ocr"},
                ))

        logger.info("OCR processed %s pages", len(images))
        return sections

    except Exception as e:
        logger.exception("OCR failed: %s", e)
        return []


def extract_pdf_sections(file_path: str) -> list[TextSection]:
    """
    Extract PDF text using a quality-aware cascade:

        PyMuPDF → Docling → OCR

    PyMuPDF is fast (< 1s for most PDFs). Docling and OCR are only
    triggered when native text extraction quality is genuinely poor.
    """

    # ------------------------------------------------------------
    # 1. PyMuPDF (fast path — handles most text-based PDFs)
    # ------------------------------------------------------------

    sections = (
        extract_pdf_text_pymupdf(file_path)
        if fitz
        else extract_pdf_text_pypdf(file_path)
    )

    score = extraction_quality(sections)

    logger.info(
        "PyMuPDF extraction quality: %.2f (%s chars)",
        score,
        sum(len(s.text) for s in sections),
    )

    # Lowered from 0.6 to 0.4 — PyMuPDF text with any reasonable
    # content is usually good enough. Only truly broken/scanned
    # PDFs need the heavier extractors.
    if score >= 0.4:
        return sections

    # ------------------------------------------------------------
    # 2. Docling (slower, better for complex layouts)
    # ------------------------------------------------------------

    logger.info(
        "PyMuPDF extraction quality insufficient, trying Docling"
    )

    sections = extract_pdf_docling(file_path)

    score = extraction_quality(sections)

    logger.info(
        "Docling extraction quality: %.2f (%s chars)",
        score,
        sum(len(s.text) for s in sections),
    )

    if score >= 0.3:
        return sections

    # ------------------------------------------------------------
    # 3. OCR (slowest — only for scanned/image PDFs)
    # ------------------------------------------------------------

    logger.info(
        "Docling extraction quality insufficient, triggering OCR"
    )

    sections = extract_pdf_ocr(file_path)

    score = extraction_quality(sections)

    logger.info(
        "OCR extraction quality: %.2f (%s chars)",
        score,
        sum(len(s.text) for s in sections),
    )

    return sections

# ---------------------------------------------------------------------------
# Other format extractors
# ---------------------------------------------------------------------------

def extract_docx_sections(file_path: str) -> list[TextSection]:
    doc = DocxDocument(file_path)
    text = "\n".join(para.text for para in doc.paragraphs)
    return [TextSection(text=clean_text(text), metadata={"element_type": "paragraph"})]


def extract_pptx_sections(file_path: str) -> list[TextSection]:
    if Presentation is None:
        raise TextExtractionError("PPTX extraction dependency is not installed")

    presentation = Presentation(file_path)
    sections = []

    for slide_number, slide in enumerate(presentation.slides, start=1):
        parts = []

        for shape in slide.shapes:
            if hasattr(shape, "text") and shape.text:
                parts.append(shape.text)

            if getattr(shape, "has_table", False):
                for row in shape.table.rows:
                    cells = [cell.text.strip() for cell in row.cells if cell.text.strip()]
                    if cells:
                        parts.append(" | ".join(cells))

        try:
            notes = slide.notes_slide.notes_text_frame.text
            if notes:
                parts.append(notes)
        except Exception:
            pass

        text = clean_text("\n".join(parts))
        if text:
            sections.append(TextSection(
                text=text,
                metadata={"slide": slide_number, "element_type": "slide"},
            ))

    return sections


def extract_txt_sections(file_path: str) -> list[TextSection]:
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        return [TextSection(text=clean_text(f.read()), metadata={"element_type": "paragraph"})]


def extract_csv_sections(file_path: str) -> list[TextSection]:
    return _dataframe_to_sections(pd.read_csv(file_path))


def extract_xlsx_sections(file_path: str) -> list[TextSection]:
    sheets = pd.read_excel(file_path, sheet_name=None)
    sections = []

    for sheet_name, df in sheets.items():
        for section in _dataframe_to_sections(df):
            section.metadata["sheet"] = sheet_name
            sections.append(section)

    return sections


def _dataframe_to_sections(df) -> list[TextSection]:
    sections = []
    df = df.fillna("")

    for row_number, row in df.iterrows():
        parts = [f"{col}: {str(val).strip()}" for col, val in row.items() if str(val).strip()]
        text = clean_text("\n".join(parts))
        if text:
            sections.append(TextSection(
                text=text,
                metadata={"row": int(row_number) + 1, "element_type": "table_row"},
            ))

    return sections


def extract_image_sections(file_path: str) -> list[TextSection]:
    try:
        text = clean_text(pytesseract.image_to_string(Image.open(file_path)))
        return [TextSection(text=text, metadata={"ocr": True, "element_type": "image"})]
    except Exception as e:
        logger.exception("Image OCR failed: %s", e)
        return []


def extract_code_sections(file_path: str) -> list[TextSection]:
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        text = clean_text(f.read())

    if not text:
        return []

    language_map = {
        ".py": "python", ".js": "javascript", ".jsx": "javascript",
        ".ts": "typescript", ".tsx": "typescript", ".java": "java",
        ".cpp": "cpp", ".c": "c", ".h": "c_header", ".go": "go",
        ".rs": "rust", ".json": "json", ".yaml": "yaml", ".yml": "yaml",
        ".sql": "sql", ".css": "css", ".html": "html",
    }

    extension = os.path.splitext(file_path)[1].lower()

    return [TextSection(
        text=text,
        metadata={"element_type": "code", "language": language_map.get(extension, extension.lstrip("."))},
    )]


# ---------------------------------------------------------------------------
# Extractor registry
# ---------------------------------------------------------------------------

EXTRACTORS = {
    ".pdf": extract_pdf_sections,
    ".docx": extract_docx_sections,
    ".pptx": extract_pptx_sections,
    ".txt": extract_txt_sections,
    ".md": extract_txt_sections,
    ".csv": extract_csv_sections,
    ".xlsx": extract_xlsx_sections,
    ".py": extract_code_sections,
    ".js": extract_code_sections,
    ".jsx": extract_code_sections,
    ".ts": extract_code_sections,
    ".tsx": extract_code_sections,
    ".java": extract_code_sections,
    ".cpp": extract_code_sections,
    ".c": extract_code_sections,
    ".h": extract_code_sections,
    ".go": extract_code_sections,
    ".rs": extract_code_sections,
    ".json": extract_code_sections,
    ".yaml": extract_code_sections,
    ".yml": extract_code_sections,
    ".sql": extract_code_sections,
    ".css": extract_code_sections,
    ".html": extract_code_sections,
    ".png": extract_image_sections,
    ".jpg": extract_image_sections,
    ".jpeg": extract_image_sections,
    ".webp": extract_image_sections,
}


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def extract_text_sections(file_path: str, mime_type: str = None) -> list[TextSection]:
    """
    Primary public API.
    Returns a list of TextSection objects extracted from the given file.
    Raises TextExtractionError on unsupported or unreadable files.
    """
    extension, detected_mime = detect_type(file_path, mime_type)
    extractor = EXTRACTORS.get(extension)

    if not extractor:
        raise TextExtractionError(f"Unsupported file type: {extension or detected_mime}")

    logger.info("Extractor: %s  extension=%s  mime=%s", extractor.__name__, extension, detected_mime)

    try:
        sections = extractor(file_path)
    except TextExtractionError:
        raise
    except Exception as e:
        logger.exception("Extraction failed: %s", e)
        raise TextExtractionError("Could not extract text from document") from e

    cleaned = [s for s in (TextSection(text=clean_text(s.text), metadata=s.metadata) for s in sections) if s.text]

    if not cleaned:
        raise TextExtractionError(
            "Could not extract text from PDF" if extension == ".pdf"
            else "Could not extract text from document"
        )

    logger.info("Extracted %s sections, %s chars", len(cleaned), sum(len(s.text) for s in cleaned))
    return cleaned


def extract_text(file_path: str, mime_type: str = None) -> str:
    """Convenience API. Returns all extracted text as a single string."""
    return "\n\n".join(s.text for s in extract_text_sections(file_path, mime_type))