#!/usr/bin/env python3
"""Ringside Exchange installer: add shared AEW: Fight Forever creations to your save folder.

File-level only. It copies whole one-wrestler UGCWrestler_*.sav files into the
next free number; it never reads or changes what is inside a save, and it never
overwrites an existing file. Works on Windows, Linux (Proton) and Steam Deck.

  ringside.py where                    show the detected SaveGames folder(s)
  ringside.py list                     list wrestler files and free numbers
  ringside.py backup                   zip the whole SaveGames folder
  ringside.py install FILE.sav [...]   add files into free numbers (backs up first)
  ringside.py get ID [...]             download wrestlers from the sharing site and add them
  ringside.py swap ID                  use a shared tag team or arena file in place of yours (backs up first)
  ringside.py unswap teams|arenas      put your own tag teams or arenas back, touching nothing else
  ringside.py restore [BACKUP.zip]     roll the whole SaveGames folder back to a backup (asks first; newest if none given)

Use --saves PATH to point at a specific folder (for example a test copy).
Close the game before installing. On Linux/Deck run it from Desktop Mode.
"""
import argparse
import datetime
import os
import re
import shutil
import sys
import zipfile
from pathlib import Path

APP_ID = "1913210"
GAME_DIR = "AEWFightForever"
MAX_WRESTLER_FILES = 100      # the game ignores UGCWrestler_101.sav and above (per Nexus mod 351)
STUB_BYTES = 8 * 1024         # deleted CAWs leave ~4 KB files behind; these slots are not reused by us
WRESTLER_RE = re.compile(r"^UGCWrestler_(\d+)\.sav$", re.IGNORECASE)
GVAS_MAGIC = b"GVAS"
SAVE_CLASS = b"/Script/ABP_200508.SaveGame_RawData"
PROTECTED = {"onlineinfo.sav"}  # never copied, restored or overwritten by this tool
DEFAULT_SERVER = "https://ringside.exchange"  # override with --server or RINGSIDE_SERVER


# ------------------------------------------------------------------ locating saves

def _rel_saves():
    return Path("AppData") / "Local" / GAME_DIR / "Saved" / "SaveGames"


def _steam_roots(home):
    roots = [
        home / ".steam" / "steam",
        home / ".local" / "share" / "Steam",
        home / ".var" / "app" / "com.valvesoftware.Steam" / ".local" / "share" / "Steam",  # Flatpak
        home / "snap" / "steam" / "common" / ".local" / "share" / "Steam",
    ]
    seen, out = set(), []
    for r in roots:
        try:
            key = r.resolve()
        except OSError:
            continue
        if r.is_dir() and key not in seen:
            seen.add(key)
            out.append(r)
    return out


def _library_folders(steam_root):
    """Steam library paths from libraryfolders.vdf (covers Steam Deck SD cards)."""
    libs = [steam_root]
    vdf = steam_root / "steamapps" / "libraryfolders.vdf"
    try:
        text = vdf.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return libs
    for m in re.finditer(r'"path"\s+"([^"]+)"', text):
        p = Path(m.group(1).replace("\\\\", "\\"))
        if p.is_dir() and p not in libs:
            libs.append(p)
    return libs


def proton_save_dirs(home):
    """SaveGames folders inside Proton prefixes, across every Steam library."""
    return [lib / "steamapps" / "compatdata" / APP_ID / "pfx" / "drive_c" / "users" / "steamuser" / _rel_saves()
            for root in _steam_roots(home) for lib in _library_folders(root)]


def find_save_dirs():
    found = []
    if os.name == "nt":
        local = os.environ.get("LOCALAPPDATA")
        if local:
            found.append(Path(local) / GAME_DIR / "Saved" / "SaveGames")
    else:
        found += proton_save_dirs(Path.home())
    unique = []
    for p in found:
        if p.is_dir() and p not in unique:
            unique.append(p)
    return unique


