#!/usr/bin/env python3
"""
App Store Connect — stuck screenshot remover

Safely finds screenshots in an App Store version and removes ONLY screenshots
whose assetDeliveryState is not COMPLETE.

Default mode is DRY RUN: nothing is deleted.
Use --delete to actually delete stuck screenshots.

Requirements:
    python3 -m pip install pyjwt cryptography requests

Authentication:
    App Store Connect API Team Key:
      - Issuer ID
      - Key ID
      - .p8 private key

The script intentionally does NOT delete screenshot sets.
"""

import argparse
import getpass
import json
import os
import sys
import time
from pathlib import Path

import jwt
import requests


BASE_URL = "https://api.appstoreconnect.apple.com/v1"


def make_token(issuer_id: str, key_id: str, private_key_path: str) -> str:
    with open(os.path.expanduser(private_key_path), "r", encoding="utf-8") as f:
        private_key = f.read()

    now = int(time.time())

    payload = {
        "iss": issuer_id,
        "iat": now,
        "exp": now + 1200,
        "aud": "appstoreconnect-v1",
    }

    headers = {
        "alg": "ES256",
        "kid": key_id,
        "typ": "JWT",
    }

    return jwt.encode(
        payload,
        private_key,
        algorithm="ES256",
        headers=headers,
    )


class ASC:
    def __init__(self, token: str):
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        })

    def get(self, path, params=None):
        r = self.session.get(BASE_URL + path, params=params)
        self._check(r)
        return r.json()

    def delete(self, path):
        r = self.session.delete(BASE_URL + path)
        self._check(r, allow_204=True)
        return r

    @staticmethod
    def _check(response, allow_204=False):
        if allow_204 and response.status_code == 204:
            return

        if response.status_code >= 400:
            print(f"\nApple API error: HTTP {response.status_code}")
            try:
                print(json.dumps(response.json(), indent=2, ensure_ascii=False))
            except Exception:
                print(response.text)
            sys.exit(1)


def find_app(api, app_name=None, app_id=None):
    if app_id:
        data = api.get(f"/apps/{app_id}")
        return data["data"]

    if not app_name:
        raise ValueError("Provide --app-id or --app-name")

    data = api.get(
        "/apps",
        {
            "filter[name]": app_name,
            "limit": 50,
        },
    )

    apps = data.get("data", [])

    if not apps:
        raise RuntimeError(f"No app found with name: {app_name}")

    if len(apps) > 1:
        print("Multiple apps matched:")
        for i, app in enumerate(apps, 1):
            print(
                f"  {i}. {app['id']} | "
                f"{app['attributes'].get('name')} | "
                f"{app['attributes'].get('bundleId')}"
            )
        raise RuntimeError(
            "Multiple apps found. Re-run with --app-id to avoid ambiguity."
        )

    return apps[0]


def find_version(api, app_id, version):
    data = api.get(
        f"/apps/{app_id}/appStoreVersions",
        {
            "filter[versionString]": version,
            "filter[platform]": "IOS",
            "limit": 50,
        },
    )

    versions = data.get("data", [])

    if not versions:
        raise RuntimeError(
            f"No iOS version '{version}' found for app {app_id}"
        )

    if len(versions) > 1:
        print("Multiple matching versions found:")
        for v in versions:
            print(
                f"  {v['id']} | "
                f"{v['attributes'].get('versionString')} | "
                f"{v['attributes'].get('appStoreState')}"
            )
        raise RuntimeError(
            "Multiple versions found. Resolve the ambiguity before continuing."
        )

    return versions[0]


def find_localization(api, version_id, locale):
    data = api.get(
        f"/appStoreVersions/{version_id}/appStoreVersionLocalizations",
        {"limit": 50},
    )

    matches = [
        x for x in data.get("data", [])
        if x["attributes"].get("locale") == locale
    ]

    if not matches:
        available = [
            x["attributes"].get("locale")
            for x in data.get("data", [])
        ]
        raise RuntimeError(
            f"Locale '{locale}' not found. Available locales: "
            + ", ".join(available)
        )

    if len(matches) > 1:
        raise RuntimeError(f"Multiple localizations found for {locale}")

    return matches[0]


def find_sets(api, localization_id, display_type=None):
    params = {"limit": 50}

    if display_type:
        params["filter[screenshotDisplayType]"] = display_type

    data = api.get(
        f"/appStoreVersionLocalizations/{localization_id}/appScreenshotSets",
        params,
    )

    return data.get("data", [])


