#!/usr/bin/env python3
"""Download Kramerius page images and assemble them into a PDF.

Default target:
https://www.digitalniknihovna.cz/mzk/view/uuid:0abf59c0-794e-11e4-8ce5-005056827e52

The script uses the same public Kramerius 7 APIs as the web viewer:
1. discover child page UUIDs through the search endpoint
2. download each page from the item image endpoint
3. embed the JPEGs directly into a PDF
"""

from __future__ import annotations

import argparse
import http.cookiejar
import json
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path


DEFAULT_URL = (
    "https://www.digitalniknihovna.cz/mzk/view/"
    "uuid:0abf59c0-794e-11e4-8ce5-005056827e52"
    "?page=uuid:e57e2820-8547-11e4-a321-5ef3fc9bb22f"
)

KRAMERIUS_BASES = {
    "mzk": "https://api.kramerius.mzk.cz",
}

DEFAULT_CLIENT_IDS = {
    "mzk": "www.digitalniknihovna.cz",
}

USER_AGENT = "Mozilla/5.0 (compatible; kramerius-pdf-downloader/1.0)"


@dataclass(frozen=True)
class Page:
    index: int
    pid: str
    label: str


@dataclass(frozen=True)
class JpegInfo:
    width: int
    height: int
    components: int


@dataclass(frozen=True)
class LibraryConfig:
    code: str
    base_url: str
    version: int
    client_id: str | None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download all page images from a Kramerius document and create a PDF."
    )
    parser.add_argument(
        "url",
        nargs="?",
        default=DEFAULT_URL,
        help="Viewer URL, document UUID, or page/document URL from digitalniknihovna.cz.",
    )
    parser.add_argument("--library", default=None, help="Library code, e.g. mzk.")
    parser.add_argument(
        "--base-url",
        default=None,
        help="Kramerius API base URL. Defaults from --library when known.",
    )
    parser.add_argument(
        "--output-dir",
        default="downloaded_pages",
        help="Directory where images will be stored.",
    )
    parser.add_argument(
        "--pdf",
        default="kramerius_download.pdf",
        help="Output PDF path.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Download only the first N pages, useful for testing.",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.2,
        help="Seconds to wait between page downloads.",
    )
    parser.add_argument(
        "--retries",
        type=int,
        default=3,
        help="Download retries per request.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Redownload images that already exist.",
    )
    parser.add_argument(
        "--cookies",
        default=None,
        help=(
            "Cookie file to use for authenticated downloads. Supports Netscape "
            "cookies.txt exports and plain Cookie header text."
        ),
    )
    parser.add_argument(
        "--auth-token",
        default=None,
        help=(
            "Kramerius auth token. The script sends it as "
            "Authorization: Bearer <token>."
        ),
    )
    parser.add_argument(
        "--auth-file",
        default=None,
        help=(
            "File containing either a raw token, an Authorization header, or "
            "copied request headers from browser devtools."
        ),
    )
    parser.add_argument(
        "--client",
        default=None,
        help="Value for the Kramerius Client header.",
    )
    return parser.parse_args()


def parse_viewer_input(value: str) -> tuple[str | None, str]:
    """Return (library_code, document_pid)."""
    parsed = urllib.parse.urlparse(value)

    if parsed.scheme and parsed.netloc:
        path_parts = [part for part in parsed.path.split("/") if part]
        library = path_parts[0] if path_parts else None

        view_uuid = None
        if "view" in path_parts:
            view_index = path_parts.index("view")
            if len(path_parts) > view_index + 1:
                view_uuid = urllib.parse.unquote(path_parts[view_index + 1])

        if view_uuid:
            return library, normalize_pid(view_uuid)

        match = re.search(r"uuid:[0-9a-fA-F-]{36}", value)
        if match:
            return library, normalize_pid(match.group(0))

    return None, normalize_pid(value)


def normalize_pid(value: str) -> str:
    value = urllib.parse.unquote(value.strip())
    match = re.search(r"uuid:[0-9a-fA-F-]{36}", value)
    if match:
        return match.group(0)
    match = re.search(r"[0-9a-fA-F]{8}(?:-[0-9a-fA-F]{4}){3}-[0-9a-fA-F]{12}", value)
    if match:
        return "uuid:" + match.group(0)
    raise ValueError(f"Could not find a UUID in: {value}")