def resolve_saves(arg):
    if arg:
        p = Path(arg).expanduser()
        if not p.is_dir():
            sys.exit(f"Not a folder: {p}")
        return p
    dirs = find_save_dirs()
    if not dirs:
        sys.exit("Could not find the AEW: Fight Forever SaveGames folder. "
                 "Launch the game once, or pass --saves PATH.")
    if len(dirs) > 1:
        print(f"Found more than one SaveGames folder; using the first. Add --saves and a path to choose ({help_url('folders')}):")
        for d in dirs:
            print("  ", d)
    return dirs[0]


# ------------------------------------------------------------------ helpers

def wrestler_files(saves):
    out = {}
    for f in saves.iterdir():
        m = WRESTLER_RE.match(f.name)
        if m and f.is_file():
            out[int(m.group(1))] = f
    return out


def free_numbers(saves):
    used = wrestler_files(saves)
    return [n for n in range(2, MAX_WRESTLER_FILES + 1) if n not in used]


def looks_like_aew_save(path):
    """Cheap sanity check on the plain wrapper header; does not read the encrypted data."""
    try:
        with open(path, "rb") as fh:
            head = fh.read(4096)
    except OSError:
        return False
    return head.startswith(GVAS_MAGIC) and SAVE_CLASS in head


def game_running():
    names = ("AEWFightForever-Win64-Shipping.exe", "AEWFightForever.exe")
    if os.name == "nt":
        try:
            import subprocess
            out = subprocess.run(["tasklist"], capture_output=True, text=True).stdout
            return any(n.lower() in out.lower() for n in names)
        except OSError:
            return False
    proc = Path("/proc")
    if not proc.is_dir():
        return False
    for p in proc.iterdir():
        if not p.name.isdigit():
            continue
        try:
            cmd = (p / "cmdline").read_bytes().decode(errors="replace")
        except OSError:
            continue
        if any(n in cmd for n in names):
            return True
    return False


def cloud_note(saves):
    if (saves / "steam_autocloud.vdf").exists():
        print("Note: Steam Cloud syncs this folder. Changes you make will upload on next launch; "
              "keep the backup zip if you want to undo them.")


def help_url(anchor):
    return f"{DEFAULT_SERVER.rstrip('/')}/how-it-works#{anchor}"


def backup(saves, dest=None, prefix="SaveGames"):
    """Zip the folder. restore picks the newest "SaveGames-" zip; its own safety copies use "Safety-"."""
    dest = Path(dest) if dest else saves.parent / "Ringside-backups"
    dest.mkdir(parents=True, exist_ok=True)
    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    out = dest / f"{prefix}-{stamp}.zip"
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        for f in sorted(saves.iterdir()):
            if f.is_file():
                z.write(f, f.name)
    return out


# ------------------------------------------------------------------ commands

def cmd_where(args):
    dirs = find_save_dirs()
    if not dirs:
        print("No SaveGames folder found.")
    for d in dirs:
        print(d)


def cmd_list(args):
    saves = resolve_saves(args.saves)
    files = wrestler_files(saves)
    print(f"SaveGames: {saves}")
    for n in sorted(files):
        f = files[n]
        size = f.stat().st_size
        note = "roster edits + first batch" if n == 1 else ("empty stub" if size < STUB_BYTES else "")
        print(f"  UGCWrestler_{n}.sav  {size:>10,} bytes  {note}")
    free = free_numbers(saves)
    print(f"{len(free)} free numbers (of 2..{MAX_WRESTLER_FILES})" +
          (f", next is {free[0]}" if free else ""))


def cmd_backup(args):
    saves = resolve_saves(args.saves)
    print("Backup written:", backup(saves, args.to))


