import argparse
import sys
import os
import zipfile
import tempfile
import json
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError

def main():
    parser = argparse.ArgumentParser(description='Download and extract GitHub release')
    parser.add_argument('--extension', required=True, help='Extension name')
    parser.add_argument('--version', required=True, help='Version number')
    parser.add_argument('--out', required=True, help='Output directory')
    parser.add_argument('--github-token-env', default='GITHUB_TOKEN',
                        help='Environment variable name containing a personal access token, if needed. Default is "GITHUB_TOKEN".')
    parser.add_argument('--asset-name', default=None, 
                        help='Optional asset filename to match exactly. Defaults to "<extension>_<version>.zip" if not provided.')
    args = parser.parse_args()

    # Build default asset name if not specified
    if not args.asset_name:
        args.asset_name = f"{args.extension}_{args.version}.zip"

    # Ensure output directory exists
    os.makedirs(args.out, exist_ok=True)

    token = os.getenv(args.github_token_env, "")
    headers = {
        "Accept": "application/vnd.github+json",
        "X-GitHub-Api-Version": "2022-11-28"
    }
    if token:
        headers["Authorization"] = f"Bearer {token}"

    # 1) Get the list of releases for the repository
    release_url = f"https://api.github.com/repos/Keyfactor/{args.extension}/releases"
    print(f"Fetching releases for: {release_url}")
    try:
        req = Request(release_url, headers=headers)
        with urlopen(req) as response:
            if response.status != 200:
                print(f"Error: Received status code {response.status} from {release_url}", file=sys.stderr)
                sys.exit(1)
            data = response.read()
            releases = json.loads(data)
    except HTTPError as e:
        print(f"Error fetching release list from {release_url}: {e}", file=sys.stderr)
        sys.exit(1)
    except URLError as e:
        print(f"Connection error while trying to reach {release_url}: {e}", file=sys.stderr)
        sys.exit(1)

    # 2) Find the release that matches the --version (tag_name)
    release_data = None
    for rel in releases:
        if rel.get('tag_name') == args.version:
            release_data = rel
            break

    if not release_data:
        print(f"No release found with tag_name == {args.version}", file=sys.stderr)
        sys.exit(1)

    # 3) Within that release, find the asset URL corresponding to the best choice.
    assets = release_data.get('assets', [])
    if not assets:
        print(f"No assets found in release tag {args.version}", file=sys.stderr)
        sys.exit(1)

    net8_asset = None
    fallback_asset = None

    for asset in assets:
        asset_name = asset.get('name', '')
        print(f"Found asset: {asset_name}")  # For debugging: show all asset names

        # Check if asset name includes 'net8.0' (case-insensitive to be safe)
        if 'net8.0' in asset_name.lower():
            net8_asset = asset
        else:
            if fallback_asset is None:
                fallback_asset = asset

    # Prefer net8.0 if present, otherwise pick the first fallback
    chosen_asset = net8_asset if net8_asset else fallback_asset
    if not chosen_asset:
        print(f"No suitable asset found (net8.0 or otherwise) in release {args.version}", file=sys.stderr)
        sys.exit(1)

    asset_url = chosen_asset.get('url')
    if not asset_url:
        print(f"No 'url' found in chosen asset for release {args.version}", file=sys.stderr)
        sys.exit(1)

    print(f"Chosen asset: {chosen_asset['name']} ({asset_url})")

    # 4) Download the asset to a temporary directory
    with tempfile.TemporaryDirectory() as tmpdirname:
        local_zip = os.path.join(tmpdirname, args.asset_name)
        print(f"Downloading to {local_zip}...")

        # Update headers to download octet stream
        dl_headers = headers.copy()
        dl_headers['Accept'] = "application/octet-stream"

        try:
            dl_req = Request(asset_url, headers=dl_headers)
            with urlopen(dl_req) as dl_response:
                if dl_response.status != 200:
                    print(f"Error: Received status code {dl_response.status} when downloading asset from {asset_url}", file=sys.stderr)
                    sys.exit(1)

                # Write in chunks
                with open(local_zip, 'wb') as f:
                    chunk_size = 8192
                    while True:
                        chunk = dl_response.read(chunk_size)
                        if not chunk:
                            break
                        f.write(chunk)
        except HTTPError as e:
            print(f"Error downloading asset from {asset_url}:\n{e}", file=sys.stderr)
            sys.exit(1)
        except URLError as e:
            print(f"Connection error while trying to download {asset_url}: {e}", file=sys.stderr)
            sys.exit(1)

        # 5) Extract to output directory
        print(f"Extracting {local_zip} into {args.out}...")
        with zipfile.ZipFile(local_zip, 'r') as zip_ref:
            # Normalize all filenames to use forward slashes
            normalized_names = [name.replace('\\', '/') for name in zip_ref.namelist()]
            # Find common prefix (root directory if the zip has one)
            common_prefix = os.path.commonprefix(normalized_names)
            if not common_prefix.endswith('/'):
                common_prefix = os.path.dirname(common_prefix) + '/'
            prefix_length = len(common_prefix)

            # Extract each member, stripping out the common prefix
            for member in zip_ref.infolist():
                original_name = member.filename.replace('\\', '/')
                relative_path = original_name[prefix_length:]
                member.filename = relative_path
                zip_ref.extract(member, args.out)

    print(f"Extraction complete. Files are in {args.out}.")

if __name__ == '__main__':
    main()