document_loader.py

"""Load supported document formats and convert their content into RAG-ready chunks

The module extracts text and source metadata from PDF, DOCX, HTML, CSV, JSONL,
and plain-text files. Content is normalized, split into overlapping chunks,
and preserved with metadata for retrieval and citation

provides basic FAR section detection to improve regulatory source
identification and traceability
"""

from __future__ import annotations
import csv
import io
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import fitz
from bs4 import BeautifulSoup
from docx import Document
from config import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE

@dataclass
class TextUnit:
    """ extracted text and its source metadata before chunking"""

    text: str
    source_name: str
    source_type: str
    page: int | None = None
    section: str | None = None
    url: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class Chunk:
    """searchable text chunk with source and citation metadata"""

    chunk_id: str
    text: str
    source_name: str
    source_type: str
    page: int | None = None
    section: str | None = None
    url: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def citation(self) -> str:
        """readable citation from available source metadata"""
        bits = [self.source_name]

        if self.section:
            bits.append(self.section)

        if self.page is not None:
            bits.append(f"p. {self.page}")

        if self.chunk_id:
            bits.append(f"chunk {self.chunk_id}")

        return " · ".join(bits)


def _clean(text: str) -> str:
    """Normalize whitespace and remove null characters from extracted text"""
    text = text.replace("\x00", " ")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def _split_text(text: str, chunk_size: int, overlap: int) -> list[str]:
    """Split text into overlapping chunks using natural breakpoints when possible"""
    text = _clean(text)

    if not text:
        return []

    if len(text) <= chunk_size:
        return [text]

    pieces: list[str] = []
    start = 0

    while start < len(text):
        end = min(start + chunk_size, len(text))

        # prefer paragraph or sentence boundaries instead of arbitrary cuts
        if end < len(text):
            preferred = max(
                text.rfind("\n\n", start, end),
                text.rfind(". ", start, end),
                text.rfind("; ", start, end),
            )

            if preferred > start + int(chunk_size * 0.55):
                end = preferred + 1

        part = text[start:end].strip()

        if part:
            pieces.append(part)

        if end >= len(text):
            break

        # keep overlap so nearby context is preserved between chunks
        start = max(end - overlap, start + 1)

    return pieces


def units_to_chunks(
    units: list[TextUnit],
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    overlap: int = DEFAULT_CHUNK_OVERLAP,
) -> list[Chunk]:
    """Convert text units into numbered chunks while preserving source metadata"""
    chunks: list[Chunk] = []
    counter = 1

    for unit in units:
        for piece in _split_text(unit.text, chunk_size, overlap):
            chunks.append(
                Chunk(
                    chunk_id=str(counter),
                    text=piece,
                    source_name=unit.source_name,
                    source_type=unit.source_type,
                    page=unit.page,
                    section=unit.section,
                    url=unit.url,
                    metadata=dict(unit.metadata),
                )
            )
            counter += 1

    return chunks


def _load_pdf(path: Path) -> list[TextUnit]:
    """Extract text from each PDF page while preserving page numbers"""
    units: list[TextUnit] = []
    doc = fitz.open(path)

    for i, page in enumerate(doc):
        text = _clean(page.get_text("text"))

        if text:
            units.append(
                TextUnit(
                    text=text,
                    source_name=path.name,
                    source_type="PDF",
                    page=i + 1,
                )
            )

    doc.close()
    return units


def _load_docx(path: Path) -> list[TextUnit]:
    """Extract and combine non-empty paragraphs from a DOCX document"""
    doc = Document(path)
    paragraphs = [_clean(p.text) for p in doc.paragraphs if _clean(p.text)]

    return [
        TextUnit(
            text="\n\n".join(paragraphs),
            source_name=path.name,
            source_type="DOCX",
        )
    ]


def _load_html(path: Path) -> list[TextUnit]:
    """Extract readable text from HTML while removing non-content elements"""
    html = path.read_text(encoding="utf-8", errors="ignore")
    soup = BeautifulSoup(html, "html.parser")

    # removes elements that do not contribute meaningful document content
    for tag in soup(["script", "style", "noscript"]):
        tag.decompose()

    text = _clean(soup.get_text("\n"))

    return [
        TextUnit(
            text=text,
            source_name=path.name,
            source_type="HTML",
        )
    ]


