#!/usr/bin/env python3 

import boto3
import argparse
import pandas as pd  # For nicely printing the table, install with 'pip install pandas'
from operator import itemgetter

def collect_ami_information(args):

  amis = []
  # Initialize a session using Amazon EC2
  session = boto3.Session(region_name=region)
  ec2 = session.client('ec2')

  filters = [{'Name': 'architecture', 'Values': [args.arch]}]
  response = ec2.describe_images(Owners=[args.owner], Filters=filters)

  # Fetch AMI information
  # response = ec2.describe_images(Owners=['aws-marketplace'])
  # response = ec2.describe_images(Owners=[args.owner])
  # response = ec2.describe_images()

  for image in response['Images']:
    Desc  = image.get('Description', '')[:60] 
    Name  = image.get('Name', '')[:60] 
    Cdate = image.get('CreationDate', '').split('T')[0]
    AmiId = image.get('ImageId')
    if not args.desc or (args.desc and (args.desc.lower() in (image.get('Name', '').lower() + image.get('Description', '').lower()))):
        print(f"{AmiId:<21}  {Cdate:<10}  {Name:<60}  {Desc:<60}")
    if args.latest:
      if args.desc.lower() in (image.get('Name', '').lower() + image.get('Description', '').lower()):
        amis.append({'ImageId': AmiId, 'CreationDate': image.get('CreationDate')})

  if args.latest and len(amis) > 0:
    latest_amis = sorted(amis, key=itemgetter('CreationDate'), reverse=True)
    latest = latest_amis[0]
    print(f"\nLatest: {latest['ImageId']}")
    
def get_ami_owners(region):
    # Initialize a session using Amazon EC2
    ec2 = boto3.client('ec2', region_name=region)
    
    # Fetch a list of public AMIs (this can be adjusted based on your needs)
    response = ec2.describe_images(
        Filters=[{'Name': 'is-public', 'Values': ['true']}],
        MaxResults=1000  # Adjust based on your needs, note there's a max limit per call
    )
    
    # Extract owner IDs
    # owner_ids = {image['OwnerId'] for image in response['Images']}
    owners = {image['ImageOwnerAlias'] for image in response['Images']}
    
    return owners
    # return owner_ids

def get_ami_owners_and_aliases(region):
  # Initialize a session using Amazon EC2
  ec2 = boto3.client('ec2', region_name=region)
    
  # Fetch a list of public AMIs (this can be adjusted based on your needs)
  response = ec2.describe_images(
    Filters=[{'Name': 'is-public', 'Values': ['true']}],
    MaxResults=1000  # Adjust based on your needs, note there's a max limit per call
  )
    
  # Extract owner IDs and Aliases
  owners_info = set()
  for image in response['Images']:
    owner_id = image.get('OwnerId', 'No OwnerId')
    owner_alias = image.get('ImageOwnerAlias', 'No Alias')
    owners_info.add((owner_id, owner_alias))

  print("Unique AMI owner IDs and aliases in region '{}':".format(region))
  print(f"\n{'OwnerID':<12} {'Alias':<}")
  print(f"{'='*12} {'='*20}")
  for owner_id, owner_alias in owners_info:
    print(f"{owner_id:<12} {owner_alias}")
    
# Create the parser
parser = argparse.ArgumentParser(description="AWS AMI search")
parser.add_argument("--region", dest='region', action='store', default='us-east-2',       help="aws region - defaults to us-east-2")
parser.add_argument("--owner",  dest='owner',  action='store', default='aws-marketplace', help="ami owner  - defaults to aws-marketplace")
parser.add_argument("--owners", dest='owners', action='store_true',                       help="list owners and ownerids in a region")
parser.add_argument("--desc",   dest='desc',   action='store',                            help="ami description")
parser.add_argument("--arch",   dest='arch',   action='store', default='x86_64',          help="architecture - defaults to x86_64")
parser.add_argument("--latest", dest='latest', action='store_true', default=False,        help="list latest AMI as part of search")
args = parser.parse_args()

region = args.region
owner  = args.owner

if args.owners:
  get_ami_owners_and_aliases(region)
else:
  collect_ami_information(args)

# owners_info = get_ami_owners_and_aliases(region)

# owners = get_ami_owners(region)
# print("Unique AMI owners in region '{}':".format(region))
# for owner in owners:
#     print(owner)

# ami_information = collect_ami_information(region)
# # Print the information in a table format
# df = pd.DataFrame(ami_information)
# print(df.to_string(index=False))