import argparse
import io
import math
import ssl
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

from PIL import Image


def zoomify_tiers(width: int, height: int, tile: int):
    """Return tier tile counts and tile index offsets for Zoomify layout."""
    tiers = []
    edge = tile
    while width > edge or height > edge:
        tiers.append((math.ceil(width / edge), math.ceil(height / edge)))
        edge *= 2
    tiers.append((1, 1))
    tiers.reverse()

    offsets = [0]
    for i in range(1, len(tiers)):
        offsets.append(tiers[i - 1][0] * tiers[i - 1][1] + offsets[i - 1])

    return tiers, offsets


def normalize_base_url(raw_url: str) -> str:
    """Turn a mapy.mzk.cz URL into the imageserver base tile URL."""
    parsed = urllib.parse.urlparse(raw_url)
    scheme = parsed.scheme or "https"
    netloc = parsed.netloc
    path = parsed.path

    if netloc == "mapy.mzk.cz":
        netloc = "imageserver.mzk.cz"
        scheme = "https"

    if path.endswith("ImageProperties.xml"):
        path = path[: -len("ImageProperties.xml")]

    if not path.endswith("/"):
        path = f"{path}/"

    if not netloc:
        raise ValueError("URL must include a hostname")

    return urllib.parse.urlunparse((scheme, netloc, path, "", "", ""))


def fetch_image_properties(base_url: str):
    """Fetch Zoomify metadata (width, height, tile size) from the server.

    Falls back to HTTP if HTTPS negotiation fails (common on some mobile envs).
    Returns metadata plus the effective base URL actually used.
    """
    meta_url = urllib.parse.urljoin(base_url, "ImageProperties.xml")
    try:
        with urllib.request.urlopen(meta_url) as resp:
            xml_data = resp.read()
        effective_base = base_url
    except (ssl.SSLError, urllib.error.URLError):
        if base_url.startswith("https://"):
            fallback_base = base_url.replace("https://", "http://", 1)
            fallback_meta = urllib.parse.urljoin(fallback_base, "ImageProperties.xml")
            with urllib.request.urlopen(fallback_meta) as resp:
                xml_data = resp.read()
            effective_base = fallback_base
        else:
            raise

    root = ET.fromstring(xml_data)
    return (
        int(root.attrib["WIDTH"]),
        int(root.attrib["HEIGHT"]),
        int(root.attrib["TILESIZE"]),
        effective_base,
    )


def tile_url(base_url: str, x: int, y: int, z: int, tiles_x: int, offset: int) -> str:
    tile_index = x + y * tiles_x + offset
    group = tile_index // 256
    return f"{base_url}TileGroup{group}/{z}-{x}-{y}.jpg"


def output_name(base_url: str) -> str:
    stem = base_url.rstrip("/").split("/")[-1]
    return f"{stem or 'map'}.png"


def main():
    parser = argparse.ArgumentParser(
        description="Download and stitch a map from mapy.mzk.cz / imageserver.mzk.cz."
    )
    parser.add_argument(
        "url",
        help="Map URL (e.g. https://mapy.mzk.cz/mzk03/001/053/827/2619316961/)",
    )
    parser.add_argument(
        "-o",
        "--output",
        help="Output PNG filename (default: uses last path segment)",
    )
    args = parser.parse_args()

    base_url = normalize_base_url(args.url)
    width, height, tile_size, base_url = fetch_image_properties(base_url)
    tiers, offsets = zoomify_tiers(width, height, tile_size)
    z = len(tiers) - 1
    tiles_x, tiles_y = tiers[z]
    offset = offsets[z]

    canvas = Image.new("RGB", (width, height))

    for y in range(tiles_y):
        for x in range(tiles_x):
            url = tile_url(base_url, x, y, z, tiles_x, offset)
            try:
                with urllib.request.urlopen(url) as resp:
                    data = resp.read()
            except (ssl.SSLError, urllib.error.URLError):
                if base_url.startswith("https://"):
                    base_url = base_url.replace("https://", "http://", 1)
                    url = tile_url(base_url, x, y, z, tiles_x, offset)
                    with urllib.request.urlopen(url) as resp:
                        data = resp.read()
                else:
                    raise
            tile = Image.open(io.BytesIO(data))
            canvas.paste(tile, (x * tile_size, y * tile_size))

    canvas = canvas.crop((0, 0, width, height))
    filename = args.output or output_name(base_url)
    canvas.save(filename)


if __name__ == "__main__":
    main()