def cmd_install(args):
    saves = resolve_saves(args.saves)
    if game_running() and not args.force:
        sys.exit("The game is running. Close it completely, then run this again.")
    sources = [Path(s).expanduser() for s in args.files]
    for s in sources:
        if not s.is_file():
            sys.exit(f"Not found: {s}")
        if s.name.lower() in PROTECTED:
            sys.exit(f"Refusing to install {s.name}.")
        if not looks_like_aew_save(s):
            sys.exit(f"{s.name} does not look like an AEW: Fight Forever save file.")
    free = free_numbers(saves)
    if len(free) < len(sources):
        sys.exit(f"Only {len(free)} free wrestler slots left, so {len(sources)} can't be added. "
                 f"Delete wrestlers you don't use in the game's Create menu, then run this again.\nHelp: {help_url('slots')}")
    cloud_note(saves)
    print("Backup written:", backup(saves))
    for s, n in zip(sources, free):
        target = saves / f"UGCWrestler_{n}.sav"
        if target.exists():  # free_numbers already excludes these; belt and braces
            sys.exit(f"{target.name} appeared while installing; stopping.")
        shutil.copy2(s, target)
        label = (getattr(args, "titles", None) or {}).get(str(s), s.name)
        print(f"Added {label} -> {target.name}")
    print("\nDone. Start the game and look in your created wrestlers.\n"
          f"Next time, open a terminal in {Path(sys.argv[0]).resolve().parent} and run the new command.\n"
          f"To take these back out, run: {self_command()} restore   (rolls back to just before this; it asks first)\n"
          f"Help: {help_url('restore')}")


def self_command():
    """How the player ran us, so hints can be pasted as-is (py on Windows, python3 elsewhere)."""
    py = "py" if os.name == "nt" else "python3"
    return f"{py} {Path(sys.argv[0]).name}"


def cmd_get(args):
    import json
    import tempfile
    import urllib.error
    import urllib.request
    base = (args.server or os.environ.get("RINGSIDE_SERVER") or DEFAULT_SERVER).rstrip("/")
    ids = []
    for raw in args.ids:
        for part in raw.replace(",", " ").split():
            part = part.lstrip("#")
            if not part.isdigit():
                sys.exit(f"'{part}' isn't a creation number. Use the numbers from the site, like: get 12 15")
            ids.append(int(part))
    if not ids:
        sys.exit("Give at least one creation number, like: get 12")
    # Check everything that can stop us before downloading anything.
    saves = resolve_saves(args.saves)
    if game_running() and not args.force:
        sys.exit("The game is running. Close it completely, then run this again.")
    free = len(free_numbers(saves))
    if free < len(ids):
        sys.exit(f"Only {free} free wrestler slots left, so {len(ids)} can't be added. "
                 f"Delete wrestlers you don't use in the game's Create menu, then run this again.\nHelp: {help_url('slots')}")

    def fetch(url, timeout):
        try:
            with urllib.request.urlopen(url, timeout=timeout) as r:
                return r.read()
        except urllib.error.HTTPError as e:
            if e.code == 404:
                return None
            sys.exit(f"The site had a problem (error {e.code}). Nothing was changed. Try again in a few minutes.")
        except (urllib.error.URLError, TimeoutError, OSError):
            sys.exit(f"Couldn't reach {base}. Check your internet connection and try again. Nothing was changed.\nHelp: {base}/how-it-works#offline")

    paths, titles, skipped = [], {}, []
    with tempfile.TemporaryDirectory() as tmp:
        for i, cid in enumerate(ids, 1):
            body = fetch(f"{base}/api/creations/{cid}", 30)
            if body is None:
                skipped.append(f"#{cid} isn't on the site any more (deleted or hidden)")
                continue
            meta = json.loads(body)
            if not meta.get("installable"):
                skipped.append(f"#{cid} {meta['title']} is a {'tag team' if meta.get('kind') == 'team' else meta.get('kind')}; add it with: {self_command()} swap {cid}")
                continue
            print(f"Downloading {i} of {len(ids)}: {meta['title']} by {meta['author']}")
            blob = fetch(f"{base}/api/creations/{cid}/file", 120)
            if blob is None:
                skipped.append(f"#{cid} {meta['title']} couldn't be downloaded")
                continue
            dest = Path(tmp) / f"{cid}.sav"
            dest.write_bytes(blob)
            paths.append(str(dest))
            titles[str(dest)] = meta["title"]
        for note in skipped:
            print("Skipped:", note)
        if not paths:
            sys.exit("Nothing to add, so nothing was changed.")
        args.files, args.titles = paths, titles
        cmd_install(args)


