"""Resource-bounded CV document decoder invoked only as a child process."""

from __future__ import annotations

import argparse
import io
import json
import os
from pathlib import Path
import re
import sys
import zipfile


MAX_EXTRACT_CHARS = 200_000
MAX_PDF_PAGES = 40
MAX_DOCX_ENTRIES = 256
MAX_DOCX_UNCOMPRESSED_BYTES = 32 * 1024 * 1024
MAX_DOCX_ENTRY_BYTES = 8 * 1024 * 1024
MAX_DOCX_COMPRESSION_RATIO = 100
MAX_IMAGE_DIMENSION = 12_000
MAX_IMAGE_PIXELS = 40_000_000
MAX_PDF_PAGE_PIXELS = 12_000_000
MAX_PDF_OCR_PIXELS = 120_000_000
OCR_LANGUAGES = "eng"
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}


class ParserError(Exception):
    def __init__(self, code: str):
        super().__init__(code)
        self.code = code


def _apply_resource_limits(cpu_seconds: int, memory_bytes: int) -> None:
    if os.name != "posix":
        return
    import resource

    resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 1))
    resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes))
    resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024))
    resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))


def normalize_text(value: str) -> str:
    value = value.replace("\x00", " ").replace("\r\n", "\n").replace("\r", "\n")
    lines: list[str] = []
    blank = False
    for raw_line in value.split("\n"):
        line = re.sub(r"[\t\f\v]+", " ", raw_line).strip()
        if not line:
            if lines and not blank:
                lines.append("")
            blank = True
            continue
        lines.append(line)
        blank = False
    return "\n".join(lines).strip()


def _check_extracted_size(text: str, limit: int = MAX_EXTRACT_CHARS) -> str:
    if len(text) > limit:
        raise ParserError("extracted_text_limit")
    return text


