#!/usr/bin/env python3

import os
import re
from pathlib import Path
from urllib.parse import quote
from xml.sax.saxutils import escape

# ====== CONFIGURATION ======
BASE_URL = "https://www.cdn.ay1.net"
FILES_LIST = "files.txt"
SITEMAP_BASENAME = "sitemap"
MAX_URLS_PER_SITEMAP = 10000
MAX_SITEMAP_BYTES = 50 * 1024 * 1024

# The directory containing this script.
ROOT = Path(__file__).resolve().parent
INPUT = ROOT / FILES_LIST
SITEMAPS_DIR = ROOT
# This command generates files.txt while intentionally including directories
# and make-sitemaps.py.txt, but no other files directly under ./:
#
# find -L . -mindepth 1 \( -type d -o \( -type f \( -path './make-sitemaps.py.txt' -o -path './*/*' \) \) \) ! -name 'sitemap-*.xml' ! -name 'files.txt' ! -name 'du.txt' ! -name 'robots.txt' ! -name 'DMCA-take-down-requests.txt' ! -name '50GB-zero-filled-file-test.img' -print | LC_ALL=C sort -u > files.txt
#
# du -Lh . > du.txt
# ===========================

def chunks(items, size):
    """Yield successive chunks containing no more than size items."""
    for start in range(0, len(items), size):
        yield items[start:start + size]


def temporary_path(final_path):
    """Return a temporary filename in the same directory."""
    return final_path.with_name(f".{final_path.name}.tmp")


def write_sitemap(final_path, path_chunk):
    """Write one sitemap atomically."""
    temp_path = temporary_path(final_path)
    try:
        with temp_path.open("w", encoding="utf-8", newline="\n") as f:
            f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
            f.write(
                '<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">\n'
            )

            for rel_path in path_chunk:
                encoded_path = quote(rel_path.lstrip("/"), safe="/")
                url = BASE_URL.rstrip("/") + "/" + encoded_path
                f.write("  <url>\n")
                f.write(f"    <loc>{escape(url)}</loc>\n")
                f.write("  </url>\n")

            f.write("</urlset>\n")

        sitemap_size = temp_path.stat().st_size

        if sitemap_size > MAX_SITEMAP_BYTES:
            raise RuntimeError(
                f"{final_path.name} would be {sitemap_size} bytes, "
                f"which exceeds the 50 MiB sitemap limit"
            )

        os.replace(temp_path, final_path)
    finally:
        try:
            temp_path.unlink()
        except FileNotFoundError:
            pass


def write_sitemap_index(final_path, sitemap_names):
    """Write the sitemap index atomically."""
    temp_path = temporary_path(final_path)
    try:
        with temp_path.open("w", encoding="utf-8", newline="\n") as f:
            f.write('<?xml version="1.0" encoding="UTF-8"?>\n')
            f.write(
                '<sitemapindex '
                'xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">\n'
            )

            for sitemap_name in sitemap_names:
                loc = BASE_URL.rstrip("/") + "/" + sitemap_name
                f.write("  <sitemap>\n")
                f.write(f"    <loc>{escape(loc)}</loc>\n")
                f.write("  </sitemap>\n")

            f.write("</sitemapindex>\n")

        os.replace(temp_path, final_path)

    finally:
        try:
            temp_path.unlink()
        except FileNotFoundError:
            pass


if not INPUT.is_file():
    raise SystemExit(f"ERROR: Input file does not exist: {INPUT}")

paths = []
with INPUT.open("r", encoding="utf-8") as f:
    for line in f:
        # Remove only the line ending. Do not remove spaces that may
        # legitimately be part of a filename.
        path = line.rstrip("\r\n")

        if not path:
            continue

        if path.startswith("./"):
            path = path[2:]

        # Directory URLs are canonical with a trailing slash. Using the final
        # URL directly avoids an unnecessary redirect when search engines crawl
        # entries from the sitemap.
        if (ROOT / path).is_dir() and not path.endswith("/"):
            path += "/"

        paths.append(path)

total = len(paths)
print(f"Found {total} paths")
# An empty files.txt could indicate that find failed. Do not replace valid
# sitemaps with an empty sitemap index in that situation.
if total == 0:
    raise SystemExit(
        "ERROR: files.txt contains no paths; existing sitemaps were not changed"
    )

out_files = []

for number, path_chunk in enumerate(
    chunks(paths, MAX_URLS_PER_SITEMAP),
    start=1
):
    filename = f"{SITEMAP_BASENAME}-{number}.xml"
    output_path = SITEMAPS_DIR / filename
    write_sitemap(output_path, path_chunk)
    out_files.append(filename)

    print(f"Wrote {filename} with {len(path_chunk)} URLs")

index_name = f"{SITEMAP_BASENAME}-index.xml"
index_path = SITEMAPS_DIR / index_name

write_sitemap_index(index_path, out_files)

print(
    f"Wrote {index_name} referencing "
    f"{len(out_files)} sitemap files"
)
# Delete obsolete numbered sitemap files. For example, if the site once
# required sitemap-3.xml but now only requires sitemap-1.xml and sitemap-2.xml.
expected_files = set(out_files)
numbered_sitemap = re.compile(
    rf"^{re.escape(SITEMAP_BASENAME)}-[0-9]+\.xml$"
)
for old_path in SITEMAPS_DIR.glob(f"{SITEMAP_BASENAME}-*.xml"):
    if (
        numbered_sitemap.fullmatch(old_path.name)
        and old_path.name not in expected_files
    ):
        old_path.unlink()
        print(f"Removed obsolete sitemap file: {old_path.name}")