SWAP_FILES = {"team": ("UGCTeam.sav", "tag teams"), "arena": ("UGCArena.sav", "arenas")}


def cmd_swap(args):
    """Tag teams and arenas live in one file each, so sharing one means swapping the whole file."""
    import json
    import urllib.error
    import urllib.request
    base = (args.server or os.environ.get("RINGSIDE_SERVER") or DEFAULT_SERVER).rstrip("/")
    cid = str(args.id).lstrip("#")
    if not cid.isdigit():
        sys.exit(f"'{args.id}' isn't a creation number. Use the number from the site, like: swap 4")
    saves = resolve_saves(args.saves)
    if game_running() and not args.force:
        sys.exit("The game is running. Close it completely, then run this again.")
    try:
        with urllib.request.urlopen(f"{base}/api/creations/{cid}", timeout=30) as r:
            meta = json.load(r)
        if meta.get("kind") not in SWAP_FILES:
            sys.exit(f"#{cid} {meta['title']} is a wrestler, so it doesn't need a swap. Use: get {cid}")
        with urllib.request.urlopen(f"{base}/api/creations/{cid}/file", timeout=120) as r:
            blob = r.read()
    except urllib.error.HTTPError as e:
        if e.code == 404:
            sys.exit(f"#{cid} isn't on the site any more (deleted or hidden). Nothing was changed.")
        sys.exit(f"The site had a problem (error {e.code}). Nothing was changed. Try again in a few minutes.")
    except (urllib.error.URLError, TimeoutError, OSError):
        sys.exit(f"Couldn't reach {base}. Check your internet connection and try again. Nothing was changed.\nHelp: {base}/how-it-works#offline")
    if not (blob.startswith(GVAS_MAGIC) and SAVE_CLASS in blob[:4096]):
        sys.exit("That download isn't an AEW: Fight Forever save. Nothing was changed.")
    name, plural = SWAP_FILES[meta["kind"]]
    print(f"{meta['title']} by {meta['author']} replaces ALL of your {plural} ({name}).")
    if not args.yes:
        answer = input("Your current file is backed up first. Type yes to continue: ").strip().lower()
        if answer not in ("y", "yes"):
            sys.exit("Cancelled. Nothing was changed.")
    cloud_note(saves)
    print("Backup written:", backup(saves))
    original = saves.parent / "Ringside-backups" / f"original-{name}"
    if (saves / name).exists() and not original.exists():
        # Keep the player's own file aside once, so later swaps never lose it.
        shutil.copy2(saves / name, original)
    (saves / name).write_bytes(blob)
    print(f"\nDone. Start the game and look in your {plural}.\n"
          f"To get your own {plural} back (nothing else changes), run: {self_command()} unswap {plural.split()[-1]}")


def cmd_unswap(args):
    saves = resolve_saves(args.saves)
    if game_running() and not args.force:
        sys.exit("The game is running. Close it completely, then run this again.")
    which = args.what.lower().replace("tag", "").strip(" -_") or "teams"
    name = {"teams": "UGCTeam.sav", "team": "UGCTeam.sav", "arenas": "UGCArena.sav", "arena": "UGCArena.sav"}.get(which)
    if not name:
        sys.exit("Say which to put back: unswap teams, or unswap arenas")
    original = saves.parent / "Ringside-backups" / f"original-{name}"
    if not original.exists():
        sys.exit(f"There's no saved original {name}, so nothing was changed. It's kept the first time you use swap.\n"
                 f"To roll the whole folder back to before your last swap instead: {self_command()} restore\nHelp: {help_url('restore')}")
    print("Backup written:", backup(saves))
    shutil.copy2(original, saves / name)
    original.unlink()
    print(f"Done. Your own {name} is back. Nothing else was changed.")


