#!/bin/python3
SOURCE = "https://go.dev/dl/"
DOWNLOAD_TO = '/tmp'
INSTALL_TO = '/opt'
PROFILE_PATH = '/etc/profile'

import sys
import re
import urllib.request
import platform
import subprocess

GO_VERSION_REGEX = re.compile(r'/dl/(go(\d+\.\d+\.\d+)\.linux-(.*?)\.tar\.gz)')

GO_MACHINE_TO_PLATFORM = {
    'x86_64': 'amd64',
    'aarch64': 'arm64',
    # TODO: add more mappings
}

GO_BIN_DIR = f'{INSTALL_TO}/go/bin'

# shr = subprocess.run and check returncode == 0
def shr(cmd):
    if type(cmd) != list:
        cmd = ['bash', '-c', cmd]
    return subprocess.run(cmd).returncode == 0

# shrx = shr but print the command first
def shrx(cmd):
    print(f">>> {cmd}")
    return shr(cmd)

def require_input(prompt, valid=None, default_value=None):
    my_prompt = prompt
    if default_value is not None:
        my_prompt += f' (default: {default_value})'
    user_input = input(f'{my_prompt}: ').strip()
    if not user_input and default_value is not None:
        return default_value
    if valid is None or user_input in valid:
        return user_input
    print(f"Invalid input. Please enter one of: {', '.join(valid)}")
    sys.exit(1)

def get_platform(supported_platforms):
    machine = platform.machine()
    detected_platform = GO_MACHINE_TO_PLATFORM.get(machine, None)
    print(f'Detected platform: {machine} == {detected_platform}')
    print('Please specify your platform.')
    print('Supported platforms are:')
    print(' '.join(supported_platforms))
    user_platform = require_input(
        'Please specify your platform',
        valid=supported_platforms,
        default_value=detected_platform,
    )
    return user_platform

def load_url_as_string(url):
    with urllib.request.urlopen(url) as response:
        return response.read().decode('utf-8')

def download_url(url, destination):
    print(f'Downloading {url} to {destination}...')
    urllib.request.urlretrieve(url, destination)
    print('Download complete.')

def check_installed_go_version():
    print('Checking installed go version...')
    try:
        shr([f'{GO_BIN_DIR}/go','version'])
    except:
        print('go not installed')

def main():
    check_installed_go_version()
    print('Fetching latest Go version...')
    go_versions_page = load_url_as_string(SOURCE)
    go_versions = GO_VERSION_REGEX.findall(go_versions_page)
    if not go_versions:
        print(go_versions_page)
        print(f'ERROR: No Go versions found on {SOURCE}')
        return 1
    # the latest version is the first one
    _, latest_version, _ = go_versions[0]
    print(f'Latest linux Go version: {latest_version}')
    supported_platforms = {}
    for version_path, version, platform in go_versions:
        if version == latest_version:
            supported_platforms[platform] = version_path
    current_platform = get_platform(supported_platforms)
    latest_version_path = supported_platforms.get(current_platform, None)
    if not latest_version_path:
        print(f'ERROR: No Go version found for platform {current_platform}')
        print('Supported platforms are:')
        print(' '.join(supported_platforms))
        return 1
    path_to_remove = f'{INSTALL_TO}/go'
    url_to_download = f'{SOURCE}{latest_version_path}'
    download_path = f'{DOWNLOAD_TO}/{latest_version_path}'
    install_to = INSTALL_TO
    add_to_path = GO_BIN_DIR
    print(f'REMOVE: {path_to_remove}')
    print(f'DOWNLOAD: {url_to_download}')
    print(f'TO: {download_path}')
    print(f'INSTALL TO: {install_to}')
    print(f'ADD TO PATH: {add_to_path}')
    proceed = require_input('Proceed with installation?', default_value='y', valid=['y', 'n'])
    if proceed != 'y':
        print('Installation aborted by user.')
        return 0
    download_url(url_to_download, download_path)
    print('Extracting Go archive...')
    if not shrx(['rm', '-rf', path_to_remove]):
        print('Failed to remove existing Go installation.')
        return 1
    if not shrx(['mkdir', '-p', install_to]):
        print('Failed to create installation directory.')
        return 1
    if not shrx(['tar', '-C', install_to, '-xzvf', download_path]):
        print('Failed to extract Go archive.')
        return 1
    check_installed_go_version()
    need_to_add_to_path = False
    if not shrx(['grep', '-F', f'export PATH={add_to_path}:$PATH', PROFILE_PATH]):
        print(f'{add_to_path} is not in {PROFILE_PATH}.')
        need_to_add_to_path = True
    if need_to_add_to_path:
        print(f'Adding {add_to_path} to PATH in {PROFILE_PATH}...')
        with open(PROFILE_PATH, 'a') as profile_file:
            profile_file.write(f'\nexport PATH={add_to_path}:$PATH\n')
    print('script done')
    return 0

try:
    sys.exit(main())
except:
    e = sys.exception()
    print(str(e) or repr(e))

# vim: filetype=python