def _check_image_size(image, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> None:
    width, height = image.size
    if width <= 0 or height <= 0 or width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION:
        raise ParserError("image_dimensions")
    if width * height > pixel_limit:
        raise ParserError("image_pixel_limit")
    if getattr(image, "n_frames", 1) != 1:
        raise ParserError("image_frames")


def validate_docx_container(
    data: bytes,
    *,
    entry_limit: int = MAX_DOCX_ENTRIES,
    total_limit: int = MAX_DOCX_UNCOMPRESSED_BYTES,
    entry_size_limit: int = MAX_DOCX_ENTRY_BYTES,
    ratio_limit: int = MAX_DOCX_COMPRESSION_RATIO,
) -> None:
    try:
        with zipfile.ZipFile(io.BytesIO(data)) as archive:
            entries = archive.infolist()
            names = {entry.filename for entry in entries}
            if "[Content_Types].xml" not in names or "word/document.xml" not in names:
                raise ParserError("docx_signature")
            if len(entries) > entry_limit:
                raise ParserError("docx_entry_limit")
            total_size = 0
            for entry in entries:
                total_size += entry.file_size
                if entry.file_size > entry_size_limit:
                    raise ParserError("docx_entry_size")
                if entry.file_size > 0 and entry.compress_size == 0:
                    raise ParserError("docx_ratio")
                if entry.compress_size > 0 and entry.file_size / entry.compress_size > ratio_limit:
                    raise ParserError("docx_ratio")
            if total_size > total_limit:
                raise ParserError("docx_expanded_size")
    except ParserError:
        raise
    except (zipfile.BadZipFile, OSError) as exc:
        raise ParserError("docx_signature") from exc


def extract_plain_text(data: bytes) -> str:
    try:
        decoded = data.decode("utf-8-sig")
    except UnicodeDecodeError as exc:
        raise ParserError("text_encoding") from exc
    return _check_extracted_size(normalize_text(decoded))


def extract_docx_text(data: bytes) -> str:
    validate_docx_container(data)
    from docx import Document

    document = Document(io.BytesIO(data))
    parts: list[str] = []
    blocks = document.iter_inner_content() if hasattr(document, "iter_inner_content") else document.paragraphs
    for block in blocks:
        if hasattr(block, "rows"):
            for row in block.rows:
                cells = [cell.text.strip() for cell in row.cells if cell.text and cell.text.strip()]
                if cells:
                    parts.append(" | ".join(cells))
        elif getattr(block, "text", "").strip():
            text = block.text.strip()
            style = (getattr(getattr(block, "style", None), "name", "") or "").lower()
            if "role" in style and parts:
                parts.append("")
            parts.append(f"- {text}" if "bullet" in style else text)
    return _check_extracted_size(normalize_text("\n".join(parts)))


def _ocr_image(image) -> str:
    import pytesseract

    if image.mode not in ("RGB", "L"):
        image = image.convert("RGB")
    return normalize_text(pytesseract.image_to_string(image, lang=OCR_LANGUAGES))


def extract_image_text(data: bytes, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> str:
    from PIL import Image

    with Image.open(io.BytesIO(data)) as image:
        _check_image_size(image, pixel_limit=pixel_limit)
        return _check_extracted_size(_ocr_image(image))


def extract_pdf_text(data: bytes, *, page_limit: int = MAX_PDF_PAGES) -> tuple[str, bool, int]:
    from pypdf import PdfReader

    page_count = 0
    extracted_pages: list[str] = []
    try:
        reader = PdfReader(io.BytesIO(data))
        page_count = len(reader.pages)
        if page_count > page_limit:
            raise ParserError("pdf_page_limit")
        for page in reader.pages:
            try:
                extracted_pages.append(page.extract_text(extraction_mode="layout") or "")
            except TypeError:
                extracted_pages.append(page.extract_text() or "")
    except ParserError:
        raise
    except Exception:
        extracted_pages = []

    text = _check_extracted_size(normalize_text("\n".join(extracted_pages)))
    if len(text) >= 80:
        return text, False, page_count

    import fitz
    from PIL import Image

    document = fitz.open(stream=data, filetype="pdf")
    try:
        page_count = document.page_count
        if page_count > page_limit:
            raise ParserError("pdf_page_limit")
        ocr_pages: list[str] = []
        total_pixels = 0
        for page in document:
            pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
            page_pixels = pixmap.width * pixmap.height
            total_pixels += page_pixels
            if page_pixels > MAX_PDF_PAGE_PIXELS or total_pixels > MAX_PDF_OCR_PIXELS:
                raise ParserError("pdf_ocr_limit")
            with Image.open(io.BytesIO(pixmap.tobytes("png"))) as image:
                _check_image_size(image, pixel_limit=MAX_PDF_PAGE_PIXELS)
                ocr_pages.append(_ocr_image(image))
            _check_extracted_size("\n".join(ocr_pages))
        return _check_extracted_size(normalize_text("\n".join(ocr_pages))), True, page_count
    finally:
        document.close()


def extract(extension: str, data: bytes) -> dict[str, object]:
    if extension in {".txt", ".md"}:
        text = extract_plain_text(data)
        return {"text": text, "ocr_used": False, "page_count": None}
    if extension == ".docx":
        text = extract_docx_text(data)
        return {"text": text, "ocr_used": False, "page_count": None}
    if extension == ".pdf":
        text, ocr_used, page_count = extract_pdf_text(data)
        return {"text": text, "ocr_used": ocr_used, "page_count": page_count}
    if extension in IMAGE_EXTENSIONS:
        text = extract_image_text(data)
        return {"text": text, "ocr_used": True, "page_count": 1}
    raise ParserError("unsupported_type")


def _write_result(path: Path, result: dict[str, object]) -> None:
    path.write_text(json.dumps(result, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--extension", required=True)
    parser.add_argument("--cpu-seconds", type=int, default=20)
    parser.add_argument("--memory-bytes", type=int, default=512 * 1024 * 1024)
    args = parser.parse_args()
    output_path = Path(args.output)
    try:
        _apply_resource_limits(args.cpu_seconds, args.memory_bytes)
        data = Path(args.input).read_bytes()
        result = extract(args.extension.lower(), data)
        if not result.get("text"):
            raise ParserError("no_text")
        _write_result(output_path, {"ok": True, **result})
        return 0
    except ParserError as exc:
        _write_result(output_path, {"ok": False, "code": exc.code})
        return 2
    except BaseException:
        _write_result(output_path, {"ok": False, "code": "parser_failed"})
        return 3


if __name__ == "__main__":
    sys.exit(main())