def resolve_library_config(
    opener: urllib.request.OpenerDirector,
    library: str,
    explicit_base_url: str | None,
    retries: int,
) -> LibraryConfig:
    if explicit_base_url:
        return LibraryConfig(
            code=library,
            base_url=explicit_base_url.rstrip("/"),
            version=7,
            client_id=DEFAULT_CLIENT_IDS.get(library),
        )

    if library in KRAMERIUS_BASES:
        return LibraryConfig(
            code=library,
            base_url=KRAMERIUS_BASES[library],
            version=7,
            client_id=DEFAULT_CLIENT_IDS.get(library),
        )

    try:
        config = fetch_digitalniknihovna_config(opener, library, retries)
    except RuntimeError as exc:
        raise SystemExit(
            f"Could not resolve library {library!r} from digitalniknihovna.cz. "
            "Pass --base-url explicitly.\n"
            f"Reason: {exc}"
        ) from exc

    return config


def fetch_digitalniknihovna_config(
    opener: urllib.request.OpenerDirector, library: str, retries: int
) -> LibraryConfig:
    url = "https://www.digitalniknihovna.cz/assets/shared/globals.js"
    data, content_type = request_bytes(opener, url, retries)
    if "javascript" not in content_type.lower() and "text" not in content_type.lower():
        raise RuntimeError(f"Expected JavaScript config, got {content_type or 'unknown'}")

    text = data.decode("utf-8", errors="replace")
    block = find_library_config_block(text, library)
    if block is None:
        raise RuntimeError(f"Library code {library!r} was not found in globals.js")

    base_url = extract_js_string_property(block, "url")
    if not base_url:
        raise RuntimeError(f"Library code {library!r} has no API URL in globals.js")

    version = extract_js_int_property(block, "version") or 7
    client_id = extract_js_string_property(block, "ga4clientId") or DEFAULT_CLIENT_IDS.get(library)
    return LibraryConfig(
        code=library,
        base_url=base_url.rstrip("/"),
        version=version,
        client_id=client_id,
    )


def find_library_config_block(text: str, library: str) -> str | None:
    code_pattern = re.compile(rf"\bcode\s*:\s*['\"]{re.escape(library)}['\"]")
    match = code_pattern.search(text)
    if not match:
        return None

    start = text.rfind("{", 0, match.start())
    if start == -1:
        return None

    depth = 0
    quote: str | None = None
    escaped = False
    for index in range(start, len(text)):
        char = text[index]
        if quote:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == quote:
                quote = None
            continue

        if char in {"'", '"'}:
            quote = char
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]

    return None


def extract_js_string_property(block: str, name: str) -> str | None:
    pattern = re.compile(rf"\b{re.escape(name)}\s*:\s*(['\"])(.*?)\1", re.S)
    match = pattern.search(block)
    return match.group(2) if match else None


def extract_js_int_property(block: str, name: str) -> int | None:
    pattern = re.compile(rf"\b{re.escape(name)}\s*:\s*(\d+)")
    match = pattern.search(block)
    return int(match.group(1)) if match else None


def build_opener(
    cookie_path: str | None,
) -> tuple[urllib.request.OpenerDirector, int, list[str]]:
    if not cookie_path:
        return urllib.request.build_opener(), 0, []

    path = Path(cookie_path)
    if not path.exists():
        raise FileNotFoundError(f"Cookie file not found: {path}")

    cookie_jar = http.cookiejar.MozillaCookieJar()
    cookie_count = 0

    try:
        cookie_jar.load(path, ignore_discard=True, ignore_expires=True)
        cookie_count = len(cookie_jar)
    except (http.cookiejar.LoadError, OSError):
        raw_cookie = parse_raw_cookie_file(path)
        if raw_cookie:
            opener = urllib.request.build_opener(RawCookieProcessor(raw_cookie))
            return opener, len([part for part in raw_cookie.split(";") if part.strip()]), ["*"]
        raise

    domains = sorted({cookie.domain for cookie in cookie_jar})
    return urllib.request.build_opener(urllib.request.HTTPCookieProcessor(cookie_jar)), cookie_count, domains


