import aiohttp
import asyncio
import time
from io import BytesIO
from lxml import etree as etree

###########################
#                         #
#  Green Week Data Miner  #
#          v0.3           #
#                         #
# Yet another niche data  #
# collection script for   #
# NationStates stats by   #
#     Luca McGrath        #
#                         #
###########################

async def main(useragent = input("Enter a User-Agent: "), text = input("Enter the target region: ")):
    text = text.lower().replace(" ", "_")

    #API CALL
    async def api_call(url, useragent):
        async with aiohttp.ClientSession() as cs:
            async with cs.get(url, headers = {"""User-Agent""": f"""{useragent}, running GreenWeek v0.3 by Refuge Isle"""}) as r:
                response = await r.content.read()
        data = BytesIO(response)
        time.sleep(1.2) # Luca why are you using blocking code on an async function? Did you really design this around an avoidance of requests?
        return data     # Yes

    async def verification(data):

        # Check that NationStates is alive
        if "Server Error!" in str(data):
            print("API Error -- Sorry, NationStates appears to be down. Please try your call again later.")
            print("*obnoxious phone noises*")
            return 1

        # Check that the region is valid
        elif "Unknown region" in str(data):
            print("Region Not Found -- This doesn't look like a vaild region. Or if it was, it doesn't exist anymore.")
            return 1

        else:
            return 0

    # Initial API Call
    print("Pulling region population data...")
    call = await api_call(f"""https://www.nationstates.net/cgi-bin/api.cgi?region={text}&q=nations""", useragent)
    check = await verification(call)
    if check > 0:
        return

    # We assume all is well at this point, so collect population
    nations = []
    for event, element in etree.iterparse(call):
        try:
            if element.tag == "NATIONS":
                nations = element.text.split(":")
            element.clear(keep_tail=True)
        except:
            continue

    # On the off chance someone uses this on a wasteland region
    if len(nations) == 0:
        print("No nations could be found in the region you specified. So that's strange.")
        return

    # Provide an estimate
    est = round((len(nations) * 1.2)/60)
    print(f"Pulling eco-data now. This should take approximately {est} minutes.")

    # Cycle through population with census API calls
    prepared = ""
    for x in nations:

        call = await api_call(f"""https://www.nationstates.net/cgi-bin/api.cgi?nation={x};q=name+census;scale=7""", useragent)
        check = await verification(call)
        if check > 0:
            return

        # Make sure the nation hasn't CTE'd or been DEAT since the script started
        if "Unknown nation" in str(call):
            print(f"Nation CTE'd: {x}")
            prepared += f"{x},CTE\n"

        # Parse and record
        else:
            for event, element in etree.iterparse(call):
                try:
                    if element.tag == "NAME":
                        name = element.text
                    if element.tag == "SCORE":
                        value = element.text
                    element.clear(keep_tail=True)
                except:
                    continue

            print(f"Recorded: {name} ({value})")
            prepared += f"{name},{value}\n"

    # Write data on conclusion
    with open("GreenWeekOut.csv", "w") as f:
        f.write(prepared)
        f.close()

    print("""Done! There is a file called GreenWeekOut.csv in the same folder where you activated this script.""")

asyncio.get_event_loop().run_until_complete(main())