#!/usr/bin/env python3

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable
from urllib.parse import urlencode
from urllib.request import urlopen
import shutil
import xml.etree.ElementTree as ET

BASE_URL = "https://media.battlesnake.com/"
S3_NS = {"s3": "http://doc.s3.amazonaws.com/2006-03-01"}
DEFAULT_PREFIXES = ("snakes/heads/", "snakes/tails/")
ALLOWED_EXTENSIONS = {".svg", ".png", ".webp"}

def build_list_url(prefix:str, marker:str|None) -> str:
  query = {"prefix": prefix}
  if marker:
    query["marker"] = marker
  return f"{BASE_URL}?{urlencode(query)}"

def list_keys_for_prefix(prefix:str) -> list[str]:
  keys: list[str] = []
  marker: str | None = None

  while True:
    url = build_list_url(prefix=prefix, marker=marker)
    with urlopen(url) as response:
      xml_bytes = response.read()

    root = ET.fromstring(xml_bytes)
    for key_node in root.findall("s3:Contents/s3:Key", S3_NS):
      key = (key_node.text or "").strip()
      if key and not key.endswith("/"):
        keys.append(key)

    truncated_text = (
      root.findtext("s3:IsTruncated", default="false", namespaces=S3_NS)
      or "false"
    ).lower()
    is_truncated = truncated_text == "true"
    if not is_truncated:
      break

    next_marker = (
      root.findtext("s3:NextMarker", default="", namespaces=S3_NS) or ""
    ).strip()
    if next_marker:
      marker = next_marker
      continue

    last_key = keys[-1] if keys else None
    if not last_key:
      break
    marker = last_key

  return keys

def keep_customization_key(key:str) -> bool:
  if not key.startswith(DEFAULT_PREFIXES):
    return False
  suffix = Path(key).suffix.lower()
  return suffix in ALLOWED_EXTENSIONS

def to_output_path(output_root:Path, key:str) -> Path:
  if key.startswith("snakes/heads/"):
    relative = key.removeprefix("snakes/")
  elif key.startswith("snakes/tails/"):
    relative = key.removeprefix("snakes/")
  else:
    relative = key
  return output_root / relative

def download_file(url:str, output_file:Path) -> None:
  output_file.parent.mkdir(parents=True, exist_ok=True)
  with urlopen(url) as response, output_file.open("wb") as target:
    shutil.copyfileobj(response, target)

def prune_output(output_root:Path, wanted_files:set[Path]) -> int:
  removed = 0
  if not output_root.exists():
    return 0

  for file_path in output_root.rglob("*"):
    if not file_path.is_file():
      continue
    if file_path not in wanted_files:
      file_path.unlink()
      removed += 1

  for directory in sorted(
    (p for p in output_root.rglob("*") if p.is_dir()), reverse=True
  ):
    if any(directory.iterdir()):
      continue
    directory.rmdir()

  return removed

def collect_customization_keys(prefixes:Iterable[str]) -> list[str]:
  all_keys: list[str] = []
  for prefix in prefixes:
    all_keys.extend(list_keys_for_prefix(prefix))
  return [key for key in sorted(set(all_keys)) if keep_customization_key(key)]

def parse_args() -> argparse.Namespace:
  parser = argparse.ArgumentParser(
    description="Download Battlesnake snake customization assets (heads/tails) from media.battlesnake.com",
  )
  parser.add_argument(
    "--output",
    default="data/battlesnake-customizations",
    help="Output directory (default: data/battlesnake-customizations)",
  )
  parser.add_argument(
    "--overwrite",
    action="store_true",
    help="Overwrite files that already exist",
  )
  parser.add_argument(
    "--prune",
    action="store_true",
    help="Delete files in output directory that are not snake customizations",
  )
  return parser.parse_args()

def main() -> None:
  args = parse_args()
  output_root = Path(args.output).resolve()

  keys = collect_customization_keys(DEFAULT_PREFIXES)
  if not keys:
    print("No customization files found.")
    return

  downloaded = 0
  skipped = 0
  wanted_files: set[Path] = set()

  for key in keys:
    file_url = f"{BASE_URL}{key}"
    output_file = to_output_path(output_root, key)
    wanted_files.add(output_file)

    if output_file.exists() and not args.overwrite:
      skipped += 1
      continue

    download_file(file_url, output_file)
    downloaded += 1

  removed = prune_output(output_root, wanted_files) if args.prune else 0

  print(f"Output directory : {output_root}")
  print(f"Files discovered : {len(keys)}")
  print(f"Downloaded     : {downloaded}")
  print(f"Skipped existing : {skipped}")
  if args.prune:
    print(f"Removed non-customization files: {removed}")

if __name__ == "__main__":
  main()