def cookie_domains_match_host(domains: list[str], host: str) -> bool:
    if "*" in domains:
        return True

    host = host.lower()
    for domain in domains:
        domain = domain.lower().lstrip(".")
        if host == domain or host.endswith("." + domain):
            return True
    return False


def parse_raw_cookie_file(path: Path) -> str:
    lines = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if line.lower().startswith("cookie:"):
            line = line.split(":", 1)[1].strip()
        lines.append(line)
    return "; ".join(lines)


class RawCookieProcessor(urllib.request.BaseHandler):
    def __init__(self, cookie_header: str) -> None:
        self.cookie_header = cookie_header

    def http_request(self, request: urllib.request.Request) -> urllib.request.Request:
        request.add_unredirected_header("Cookie", self.cookie_header)
        return request

    https_request = http_request


class StaticHeaderProcessor(urllib.request.BaseHandler):
    def __init__(self, headers: dict[str, str]) -> None:
        self.headers = headers

    def http_request(self, request: urllib.request.Request) -> urllib.request.Request:
        for key, value in self.headers.items():
            request.add_unredirected_header(key, value)
        return request

    https_request = http_request


def parse_auth_file(path: str) -> dict[str, str]:
    text = Path(path).read_text(encoding="utf-8").strip()
    if not text:
        return {}

    headers: dict[str, str] = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if ":" not in line:
            continue
        key, value = line.split(":", 1)
        key = key.strip()
        value = value.strip()
        if key.lower() in {"authorization", "client"} and value:
            headers[header_display_name(key)] = value

    if headers:
        return headers

    token = text
    if token.lower().startswith("bearer "):
        token = token.split(None, 1)[1].strip()
    return {"Authorization": f"Bearer {token}"}


def header_display_name(key: str) -> str:
    lower = key.lower()
    if lower == "authorization":
        return "Authorization"
    if lower == "client":
        return "Client"
    return key


def auth_headers(auth_token: str | None, auth_file: str | None, client: str | None) -> dict[str, str]:
    headers = parse_auth_file(auth_file) if auth_file else {}

    if auth_token:
        token = auth_token.strip()
        if token.lower().startswith("bearer "):
            headers["Authorization"] = token
        else:
            headers["Authorization"] = f"Bearer {token}"

    if client:
        headers["Client"] = client

    return headers


def install_headers(
    opener: urllib.request.OpenerDirector, headers: dict[str, str]
) -> urllib.request.OpenerDirector:
    if not headers:
        return opener
    opener.add_handler(StaticHeaderProcessor(headers))
    return opener


def request_bytes(
    opener: urllib.request.OpenerDirector, url: str, retries: int
) -> tuple[bytes, str]:
    headers = {"User-Agent": USER_AGENT, "Accept": "*/*"}
    last_error: Exception | None = None

    for attempt in range(1, retries + 1):
        request = urllib.request.Request(url, headers=headers)
        try:
            with opener.open(request, timeout=60) as response:
                content_type = response.headers.get("Content-Type", "")
                return response.read(), content_type
        except (urllib.error.URLError, TimeoutError) as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(2**attempt, 10))

    raise RuntimeError(f"Failed to fetch {url}: {last_error}")


def request_json(opener: urllib.request.OpenerDirector, url: str, retries: int) -> dict:
    data, content_type = request_bytes(opener, url, retries)
    if "json" not in content_type.lower():
        raise RuntimeError(f"Expected JSON from {url}, got {content_type or 'unknown'}")
    return json.loads(data.decode("utf-8"))


def discover_pages(
    opener: urllib.request.OpenerDirector, config: LibraryConfig, document_pid: str, retries: int
) -> list[Page]:
    if config.version >= 7:
        query = {
            "fl": "pid,accessibility,model,title.search,page.type,page.number,page.placement",
            "q": f'own_parent.pid:"{document_pid}"',
            "sort": "rels_ext_index.sort asc",
            "rows": "10000",
            "wt": "json",
        }
        pid_field = "pid"
        label_fields = ("page.number", "title.search")
    else:
        query = {
            "fl": "PID,dostupnost,fedora.model,dc.title",
            "q": f'parent_pid:"{document_pid}"',
            "sort": "rels_ext_index asc",
            "rows": "10000",
            "wt": "json",
        }
        pid_field = "PID"
        label_fields = ("dc.title",)

    url = f"{api_url(config)}/search?{urllib.parse.urlencode(query)}"
    payload = request_json(opener, url, retries)
    docs = payload.get("response", {}).get("docs", [])

    pages: list[Page] = []
    for index, doc in enumerate(docs, start=1):
        pid = doc.get(pid_field)
        if not pid:
            continue
        label = next((doc.get(field) for field in label_fields if doc.get(field)), str(index))
        pages.append(Page(index=index, pid=pid, label=str(label)))

    if not pages:
        raise RuntimeError(f"No pages found for {document_pid}. Search URL was: {url}")

    return pages


