import argparse, os, plistlib, shlex, shutil, struct, subprocess, sys, tempfile, zlib

NAME = "Pythia"
COMMENT = "UTAU-style piano roll for the Ascended Ancient energeme synthesizer"
HERE = os.path.dirname(os.path.abspath(__file__))
TARGET = os.path.join(HERE, 'interface.py')
ASSETS = os.path.join(HERE, 'assets')
WIN, MAC = sys.platform == 'win32', sys.platform == 'darwin'

def os_release():
    ids = set()
    try:
        with open('/etc/os-release', encoding='utf-8') as f:
            for line in f:
                key, _, val = line.partition('=')
                if key in ('ID', 'ID_LIKE'):
                    ids |= set(val.strip().strip('"\'').lower().split())
    except OSError:
        pass
    return ids

def dependency_hint():
    if WIN:
        return "install Python from python.org (it includes tkinter), then:  pip install numpy scipy pillow"
    if MAC:
        return "brew install python-tk   (or use the python.org installer), then:  pip3 install numpy scipy pillow"
    ids = os_release()
    if ids & {'debian', 'ubuntu', 'linuxmint', 'pop', 'raspbian'}:
        return "sudo apt install python3-tk python3-numpy python3-scipy python3-pil"
    if ids & {'fedora', 'rhel', 'centos', 'rocky', 'almalinux'}:
        return "sudo dnf install python3-tkinter python3-numpy python3-scipy python3-pillow"
    if ids & {'arch', 'manjaro', 'endeavouros'}:
        return "sudo pacman -S tk python-numpy python-scipy python-pillow"
    if ids & {'suse', 'opensuse', 'opensuse-leap', 'opensuse-tumbleweed'}:
        return "sudo zypper install python3-tk python3-numpy python3-scipy python3-Pillow"
    return "install tkinter, numpy and scipy for Python 3 with your package manager (or pip)"

def check_dependencies(python):
    probe = subprocess.run([python, '-c', 'import tkinter, numpy, scipy'], capture_output=True, text=True)
    if probe.returncode:
        last = (probe.stderr.strip().splitlines() or ['import failed'])[-1]
        print(f"warning: {python} cannot import tkinter/numpy/scipy yet:\n  {last}\n  fix: {dependency_hint()}")

def get_icon(ext, dest):
    path = os.path.join(ASSETS, 'icon.' + ext)
    if os.path.isfile(path):
        return path
    other = os.path.join(ASSETS, 'icon.' + ('ico' if ext == 'png' else 'png'))
    if os.path.isfile(other):
        try:
            from PIL import Image
            im = Image.open(other)
            if im.info.get('sizes'):
                im.size = max(im.info['sizes'])
            im = im.convert('RGBA')
            if ext == 'ico':
                im.save(dest, sizes=[(n, n) for n in (16, 32, 48, 64, 128, 256)])
            else:
                im.save(dest)
            return dest
        except Exception:
            pass
    return None

def make_png(path, size=256):
    s = size / 64
    BG, BORDER, GRID, NOTE, HI, CURSOR = ((0x15, 0x16, 0x1a), (0x3b, 0x82, 0xf6), (0x2f, 0x32, 0x39),
                                          (0x2f, 0x6f, 0xe0), (0x9c, 0xc4, 0xff), (0xff, 0xcf, 0x5c))
    px = [[BG] * size for _ in range(size)]
    def rect(x0, y0, x1, y1, c):
        a, b, e = int(x0 * s), int(y0 * s), min(size, max(int(x1 * s), int(x0 * s) + 1))
        for y in range(b, min(size, max(int(y1 * s), b + 1))):
            px[y][a:e] = [c] * (e - a)
    rect(0, 0, 64, 64, BORDER); rect(2, 2, 62, 62, BG)
    for y in (18, 30, 42, 54):
        rect(4, y, 60, y + 1, GRID)
    for x0, y0, x1 in ((7, 38, 25), (27, 26, 45), (47, 14, 59)):
        rect(x0, y0, x1, y0 + 8, NOTE); rect(x0, y0, x1, y0 + 2, HI)
    rect(25, 4, 27, 60, CURSOR)
    raw = b''.join(b'\x00' + bytes(v for p in row for v in p) for row in px)
    def chunk(tag, data):
        body = tag + data
        return struct.pack('>I', len(data)) + body + struct.pack('>I', zlib.crc32(body) & 0xffffffff)
    with open(path, 'wb') as f:
        f.write(b'\x89PNG\r\n\x1a\n' + chunk(b'IHDR', struct.pack('>IIBBBBB', size, size, 8, 2, 0, 0, 0))
                + chunk(b'IDAT', zlib.compress(raw, 9)) + chunk(b'IEND', b''))