def _load_csv(path: Path) -> list[TextUnit]:
    """Convert each populated CSV row into a metadata-aware text unit"""
    units: list[TextUnit] = []

    with path.open(
        "r",
        encoding="utf-8-sig",
        errors="ignore",
        newline="",
    ) as f:
        reader = csv.DictReader(f)

        for row_num, row in enumerate(reader, start=2):
            rendered = "\n".join(
                f"{k}: {v}"
                for k, v in row.items()
                if v not in (None, "")
            )

            if rendered.strip():
                units.append(
                    TextUnit(
                        text=rendered,
                        source_name=path.name,
                        source_type="CSV",
                        section=f"row {row_num}",
                    )
                )

    return units


def _load_jsonl(path: Path) -> list[TextUnit]:
    """Load JSONL records while preserving structured source metadata when available"""
    import json

    units: list[TextUnit] = []

    with path.open("r", encoding="utf-8", errors="ignore") as f:
        for row_num, line in enumerate(f, start=1):
            line = line.strip()

            if not line:
                continue

            try:
                obj = json.loads(line)

            # preserves malformed JSONL lines as plain text instead of discarding them
            except json.JSONDecodeError:
                units.append(
                    TextUnit(
                        text=line,
                        source_name=path.name,
                        source_type="JSONL",
                        section=f"line {row_num}",
                    )
                )
                continue

            # restore previously stored chunk/source metadata when present
            if isinstance(obj, dict) and obj.get("text"):
                units.append(
                    TextUnit(
                        text=str(obj["text"]),
                        source_name=str(obj.get("source_name") or path.name),
                        source_type=str(obj.get("source_type") or "JSONL"),
                        page=(
                            int(obj["page"])
                            if obj.get("page") not in (None, "")
                            else None
                        ),
                        section=obj.get("section"),
                        url=obj.get("url"),
                        metadata=obj.get("metadata") or {},
                    )
                )

            # serialize other JSON values into searchable text
            else:
                units.append(
                    TextUnit(
                        text=json.dumps(obj, ensure_ascii=False),
                        source_name=path.name,
                        source_type="JSONL",
                        section=f"line {row_num}",
                    )
                )

    return units


def load_file(path_like: str | Path) -> list[Chunk]:
    """Load a supported file type and return normalized, searchable chunks"""
    path = Path(path_like)
    suffix = path.suffix.lower()

    # route the file to the appropriate format-specific loader
    if suffix == ".pdf":
        units = _load_pdf(path)
    elif suffix == ".docx":
        units = _load_docx(path)
    elif suffix in {".html", ".htm"}:
        units = _load_html(path)
    elif suffix == ".csv":
        units = _load_csv(path)
    elif suffix == ".jsonl":
        units = _load_jsonl(path)
    elif suffix in {".txt", ".md", ".rst", ".json", ".yaml", ".yml"}:
        text = path.read_text(encoding="utf-8", errors="ignore")
        units = [
            TextUnit(
                text=text,
                source_name=path.name,
                source_type=suffix.lstrip(".").upper(),
            )
        ]
    else:
        raise ValueError(f"Unsupported file type: {suffix or 'unknown'}")

    return units_to_chunks(units)


def infer_far_section(text: str) -> str | None:
    """Detect common FAR section, subpart, or part references in text"""
    patterns = [
        r"\b(?:FAR\s+)?(\d{1,2}\.\d{3}(?:-\d+)?)\b",
        r"\b(Subpart\s+\d{1,2}\.\d+)\b",
        r"\b(Part\s+\d{1,2})\b",
    ]

    # return the first recognizable FAR reference
    for pattern in patterns:
        match = re.search(pattern, text, flags=re.IGNORECASE)

        if match:
            value = match.group(1)

            if re.match(r"^\d", value):
                return f"FAR {value}"

            return value

    return None