def api_url(config: LibraryConfig) -> str:
    if config.version >= 7:
        return f"{config.base_url}/search/api/client/v7.0"
    return f"{config.base_url}/search/api/v5.0"


def page_image_url(config: LibraryConfig, pid: str) -> str:
    encoded_pid = urllib.parse.quote(pid, safe=":")
    if config.version >= 7:
        return f"{api_url(config)}/items/{encoded_pid}/image"
    return f"{api_url(config)}/item/{encoded_pid}/streams/IMG_FULL"


def download_pages(
    opener: urllib.request.OpenerDirector,
    config: LibraryConfig,
    pages: list[Page],
    output_dir: Path,
    retries: int,
    delay: float,
    overwrite: bool,
) -> list[Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []

    for page in pages:
        path = output_dir / f"{page.index:04d}.jpg"
        paths.append(path)

        if path.exists() and not overwrite:
            print(f"[{page.index:03d}/{len(pages):03d}] exists {path}")
            continue

        url = page_image_url(config, page.pid)
        data, content_type = request_bytes(opener, url, retries)
        if "image" not in content_type.lower():
            preview = data[:200].decode("utf-8", errors="replace")
            raise RuntimeError(
                f"Expected image for {page.pid}, got {content_type or 'unknown'}.\n"
                f"Response starts with: {preview}"
            )

        path.write_bytes(data)
        print(f"[{page.index:03d}/{len(pages):03d}] saved {path} ({len(data) // 1024} KiB)")
        if delay > 0:
            time.sleep(delay)

    return paths


def read_jpeg_info(data: bytes) -> JpegInfo:
    if not data.startswith(b"\xff\xd8"):
        raise ValueError("Not a JPEG file")

    offset = 2
    while offset < len(data):
        while offset < len(data) and data[offset] == 0xFF:
            offset += 1
        if offset >= len(data):
            break

        marker = data[offset]
        offset += 1

        if marker in {0xD8, 0xD9}:
            continue
        if offset + 2 > len(data):
            break

        segment_length = int.from_bytes(data[offset : offset + 2], "big")
        if segment_length < 2 or offset + segment_length > len(data):
            break

        if marker in {
            0xC0,
            0xC1,
            0xC2,
            0xC3,
            0xC5,
            0xC6,
            0xC7,
            0xC9,
            0xCA,
            0xCB,
            0xCD,
            0xCE,
            0xCF,
        }:
            if segment_length < 8:
                break
            height = int.from_bytes(data[offset + 3 : offset + 5], "big")
            width = int.from_bytes(data[offset + 5 : offset + 7], "big")
            components = data[offset + 7]
            return JpegInfo(width=width, height=height, components=components)

        offset += segment_length

    raise ValueError("Could not find JPEG dimensions")


def pdf_escape(value: str) -> str:
    return value.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)")


