#!/usr/bin/env python3
"""
PDF-Decrypt - headless worker (Windows).

Removes password protection from encrypted PDF files.
Only .pdf inputs are accepted; the original file is never modified.
The decrypted copy is saved next to the original with a configurable suffix.
If the original's folder is read-only, the copy is saved to a writable
fallback location (Desktop, Documents, home, or the temp folder) instead.

Status tokens returned by process_item() (stable identifiers mapped to
localized text in ui.py):
    ok               - decrypted and saved successfully
    not_pdf          - file does not look like a valid PDF
    not_encrypted    - file is a valid PDF but has no encryption
    wrong_password   - password was rejected by the PDF reader
    not_supported    - unsupported file type (non-.pdf)
    error            - unexpected runtime failure
"""

import os
import sys
import tempfile

from pypdf import PdfReader, PdfWriter

IS_WINDOWS = os.name == "nt"

PDF_EXT = (".pdf",)

DEFAULT_POSTFIX = "_decrypted"


def _looks_like_pdf(path):
    """True if the file starts with a PDF magic header."""
    try:
        with open(path, "rb") as fh:
            return b"%PDF-" in fh.read(1024)
    except OSError:
        return False


def _derive_output(src, postfix, used):
    """<stem><postfix>.pdf next to the original; never overwrites an
    existing file (appends ' (2)', ' (3)', ... if needed)."""
    d = os.path.dirname(os.path.abspath(src))
    stem = os.path.splitext(os.path.basename(src))[0]
    base = stem + (postfix or DEFAULT_POSTFIX)
    out = os.path.join(d, base + ".pdf")
    n = 2
    while os.path.exists(out) or out in used:
        out = os.path.join(d, "{0} ({1}).pdf".format(base, n))
        n += 1
    used.add(out)
    return out


def needs_password(src):
    """True only if *src* is a readable, encrypted PDF.

    False for non-PDFs, unreadable files, and unencrypted PDFs. Used by the
    UI to decide whether a password prompt is worth showing and, for a
    'same password for all' batch, which file to validate the password
    against.
    """
    if not _looks_like_pdf(src):
        return False
    try:
        return bool(PdfReader(src).is_encrypted)
    except Exception:
        return False


def check_password(src, password):
    """Test whether *password* opens *src*, writing nothing.

    Returns one of:
      "ok"             - password (or empty password) opens the file
      "not_encrypted"  - file needs no password
      "wrong_password" - password was rejected
      "not_pdf"        - magic header missing
      "error"          - the file could not be read/parsed

    Both "ok" and "not_encrypted" mean 'no further password is needed'.
    """
    if not _looks_like_pdf(src):
        return "not_pdf"
    try:
        reader = PdfReader(src)
    except Exception:
        return "error"
    if not reader.is_encrypted:
        return "not_encrypted"
    try:
        return "ok" if reader.decrypt(password) else "wrong_password"
    except Exception:
        return "error"


def _fallback_dirs():
    """Writable directories to try, in order, when the original's folder
    cannot be written to. Existing directories only; de-duplicated."""
    candidates = []
    profile = os.environ.get("USERPROFILE") or os.path.expanduser("~")
    onedrive = os.environ.get("OneDrive") or os.environ.get("OneDriveConsumer")
    for base in (onedrive, profile):
        if base:
            for sub in ("Desktop", "Documents"):
                candidates.append(os.path.join(base, sub))
    if profile:
        candidates.append(profile)
    candidates.append(tempfile.gettempdir())

    seen = set()
    dirs = []
    for d in candidates:
        if d and d not in seen and os.path.isdir(d):
            seen.add(d)
            dirs.append(d)
    return dirs


def _unique_in_dir(directory, filename):
    """A non-colliding path for *filename* inside *directory* (appends
    ' (2)', ' (3)', ...). Relies on os.path.exists; safe because the worker
    writes files one at a time."""
    stem, ext = os.path.splitext(filename)
    out = os.path.join(directory, filename)
    n = 2
    while os.path.exists(out):
        out = os.path.join(directory, "{0} ({1}){2}".format(stem, n, ext))
        n += 1
    return out


def _try_write(writer, path):
    """Atomically write *writer* to *path* (.part then os.replace).

    Returns True on success, False if the destination could not be written
    (e.g. a read-only folder). A crash mid-write cannot leave a truncated
    file where a good one should be.
    """
    tmp = path + ".part"
    try:
        with open(tmp, "wb") as fh:
            writer.write(fh)
        os.replace(tmp, path)
        return True
    except OSError:
        try:
            os.remove(tmp)
        except OSError:
            pass
        return False


def decrypt_pdf(src, dst, password):
    """Remove encryption from *src* and write the plain PDF.

    Writes to *dst* (next to the original) when possible; if that folder is
    read-only, falls back to the first writable location from
    _fallback_dirs().

    Returns (status, extra):
      ("ok",             path)    decrypted and saved (path = where it landed)
      ("not_pdf",        "")      magic header missing
      ("not_encrypted",  "")      PDF is not encrypted - nothing to do
      ("wrong_password", "")      password was rejected
      ("error",          message) unexpected failure
    """
    if not _looks_like_pdf(src):
        return ("not_pdf", "")

    try:
        reader = PdfReader(src)
    except Exception as exc:
        return ("error", str(exc))

    if not reader.is_encrypted:
        return ("not_encrypted", "")

    try:
        # PdfReader.decrypt() returns:
        #   PasswordType.NOT_DECRYPTED (0) - wrong password
        #   PasswordType.USER_PASSWORD  (1) - user password accepted
        #   PasswordType.OWNER_PASSWORD (2) - owner password accepted
        match = reader.decrypt(password)
    except Exception as exc:
        return ("error", str(exc))

    if not match:
        return ("wrong_password", "")

    try:
        writer = PdfWriter(clone_from=reader)
    except Exception as exc:
        return ("error", str(exc))

    # Preferred location first: right next to the original.
    if _try_write(writer, dst):
        return ("ok", dst)

    # The original's folder is not writable - save the copy elsewhere.
    name = os.path.basename(dst)
    tried = os.path.dirname(os.path.abspath(dst))
    for directory in _fallback_dirs():
        if os.path.abspath(directory) == tried:
            continue  # already tried this folder
        alt = _unique_in_dir(directory, name)
        if _try_write(writer, alt):
            return ("ok", alt)

    return ("error", "could not write the decrypted file to any writable folder")


def process_item(item):
    """Process one {input, output, password} dict.

    Returns a status tuple exactly like the encrypt worker's process_item,
    so the UI pumping loop can treat both apps identically. For "ok", the
    third element is the path the decrypted file actually landed at (which
    may differ from the requested output if a read-only-folder fallback
    kicked in).
    """
    src = item.get("input", "")
    dst = item.get("output", "")
    password = item.get("password", "")
    ext = os.path.splitext(src)[1].lower()

    if ext not in PDF_EXT:
        return ("not_supported", src)

    status, extra = decrypt_pdf(src, dst, password)
    if status == "ok":
        return ("ok", src, extra)
    if status == "error":
        return ("error", src, extra)
    return (status, src)
