from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import fitz  # type: ignore
import pikepdf

from .utils import add_pdf_extension, ensure_parent_dir, resolve_output_path


@dataclass
class EncryptResult:
    output_path: Path
    page_count: int
    mode: str
    permissions: dict[str, bool]


def _to_path(path_like: str) -> Path:
    return Path(path_like).expanduser().resolve()


def _permissions_dict(permissions: Optional[dict]) -> dict[str, bool]:
    return {
        "printing": bool(permissions.get("printing")) if permissions else False,
        "modifying": bool(permissions.get("modifying")) if permissions else False,
        "copying": bool(permissions.get("copying")) if permissions else False,
        "annotating": bool(permissions.get("annotating")) if permissions else False,
        "fillingForms": bool(permissions.get("fillingForms")) if permissions else False,
    }


def encrypt_pdf(
    input_path: str,
    output_path: str,
    user_password: str,
    owner_password: Optional[str] = None,
    permissions: Optional[dict] = None,
    mode: str = "basic",
    dpi: int = 144,
    image_format: str = "png",
    quality: int = 85,
) -> EncryptResult:
    if not user_password:
        raise ValueError("User password is required")

    source_path = _to_path(input_path)
    if not source_path.exists():
        raise FileNotFoundError(f"File not found: {source_path}")
    if source_path.suffix.lower() != ".pdf":
        raise ValueError("Input file must be a PDF")

    perms = _permissions_dict(permissions)

    if mode == "rasterize":
        doc = fitz.open(source_path)
        page_count = doc.page_count
        if page_count == 0:
            raise ValueError("Input PDF is empty")

        new_doc = fitz.open()
        scale = dpi / 72.0
        matrix = fitz.Matrix(scale, scale)

        for page in doc:
            pix = page.get_pixmap(matrix=matrix, alpha=image_format == "png")
            img_bytes = pix.tobytes("png" if image_format == "png" else "jpg", jpg_quality=quality)
            new_page = new_doc.new_page(width=page.rect.width, height=page.rect.height)
            new_page.insert_image(new_page.rect, stream=img_bytes)

        resolved_output_path = add_pdf_extension(resolve_output_path(source_path, output_path))
        ensure_parent_dir(resolved_output_path)

        permission_bits = 0
        if perms["printing"]:
            permission_bits |= fitz.PDF_PERM_PRINT
        if perms["modifying"]:
            permission_bits |= fitz.PDF_PERM_MODIFY
        if perms["copying"]:
            permission_bits |= fitz.PDF_PERM_COPY
        if perms["annotating"]:
            permission_bits |= fitz.PDF_PERM_ANNOTATE
        if perms["fillingForms"]:
            permission_bits |= fitz.PDF_PERM_FILL_FORM

        new_doc.save(
            resolved_output_path,
            encryption=fitz.PDF_ENCRYPT_AES_256,
            owner_pw=owner_password or user_password,
            user_pw=user_password,
            permissions=permission_bits,
        )
        doc.close()
        new_doc.close()

        return EncryptResult(
            output_path=resolved_output_path,
            page_count=page_count,
            mode="rasterize",
            permissions=perms,
        )

    with pikepdf.Pdf.open(source_path) as pdf:
        page_count = len(pdf.pages)
        if page_count == 0:
            raise ValueError("Input PDF is empty")

        resolved_output_path = add_pdf_extension(resolve_output_path(source_path, output_path))
        ensure_parent_dir(resolved_output_path)

        allow = pikepdf.Permissions(
            print_lowres=perms["printing"],
            print_highres=perms["printing"],
            extract=perms["copying"],
            modify_annotation=perms["annotating"],
            modify_form=perms["fillingForms"],
            modify_other=perms["modifying"],
            modify_assembly=perms["modifying"],
            accessibility=True,
        )

        pdf.save(
            resolved_output_path,
            encryption=pikepdf.Encryption(
                owner=owner_password or user_password,
                user=user_password,
                allow=allow,
                R=6,
            ),
        )

    return EncryptResult(
        output_path=resolved_output_path,
        page_count=page_count,
        mode="basic",
        permissions=perms,
    )
