#!/usr/bin/env python3
"""JSON-stdin PDF worker for pi-word-pdf. No OCR and no certificate signing."""

from __future__ import annotations

import io
import json
import os
import sys
from pathlib import Path
from typing import Any


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


def require_module(import_name: str, package_name: str | None = None):
    try:
        return __import__(import_name)
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", f"Missing Python package: {package_name or import_name}") from exc


def require_path(value: Any, extensions: tuple[str, ...]) -> Path:
    if not isinstance(value, str) or not value.strip():
        raise WorkerError("INVALID_PATH", "Path is required")
    path = Path(value).resolve()
    if path.suffix.lower() not in extensions:
        raise WorkerError("UNSUPPORTED_EXTENSION", f"Expected one of: {', '.join(extensions)}")
    return path


def reader_for(path: Path, password: str | None = None):
    require_module("pypdf")
    from pypdf import PdfReader

    reader = PdfReader(str(path), strict=False)
    if reader.is_encrypted:
        if not password:
            raise WorkerError("PASSWORD_REQUIRED", "The PDF is encrypted")
        if reader.decrypt(password) == 0:
            raise WorkerError("INVALID_PASSWORD", "The PDF password is incorrect")
    return reader


def clone_writer(reader):
    from pypdf import PdfWriter

    writer = PdfWriter()
    writer.clone_document_from_reader(reader)
    return writer


def write_pdf(writer, output: Path) -> None:
    output.parent.mkdir(parents=False, exist_ok=True)
    with output.open("wb") as handle:
        writer.write(handle)


def require_page_numbers(value: Any, total: int, default: set[int]) -> set[int]:
    if value is None:
        return default
    if not isinstance(value, list):
        raise WorkerError("INVALID_PAGE", "Page selection must be a list")
    if not value:
        return default
    pages: set[int] = set()
    for item in value:
        if isinstance(item, bool):
            raise WorkerError("INVALID_PAGE", "Page number must be an integer")
        try:
            page = int(item)
        except (TypeError, ValueError) as exc:
            raise WorkerError("INVALID_PAGE", "Page number must be an integer") from exc
        if page < 1 or page > total:
            raise WorkerError("INVALID_PAGE", f"Page {page} is outside the document")
        pages.add(page)
    return pages


def require_page_number(value: Any, total: int) -> int:
    if isinstance(value, bool):
        raise WorkerError("INVALID_PAGE", "Page number must be an integer")
    try:
        page = int(value)
    except (TypeError, ValueError) as exc:
        raise WorkerError("INVALID_PAGE", "Page number must be an integer") from exc
    if page < 1 or page > total:
        raise WorkerError("INVALID_PAGE", "Target page is outside the document")
    return page


def validate_pdf(path: Path, password: str | None = None) -> dict[str, Any]:
    if not path.exists() or path.stat().st_size < 8:
        raise WorkerError("INVALID_PDF", "PDF file is missing or empty")
    with path.open("rb") as handle:
        if handle.read(5) != b"%PDF-":
            raise WorkerError("INVALID_PDF", "PDF header is invalid")
    reader = reader_for(path, password)
    metadata = {str(k): str(v) for k, v in (reader.metadata or {}).items() if v is not None}
    return {
        "path": str(path),
        "valid": True,
        "pages": len(reader.pages),
        "encrypted": bool(reader.is_encrypted),
        "metadata": metadata,
        "size": path.stat().st_size,
    }


