pruncalc/planet_bases.py

from __future__ import annotations

import collections
import concurrent.futures
import sys
import typing

import tap

import cache

class Args(tap.Tap):
	planets: list[str] = []
	sector: tuple[str, ...] = ()

	def configure(self) -> None:
		self.add_argument('planets', nargs='*', metavar='planet') # take planets as positional args instead of flag

def main() -> None:
	args = Args().parse_args()
	if args.sector:
		sectors = []
		all_sector_ids, sector_systems = get_sectors()
		for sector in args.sector:
			if len(sector) == 2:
				if len(sector_ids := all_sector_ids[sector]) > 1:
					print(sector, 'has multiple sector ids')
					for sector_id in sector_ids:
						print(f'\t{sector_id}:', ', '.join(system['Name'] for system in sector_systems[sector_id]))
					return
				else:
					(sector_id,) = sector_ids
					sectors.append(sector_id)
			elif sector.startswith('sector-'):
				sectors.append(sector)
			else:
				sys.exit(f'invalid sector: {sector}')

		system_ids = set()
		for sector_id in sectors:
			system_ids.update(system['SystemId'] for system in sector_systems[sector_id])

		all_planets: typing.Sequence[Planet] = cache.get('https://rest.fnar.net/planet/allplanets/full')
		for planet in all_planets:
			if planet['SystemId'] in system_ids:
				args.planets.append(planet['PlanetName'])

	with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
		for planet, bases in zip(args.planets, executor.map(get_bases, args.planets)):
			print(planet)
			companies = executor.map(get_company, (base['OwnerCode'] for base in bases if base['OwnerCode'] is not None))
			for co in sorted(companies, key=lambda co: co['CorporationCode'] or '\uffff'):
				print('\t', end='')
				if co['CorporationCode'] is not None:
					print(f'[{co["CorporationCode"]}]', end=' ')
				print(f'{co["UserName"]}\t{co["CompanyCode"]}\t{co["CompanyName"]}')

def get_sectors() -> tuple[typing.Mapping[str, typing.Collection[str]], typing.Mapping[str, typing.Sequence[System]]]:
	systems: typing.Sequence[System] = cache.get('https://rest.fnar.net/systemstars')
	sector_ids = collections.defaultdict(set)
	sector_systems = collections.defaultdict(list)
	for system in systems:
		sector_name = system['NaturalId'][:2]
		sector_ids[sector_name].add(system['SectorId'])
		sector_systems[system['SectorId']].append(system)
	return sector_ids, sector_systems

def get_bases(planet: str) -> typing.Sequence[Site]:
	return cache.get('https://rest.fnar.net/planet/sites/' + planet, expiry=cache.ONE_DAY)

def get_company(code: str) -> typing.Mapping[str, typing.Any]:
	try:
		return cache.get('https://rest.fnar.net/company/code/' + code, expiry=cache.ONE_DAY)
	except cache.NoContent:
		return {'CorporationCode': None, 'CompanyCode': code, 'CompanyName': None, 'UserName': None}

class System(typing.TypedDict):
	SystemId: str
	NaturalId: str
	Name: str
	SectorId: str

class Planet(typing.TypedDict):
	PlanetNaturalId: str
	PlanetName: str
	SystemId: str

class Site(typing.TypedDict):
	OwnerId: str
	OwnerName: str
	OwnerCode: str | None

class Company(typing.TypedDict):
	UserName: str
	CompanyCode: str
	CompanyName: str
	CorporationCode: str

if __name__ == '__main__':
	main()