def remove(*paths):
    gone = []
    for p in paths:
        try:
            if os.path.islink(p) or os.path.isfile(p):
                os.remove(p)
            elif os.path.isdir(p):
                shutil.rmtree(p)
            else:
                continue
            gone.append(p)
        except OSError as e:
            print(f"could not remove {p}: {e}")
    return gone

def _xdg_data():
    return os.environ.get('XDG_DATA_HOME') or os.path.expanduser('~/.local/share')

def _linux_paths():
    apps = os.path.join(_xdg_data(), 'applications')
    icons = os.path.join(_xdg_data(), 'icons', 'hicolor', '256x256', 'apps')
    return apps, icons, os.path.join(apps, 'pythia.desktop'), os.path.join(icons, 'pythia.png')

def _linux_desktop_dir():
    try:
        out = subprocess.run(['xdg-user-dir', 'DESKTOP'], capture_output=True, text=True).stdout.strip()
        if out:
            return out
    except OSError:
        pass
    return os.path.expanduser('~/Desktop')

def _desktop_quote(path):
    esc = path.replace('\\', '\\\\').replace('"', '\\"').replace('`', '\\`').replace('$', '\\$')
    return f'"{esc}"'

def _refresh_linux_caches(apps):
    for cmd in (['update-desktop-database', apps],
                ['gtk-update-icon-cache', '-f', '-t', os.path.join(_xdg_data(), 'icons', 'hicolor')]):
        if shutil.which(cmd[0]):
            subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def linux_install(python, desktop):
    apps, icons, entry_path, icon_path = _linux_paths()
    os.makedirs(apps, exist_ok=True); os.makedirs(icons, exist_ok=True)
    src = get_icon('png', icon_path)
    if src is None:
        make_png(icon_path)
    elif src != icon_path:
        shutil.copyfile(src, icon_path)
    entry = "\n".join([
        "[Desktop Entry]", "Type=Application", "Version=1.0", f"Name={NAME}",
        "GenericName=Energeme piano roll", f"Comment={COMMENT}",
        f"Exec={_desktop_quote(python)} {_desktop_quote(TARGET)}", f"Path={HERE}", f"Icon={icon_path}",
        "Terminal=false", "Categories=AudioVideo;Audio;Music;",
        "Keywords=utau;vocaloid;conlang;energeme;synth;piano roll;",
        f"StartupWMClass={NAME}", "StartupNotify=true", ""])
    def write(path):
        with open(path, 'w', encoding='utf-8') as f:
            f.write(entry)
        os.chmod(path, 0o755)
    write(entry_path)
    _refresh_linux_caches(apps)
    print(f"menu entry: {entry_path}\nicon:       {icon_path}\nPress the Super key and type '{NAME}'. "
          "(If it doesn't show up at once, log out and back in.)")
    if desktop:
        folder = _linux_desktop_dir()
        if not os.path.isdir(folder):
            print(f"no desktop folder at {folder}; skipped the desktop shortcut")
            return
        path = os.path.join(folder, 'Pythia.desktop')
        write(path)
        if shutil.which('gio'):
            subprocess.run(['gio', 'set', path, 'metadata::trusted', 'true'],
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print(f"desktop:    {path}\n(If it shows a warning icon, right-click it > Allow Launching / Trust.)")

def linux_uninstall():
    apps, _icons, entry_path, icon_path = _linux_paths()
    gone = remove(entry_path, icon_path, os.path.join(_linux_desktop_dir(), 'Pythia.desktop'))
    _refresh_linux_caches(apps)
    return gone

def _powershell(script, **env):
    exe = shutil.which('powershell') or shutil.which('pwsh') or 'powershell'
    res = subprocess.run([exe, '-NoProfile', '-NonInteractive', '-ExecutionPolicy', 'Bypass', '-Command', script],
                         capture_output=True, text=True, env={**os.environ, **env})
    if res.returncode:
        raise RuntimeError(res.stderr.strip() or 'PowerShell failed')
    return res.stdout.strip()

def _win_folder(kind):
    return _powershell(f"[Environment]::GetFolderPath('{kind}')")

def _windowless(python):
    head, tail = os.path.split(python)
    if tail.lower() == 'python.exe' and os.path.isfile(os.path.join(head, 'pythonw.exe')):
        return os.path.join(head, 'pythonw.exe')
    return python

def win_install(python, desktop):
    icon = get_icon('ico', os.path.join(ASSETS, 'icon.ico')) or ''
    make = ("$s=(New-Object -ComObject WScript.Shell).CreateShortcut($env:PY_LNK);"
            "$s.TargetPath=$env:PY_TARGET;$s.Arguments=$env:PY_ARGS;$s.WorkingDirectory=$env:PY_CWD;"
            "$s.Description=$env:PY_DESC;if($env:PY_ICON){$s.IconLocation=$env:PY_ICON};$s.Save()")
    folders = [('Programs', 'Start menu')] + ([('Desktop', 'desktop')] if desktop else [])
    for kind, label in folders:
        lnk = os.path.join(_win_folder(kind), NAME + '.lnk')
        _powershell(make, PY_LNK=lnk, PY_TARGET=_windowless(python), PY_ARGS=f'"{TARGET}"',
                    PY_CWD=HERE, PY_DESC=COMMENT, PY_ICON=icon)
        print(f"{label}: {lnk}")
    if not icon:
        print("note: assets/icon.ico not found, so the shortcut uses the default Python icon")
    print(f"Press the Windows key and type '{NAME}'.")

def win_uninstall():
    return remove(*(os.path.join(_win_folder(k), NAME + '.lnk') for k in ('Programs', 'Desktop')))

def _mac_app():
    return os.path.expanduser(f'~/Applications/{NAME}.app')

def _make_icns(png, out):
    if not (shutil.which('sips') and shutil.which('iconutil')):
        return False
    with tempfile.TemporaryDirectory() as tmp:
        iconset = os.path.join(tmp, 'Pythia.iconset')
        os.mkdir(iconset)
        for name, px in (('16x16', 16), ('16x16@2x', 32), ('32x32', 32), ('32x32@2x', 64),
                         ('128x128', 128), ('128x128@2x', 256), ('256x256', 256)):
            r = subprocess.run(['sips', '-z', str(px), str(px), png, '--out', os.path.join(iconset, f'icon_{name}.png')],
                               capture_output=True)
            if r.returncode:
                return False
        return subprocess.run(['iconutil', '-c', 'icns', iconset, '-o', out], capture_output=True).returncode == 0

def mac_install(python, desktop):
    app = _mac_app()
    remove(app)
    macos, res = os.path.join(app, 'Contents', 'MacOS'), os.path.join(app, 'Contents', 'Resources')
    os.makedirs(macos); os.makedirs(res)
    launcher = os.path.join(macos, NAME)
    with open(launcher, 'w', encoding='utf-8') as f:
        f.write(f"#!/bin/bash\ncd {shlex.quote(HERE)}\nexec {shlex.quote(python)} {shlex.quote(TARGET)}\n")
    os.chmod(launcher, 0o755)
    with tempfile.TemporaryDirectory() as tmp:
        png = get_icon('png', os.path.join(tmp, 'icon.png'))
        has_icon = bool(png) and _make_icns(png, os.path.join(res, f'{NAME}.icns'))
    info = {'CFBundleName': NAME, 'CFBundleDisplayName': NAME, 'CFBundleExecutable': NAME,
            'CFBundleIdentifier': 'local.pythia', 'CFBundlePackageType': 'APPL', 'CFBundleVersion': '1.0',
            'CFBundleShortVersionString': '1.0', 'NSHighResolutionCapable': True}
    if has_icon:
        info['CFBundleIconFile'] = NAME
    with open(os.path.join(app, 'Contents', 'Info.plist'), 'wb') as f:
        plistlib.dump(info, f)
    print(f"app: {app}\nOpen Launchpad or Spotlight and type '{NAME}'.")
    if not has_icon:
        print("note: no icon (needs assets/icon.png; sips/iconutil come with macOS)")
    if desktop:
        link = os.path.expanduser(f'~/Desktop/{NAME}')
        remove(link)
        os.symlink(app, link)
        print(f"desktop: {link}")

def mac_uninstall():
    return remove(_mac_app(), os.path.expanduser(f'~/Desktop/{NAME}'))

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--desktop', action='store_true', help="also put a shortcut on the desktop")
    ap.add_argument('--uninstall', action='store_true', help="remove the menu entry and desktop shortcut")
    ap.add_argument('--python', default=sys.executable, help="Python to launch Pythia with (default: this one)")
    a = ap.parse_args()
    if a.uninstall:
        gone = win_uninstall() if WIN else mac_uninstall() if MAC else linux_uninstall()
        print("removed:", ", ".join(gone) if gone else "nothing to remove")
        return
    if not os.path.isfile(TARGET):
        sys.exit(f"interface.py not found next to this script ({HERE}).")
    check_dependencies(a.python)
    (win_install if WIN else mac_install if MAC else linux_install)(a.python, a.desktop)

if __name__ == '__main__':
    main()