def get_screenshots(api, set_id):
    data = api.get(
        f"/appScreenshotSets/{set_id}/appScreenshots",
        {"limit": 50},
    )
    return data.get("data", [])


def describe_screenshot(screenshot):
    attrs = screenshot.get("attributes", {})
    state = (attrs.get("assetDeliveryState") or {}).get("state")

    return {
        "id": screenshot["id"],
        "fileName": attrs.get("fileName"),
        "state": state,
        "fileSize": attrs.get("fileSize"),
        "width": (attrs.get("imageAsset") or {}).get("width"),
        "height": (attrs.get("imageAsset") or {}).get("height"),
        "sourceFileChecksum": attrs.get("sourceFileChecksum"),
        "uploadOperations": attrs.get("uploadOperations"),
    }


def main():
    parser = argparse.ArgumentParser(
        description="Find and optionally remove stuck App Store Connect screenshots."
    )

    parser.add_argument("--issuer-id", required=True)
    parser.add_argument("--key-id", required=True)
    parser.add_argument("--private-key", required=True)

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--app-id")
    group.add_argument("--app-name")

    parser.add_argument("--version", required=True)
    parser.add_argument("--locale", required=True)
    parser.add_argument(
        "--display-type",
        help="Optional, e.g. APP_IPHONE_67. If omitted, all screenshot sets are checked.",
    )

    parser.add_argument(
        "--delete",
        action="store_true",
        help="Actually delete stuck screenshots. Without this flag the script is read-only.",
    )

    args = parser.parse_args()

    token = make_token(
        args.issuer_id,
        args.key_id,
        args.private_key,
    )

    api = ASC(token)

    print("=== App Store Connect stuck screenshot checker ===\n")

    app = find_app(api, args.app_name, args.app_id)
    print(
        f"App: {app['attributes'].get('name')} "
        f"({app['id']})"
    )

    version = find_version(
        api,
        app["id"],
        args.version,
    )
    print(
        f"Version: {version['attributes'].get('versionString')} "
        f"| state={version['attributes'].get('appStoreState')}"
    )

    localization = find_localization(
        api,
        version["id"],
        args.locale,
    )
    print(
        f"Locale: {localization['attributes'].get('locale')} "
        f"({localization['id']})"
    )

    sets = find_sets(
        api,
        localization["id"],
        args.display_type,
    )

    if not sets:
        print("\nNo screenshot sets found.")
        return

    stuck = []

    print("\n=== Screenshot sets ===")

    for screenshot_set in sets:
        set_id = screenshot_set["id"]
        display_type = screenshot_set["attributes"].get(
            "screenshotDisplayType"
        )

        screenshots = get_screenshots(api, set_id)

        print(
            f"\n{display_type} | set={set_id} | "
            f"screenshots={len(screenshots)}"
        )

        for screenshot in screenshots:
            info = describe_screenshot(screenshot)

            print(
                f"  {info['fileName']} | "
                f"state={info['state']} | "
                f"id={info['id']} | "
                f"{info['width']}x{info['height']}"
            )

            if info["state"] != "COMPLETE":
                stuck.append(
                    {
                        "set_id": set_id,
                        "display_type": display_type,
                        **info,
                    }
                )

    print("\n=== Result ===")

    if not stuck:
        print("No screenshots with a non-COMPLETE delivery state were found.")
        return

    print(f"Found {len(stuck)} non-COMPLETE screenshot(s):")

    for item in stuck:
        print(
            f"  - {item['display_type']} | "
            f"{item['fileName']} | "
            f"state={item['state']} | "
            f"id={item['id']}"
        )

    if not args.delete:
        print(
            "\nDRY RUN: nothing was deleted.\n"
            "If these are the stuck screenshots you want to remove, "
            "re-run the same command with --delete."
        )
        return

    print(
        "\nWARNING: The following screenshot resources will be deleted:"
    )
    for item in stuck:
        print(
            f"  {item['id']} | "
            f"{item['display_type']} | "
            f"{item['fileName']} | "
            f"{item['state']}"
        )

    answer = input("\nType DELETE to continue: ").strip()
    if answer != "DELETE":
        print("Aborted. Nothing was deleted.")
        return

    for item in stuck:
        print(f"\nDeleting {item['id']} ...")
        response = api.delete(
            f"/appScreenshots/{item['id']}"
        )

        if response.status_code == 204:
            print("  Deleted successfully.")
        else:
            print(
                f"  Unexpected status: {response.status_code}"
            )

    print("\nDone.")
    print(
        "Refresh App Store Connect and check the screenshot set again."
    )


if __name__ == "__main__":
    main()