def build_pdf(image_paths: list[Path], pdf_path: Path) -> None:
    if not image_paths:
        raise RuntimeError("No images to write to PDF")

    objects: list[bytes] = []

    def add_object(body: bytes) -> int:
        objects.append(body)
        return len(objects)

    catalog_id = add_object(b"<< /Type /Catalog /Pages 2 0 R >>")
    pages_id = add_object(b"")
    page_ids: list[int] = []

    for index, image_path in enumerate(image_paths, start=1):
        image_data = image_path.read_bytes()
        info = read_jpeg_info(image_data)

        color_space = b"/DeviceRGB"
        decode = b""
        if info.components == 1:
            color_space = b"/DeviceGray"
        elif info.components == 4:
            color_space = b"/DeviceCMYK"
            decode = b" /Decode [1 0 1 0 1 0 1 0]"

        image_id = len(objects) + 2
        content_id = len(objects) + 3
        page_id = add_object(
            (
                f"<< /Type /Page /Parent {pages_id} 0 R "
                f"/MediaBox [0 0 {info.width} {info.height}] "
                f"/Resources << /XObject << /Im{index} {image_id} 0 R >> >> "
                f"/Contents {content_id} 0 R >>"
            ).encode("ascii")
        )
        page_ids.append(page_id)

        image_header = (
            f"<< /Type /XObject /Subtype /Image /Width {info.width} "
            f"/Height {info.height} /ColorSpace {color_space.decode('ascii')} "
            f"/BitsPerComponent 8 /Filter /DCTDecode{decode.decode('ascii')} "
            f"/Length {len(image_data)} >>\nstream\n"
        ).encode("ascii")
        add_object(image_header + image_data + b"\nendstream")

        content = f"q\n{info.width} 0 0 {info.height} 0 0 cm\n/Im{index} Do\nQ\n".encode(
            "ascii"
        )
        add_object(
            b"<< /Length "
            + str(len(content)).encode("ascii")
            + b" >>\nstream\n"
            + content
            + b"endstream"
        )

    kids = " ".join(f"{page_id} 0 R" for page_id in page_ids)
    objects[pages_id - 1] = (
        f"<< /Type /Pages /Count {len(page_ids)} /Kids [{kids}] >>"
    ).encode("ascii")

    metadata = f"<< /Producer ({pdf_escape(Path(sys.argv[0]).name)}) >>".encode("ascii")
    info_id = add_object(metadata)

    pdf_path.parent.mkdir(parents=True, exist_ok=True)
    with pdf_path.open("wb") as handle:
        handle.write(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n")
        offsets = [0]
        for object_id, body in enumerate(objects, start=1):
            offsets.append(handle.tell())
            handle.write(f"{object_id} 0 obj\n".encode("ascii"))
            handle.write(body)
            handle.write(b"\nendobj\n")

        xref_offset = handle.tell()
        handle.write(f"xref\n0 {len(objects) + 1}\n".encode("ascii"))
        handle.write(b"0000000000 65535 f \n")
        for offset in offsets[1:]:
            handle.write(f"{offset:010d} 00000 n \n".encode("ascii"))
        handle.write(
            (
                f"trailer\n<< /Size {len(objects) + 1} "
                f"/Root {catalog_id} 0 R /Info {info_id} 0 R >>\n"
                f"startxref\n{xref_offset}\n%%EOF\n"
            ).encode("ascii")
        )


def main() -> int:
    args = parse_args()
    library_from_url, document_pid = parse_viewer_input(args.url)
    library = args.library or library_from_url or "mzk"
    opener, cookie_count, cookie_domains = build_opener(args.cookies)
    config = resolve_library_config(opener, library, args.base_url, args.retries)

    print(f"Library: {config.code}")
    print(f"API: {config.base_url} (Kramerius {config.version})")
    print(f"Document: {document_pid}")
    client = args.client or config.client_id
    headers = auth_headers(
        args.auth_token,
        args.auth_file,
        client if (args.auth_token or args.auth_file) else None,
    )
    opener = install_headers(opener, headers)
    if args.cookies:
        if cookie_count:
            print(f"Loaded {cookie_count} cookies from {args.cookies}")
            api_host = urllib.parse.urlparse(config.base_url).hostname or ""
            if not cookie_domains_match_host(cookie_domains, api_host):
                print(
                    "Warning: cookie domains do not match the API host "
                    f"{api_host}; these cookies probably will not authenticate downloads."
                )
        else:
            print(f"Warning: loaded 0 cookies from {args.cookies}")
    if "Authorization" in headers:
        print("Using Authorization bearer token")
    if "Client" in headers:
        print(f"Using Client header: {headers['Client']}")
    print("Discovering pages...")

    pages = discover_pages(opener, config, document_pid, args.retries)
    if args.limit:
        pages = pages[: args.limit]
    print(f"Found {len(pages)} pages")

    image_paths = download_pages(
        opener=opener,
        config=config,
        pages=pages,
        output_dir=Path(args.output_dir),
        retries=args.retries,
        delay=args.delay,
        overwrite=args.overwrite,
    )

    print(f"Writing PDF: {args.pdf}")
    build_pdf(image_paths, Path(args.pdf))
    print("Done")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())