def metadata_get(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    path = require_path(params.get("inputPath"), (".pdf",))
    return validate_pdf(path, secrets.get("password"))


def metadata_set(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    values = params.get("metadata")
    if not isinstance(values, dict):
        raise WorkerError("INVALID_METADATA", "metadata must be an object")
    reader = reader_for(input_path, secrets.get("password"))
    writer = clone_writer(reader)
    merged = dict(reader.metadata or {})
    for key, value in values.items():
        normalized = str(key) if str(key).startswith("/") else f"/{key}"
        if value is None:
            merged.pop(normalized, None)
        else:
            merged[normalized] = str(value)
    writer.add_metadata(merged)
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def extract_content(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    try:
        import pdfplumber
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", "Missing Python package: pdfplumber") from exc
    input_path = require_path(params.get("inputPath"), (".pdf",))
    include_tables = bool(params.get("tables", True))
    pages: list[dict[str, Any]] = []
    with pdfplumber.open(str(input_path), password=secrets.get("password")) as pdf:
        for number, page in enumerate(pdf.pages, 1):
            item: dict[str, Any] = {"page": number, "text": page.extract_text() or ""}
            if include_tables:
                item["tables"] = page.extract_tables() or []
            pages.append(item)
    return {"path": str(input_path), "pages": pages}


def merge_pdfs(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    from pypdf import PdfWriter

    inputs = params.get("inputPaths")
    if not isinstance(inputs, list) or len(inputs) < 2:
        raise WorkerError("INPUTS_REQUIRED", "At least two input PDFs are required")
    output_path = require_path(params.get("outputPath"), (".pdf",))
    writer = PdfWriter()
    for value in inputs:
        path = require_path(value, (".pdf",))
        reader = reader_for(path, secrets.get("password"))
        for page in reader.pages:
            writer.add_page(page)
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def split_pdf(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    from pypdf import PdfWriter

    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    start = int(params.get("startPage", 1))
    reader = reader_for(input_path, secrets.get("password"))
    end = int(params.get("endPage", len(reader.pages)))
    if start < 1 or end < start or end > len(reader.pages):
        raise WorkerError("INVALID_PAGE_RANGE", "Page range is outside the document")
    writer = PdfWriter()
    for index in range(start - 1, end):
        writer.add_page(reader.pages[index])
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def rotate_pdf(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    angle = int(params.get("angle", 0))
    if angle not in (90, 180, 270):
        raise WorkerError("INVALID_ROTATION", "Rotation must be 90, 180, or 270")
    reader = reader_for(input_path, secrets.get("password"))
    writer = clone_writer(reader)
    page_numbers = require_page_numbers(params.get("pages"), len(writer.pages), set(range(1, len(writer.pages) + 1)))
    for number, page in enumerate(writer.pages, 1):
        if number in page_numbers:
            page.rotate(angle)
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def register_font(font_path: str | None = None) -> str:
    try:
        from reportlab.pdfbase import pdfmetrics
        from reportlab.pdfbase.ttfonts import TTFont
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", "Missing Python package: reportlab") from exc
    candidates = [font_path] if font_path else []
    windir = os.environ.get("WINDIR", r"C:\Windows")
    candidates.extend([
        os.path.join(windir, "Fonts", "msyh.ttc"),
        os.path.join(windir, "Fonts", "simsun.ttc"),
        os.path.join(windir, "Fonts", "simhei.ttf"),
    ])
    for candidate in candidates:
        if candidate and os.path.isfile(candidate):
            try:
                pdfmetrics.registerFont(TTFont("PiWordPdfCJK", candidate, subfontIndex=0))
                return "PiWordPdfCJK"
            except Exception:
                continue
    raise WorkerError("CJK_FONT_MISSING", "No compatible Chinese TrueType font was found")


def make_text_overlay(width: float, height: float, text: str, options: dict[str, Any]):
    try:
        from reportlab.lib.colors import HexColor
        from reportlab.pdfgen import canvas
        from pypdf import PdfReader
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", "Missing Python package: reportlab") from exc
    buffer = io.BytesIO()
    drawing = canvas.Canvas(buffer, pagesize=(width, height))
    font_name = register_font(options.get("fontPath"))
    drawing.setFont(font_name, float(options.get("fontSize", 18)))
    drawing.setFillColor(HexColor(str(options.get("color", "666666"))))
    opacity = float(options.get("opacity", 1.0))
    if hasattr(drawing, "setFillAlpha"):
        drawing.setFillAlpha(max(0.0, min(1.0, opacity)))
    x = float(options.get("x", width / 2))
    y = float(options.get("y", height / 2))
    rotation = float(options.get("rotation", 0))
    drawing.saveState()
    drawing.translate(x, y)
    drawing.rotate(rotation)
    drawing.drawCentredString(0, 0, text)
    drawing.restoreState()
    drawing.save()
    buffer.seek(0)
    return PdfReader(buffer).pages[0]


def apply_text_overlay(params: dict[str, Any], secrets: dict[str, str], watermark: bool) -> dict[str, Any]:
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    text = params.get("text")
    if not isinstance(text, str) or not text:
        raise WorkerError("TEXT_REQUIRED", "Overlay text is required")
    reader = reader_for(input_path, secrets.get("password"))
    writer = clone_writer(reader)
    if watermark:
        selected = require_page_numbers(params.get("pages"), len(writer.pages), set(range(1, len(writer.pages) + 1)))
    else:
        selected = {require_page_number(params.get("page", 1), len(writer.pages))}
    for number, page in enumerate(writer.pages, 1):
        if number not in selected:
            continue
        width = float(page.mediabox.width)
        height = float(page.mediabox.height)
        overlay = make_text_overlay(width, height, text, params)
        page.merge_page(overlay)
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def stamp_image(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    try:
        from reportlab.lib.utils import ImageReader
        from reportlab.pdfgen import canvas
        from pypdf import PdfReader
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", "Missing Python package: reportlab") from exc
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    image_path = require_path(params.get("imagePath"), (".png", ".jpg", ".jpeg"))
    page_number = int(params.get("page", 1))
    reader = reader_for(input_path, secrets.get("password"))
    writer = clone_writer(reader)
    if page_number < 1 or page_number > len(writer.pages):
        raise WorkerError("INVALID_PAGE", "Target page is outside the document")
    page = writer.pages[page_number - 1]
    width, height = float(page.mediabox.width), float(page.mediabox.height)
    buffer = io.BytesIO()
    drawing = canvas.Canvas(buffer, pagesize=(width, height))
    drawing.drawImage(ImageReader(str(image_path)), float(params.get("x", 36)), float(params.get("y", 36)),
                      width=float(params.get("width", 120)), height=float(params.get("height", 60)),
                      preserveAspectRatio=True, mask="auto")
    drawing.save()
    buffer.seek(0)
    page.merge_page(PdfReader(buffer).pages[0])
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def fill_form(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    fields = params.get("fields")
    if not isinstance(fields, dict) or not fields:
        raise WorkerError("FIELDS_REQUIRED", "Form fields are required")
    reader = reader_for(input_path, secrets.get("password"))
    existing = reader.get_fields() or {}
    unknown = sorted(set(fields) - set(existing))
    if unknown:
        raise WorkerError("FORM_FIELD_NOT_FOUND", f"Unknown form fields: {', '.join(unknown)}")
    writer = clone_writer(reader)
    for page in writer.pages:
        writer.update_page_form_field_values(page, {str(k): str(v) for k, v in fields.items()}, auto_regenerate=False)
    try:
        writer.set_need_appearances_writer(True)
    except Exception:
        pass
    write_pdf(writer, output_path)
    return {**validate_pdf(output_path), "fieldsUpdated": sorted(fields)}


def encrypt_pdf(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    password = secrets.get("password")
    if not password:
        raise WorkerError("PASSWORD_REQUIRED", "A password is required")
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    reader = reader_for(input_path)
    writer = clone_writer(reader)
    writer.encrypt(user_password=password, owner_password=secrets.get("ownerPassword") or password, algorithm="AES-256")
    write_pdf(writer, output_path)
    return {"path": str(output_path), "encrypted": True, "pages": len(reader.pages), "size": output_path.stat().st_size}


def decrypt_pdf(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    password = secrets.get("password")
    if not password:
        raise WorkerError("PASSWORD_REQUIRED", "A password is required")
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_path = require_path(params.get("outputPath"), (".pdf",))
    reader = reader_for(input_path, password)
    writer = clone_writer(reader)
    write_pdf(writer, output_path)
    return validate_pdf(output_path)


def render_pdf(params: dict[str, Any], secrets: dict[str, str]) -> dict[str, Any]:
    try:
        import pypdfium2 as pdfium
    except ImportError as exc:
        raise WorkerError("DEPENDENCY_MISSING", "Missing Python package: pypdfium2") from exc
    input_path = require_path(params.get("inputPath"), (".pdf",))
    output_dir = Path(str(params.get("outputDirectory"))).resolve()
    output_dir.mkdir(parents=False, exist_ok=True)
    dpi = max(36, min(600, int(params.get("dpi", 144))))
    document = pdfium.PdfDocument(str(input_path), password=secrets.get("password"))
    selected = [int(value) for value in params.get("pages", [])] if isinstance(params.get("pages"), list) else list(range(1, len(document) + 1))
    paths: list[str] = []
    for number in selected:
        if number < 1 or number > len(document):
            raise WorkerError("INVALID_PAGE", "Render page is outside the document")
        page = document[number - 1]
        bitmap = page.render(scale=dpi / 72)
        output = output_dir / f"page-{number:03d}.png"
        bitmap.to_pil().save(output, format="PNG")
        paths.append(str(output))
        page.close()
    document.close()
    return {"path": str(input_path), "outputDirectory": str(output_dir), "previewPaths": paths, "dpi": dpi}


OPERATIONS = {
    "validate": lambda p, s: validate_pdf(require_path(p.get("inputPath"), (".pdf",)), s.get("password")),
    "metadata_get": metadata_get,
    "metadata_set": metadata_set,
    "extract": extract_content,
    "merge": merge_pdfs,
    "split": split_pdf,
    "rotate": rotate_pdf,
    "watermark": lambda p, s: apply_text_overlay(p, s, True),
    "overlay_text": lambda p, s: apply_text_overlay(p, s, False),
    "stamp_image": stamp_image,
    "fill_form": fill_form,
    "encrypt": encrypt_pdf,
    "decrypt": decrypt_pdf,
    "render": render_pdf,
}


def main() -> None:
    try:
        request = json.load(sys.stdin)
        if not isinstance(request, dict):
            raise WorkerError("INVALID_REQUEST", "Request must be a JSON object")
        operation = request.get("operation")
        if operation not in OPERATIONS:
            raise WorkerError("UNSUPPORTED_OPERATION", "Unsupported PDF operation")
        params = request.get("params") if isinstance(request.get("params"), dict) else {}
        secrets = request.get("secrets") if isinstance(request.get("secrets"), dict) else {}
        result = OPERATIONS[operation](params, secrets)
        print(json.dumps({"ok": True, "result": result}, ensure_ascii=False, separators=(",", ":")))
    except WorkerError as exc:
        print(json.dumps({"ok": False, "error": {"code": exc.code, "message": str(exc)}}, ensure_ascii=False, separators=(",", ":")))
        raise SystemExit(2)
    except Exception:
        print(json.dumps({"ok": False, "error": {"code": "PDF_WORKER_FAILED", "message": "PDF operation failed"}}, separators=(",", ":")))
        raise SystemExit(1)


if __name__ == "__main__":
    main()