def cmd_restore(args):
    saves = resolve_saves(args.saves)
    if game_running() and not args.force:
        sys.exit("The game is running. Close it completely, then run this again.")
    if args.zip:
        src = Path(args.zip).expanduser()
        if not src.exists():  # a bare name from the Ringside-backups folder
            src = saves.parent / "Ringside-backups" / src.name
        if not src.exists():
            sys.exit(f"Backup not found: {args.zip}")
    else:
        found = sorted((saves.parent / "Ringside-backups").glob("SaveGames-*.zip"))
        if not found:
            sys.exit("No backups found in Ringside-backups.")
        src = found[-1]
    when = datetime.datetime.fromtimestamp(src.stat().st_mtime)
    age = datetime.datetime.now() - when
    ago = f"{age.days} days ago" if age.days else f"{age.seconds // 3600} hours ago" if age.seconds >= 3600 else "minutes ago"
    print(f"This rolls your whole SaveGames folder back to {src.name}, made {when:%Y-%m-%d %H:%M} ({ago}).")
    print("Anything you created or played since then is undone, including wrestlers made after it.")
    if not getattr(args, "yes", False):
        answer = input("Your current state is backed up first. Type yes to continue: ").strip().lower()
        if answer not in ("y", "yes"):
            sys.exit("Cancelled. Nothing was changed.")
    with zipfile.ZipFile(src) as z:
        names = [n for n in z.namelist() if "/" not in n and "\\" not in n]
        print("Safety copy of your current saves:", backup(saves, prefix="Safety"))
        for f in saves.iterdir():
            if f.is_file() and WRESTLER_RE.match(f.name) and f.name not in names:
                f.unlink()  # remove wrestler files added after that backup
        for n in names:
            if n.lower() in PROTECTED:
                continue
            (saves / n).write_bytes(z.read(n))
    print(f"Restored {src.name}.")


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--saves", help="SaveGames folder to use (default: auto-detect)")
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("where").set_defaults(fn=cmd_where)
    sub.add_parser("list").set_defaults(fn=cmd_list)
    b = sub.add_parser("backup"); b.add_argument("--to"); b.set_defaults(fn=cmd_backup)
    i = sub.add_parser("install"); i.add_argument("files", nargs="+")
    i.add_argument("--force", action="store_true"); i.set_defaults(fn=cmd_install)
    g = sub.add_parser("get"); g.add_argument("ids", nargs="+")
    g.add_argument("--server"); g.add_argument("--force", action="store_true"); g.set_defaults(fn=cmd_get)
    w = sub.add_parser("swap"); w.add_argument("id"); w.add_argument("--server")
    w.add_argument("--yes", action="store_true", help="don't ask before replacing")
    w.add_argument("--force", action="store_true"); w.set_defaults(fn=cmd_swap)
    u = sub.add_parser("unswap"); u.add_argument("what", help="teams or arenas")
    u.add_argument("--force", action="store_true"); u.set_defaults(fn=cmd_unswap)
    r = sub.add_parser("restore"); r.add_argument("zip", nargs="?")
    r.add_argument("--yes", action="store_true", help="don't ask before rolling back")
    r.add_argument("--force", action="store_true"); r.set_defaults(fn=cmd_restore)
    args = ap.parse_args(argv)
    try:
        args.fn(args)
    except KeyboardInterrupt:
        sys.exit("\nStopped.")
    except Exception:  # anything unexpected: say so plainly, keep one line of detail for a bug report
        if os.environ.get("RINGSIDE_DEBUG"):
            raise
        import traceback
        sys.exit("Something unexpected went wrong. If it happened during an install, your saves were backed up "
                 f"first, so '{self_command()} restore' puts them back.\nDetails for a bug report: "
                 + traceback.format_exc().strip().splitlines()[-1])


if __name__ == "__main__":
    main()