from multiprocessing import Pool, Queue
import random

# Gender globals
GENDER_MALE = 0
GENDER_FEMALE = 1

# Ability globals
ABILITY_ONE = 0
ABILITY_TWO = 1

# Wurmple globals
WURMPLE_EVO_SILCOON = 0
WURMPLE_EVO_CASCOON = 1


def brute_force_pid(gender=None, ability=None, nature=None, is_shiny=False, tid=None, sid=None, unown_letter=None,
                    wurmple_evo=None):
    # Domain checks
    if gender is not None and gender is not GENDER_MALE and gender is not GENDER_FEMALE:
        print(f"ERROR: Supplied gender restriction <{gender}> is invalid")
        return None
    if ability is not None and ability is not ABILITY_ONE and ability is not ABILITY_TWO:
        print(f"ERROR: Supplied ability restriction <{ability}> is invalid")
        return None
    if nature is not None and not 0 <= nature <= 24:
        print(f"ERROR: Supplied nature restriction <{nature}> is invalid")
        return None
    if tid is not None and not 0 <= tid <= 0xFFFF:
        print(f"ERROR: Supplied Trainer ID <{tid}> is invalid")
        return None
    if sid is not None and not 0 <= sid <= 0xFFFF:
        print(f"ERROR: Supplied Secret ID <{sid}> is invalid")
        return None
    if unown_letter is not None and unown_letter not in "abcdefghijklmnopqrstuvwxyz!?":
        print(f"ERROR: Supplied Unown letter restriction <{unown_letter}> is invalid")
        return None
    if wurmple_evo is not None and wurmple_evo is not WURMPLE_EVO_SILCOON and wurmple_evo is not WURMPLE_EVO_CASCOON:
        print(f"ERROR: Supplied Wurmple evolution restriction <{wurmple_evo}> is invalid")
        return None

    # Sanity checks
    if is_shiny and (tid is None or sid is None):
        print("ERROR: Must supply a TID/SID to compute a shiny PID")
        return None

    if unown_letter is not None and gender is not None:
        print("WARN: Unown is genderless, setting a gender restriction for Unown can cause unnecessarily high compute "
              "times")
    if unown_letter is not None and ability is not None:
        print("ERROR: Unown does not have multiple abilities, and Unown letter restrictions are incompatible with "
              "Ability restrictions")
        return None

    if unown_letter is not None and wurmple_evo is not None:
        print("ERROR: Cannot restrict by both Unown letter and Wurmple evolution")

    if unown_letter is not None:
        unown_letter = ord(unown_letter)
        if ord("a") <= unown_letter <= ord("z"):
            unown_letter = unown_letter - 97
        elif unown_letter == ord("!"):
            unown_letter = 26
        elif unown_letter == ord("?"):
            unown_letter = 27

    result = None
    with Pool() as p:
        num_workers = p._processes

        print(f"Spawning {num_workers} workers")
        await_result = Queue(maxsize=num_workers)
        for worker_id in range(0, num_workers):
            p.apply_async(
                generate_and_test_pids,
                (worker_id, num_workers, gender, ability, nature, is_shiny, tid, sid, unown_letter, wurmple_evo),
                callback=lambda pid: await_result.put(pid))

        p.close()

        for i in range(0, num_workers):
            result = await_result.get()
            if result is not None:
                break

        p.terminate()
        p.join()

    return result


def generate_and_test_pids(worker_id, num_workers, gender, ability, nature, is_shiny, tid, sid, unown_letter,
                           wurmple_evo):
    start = int(worker_id * (0xFFFFFFFF + 1) / num_workers)
    stop = int((worker_id + 1) * (0xFFFFFFFF + 1) / num_workers)

    pid = start
    tried_msg = 1
    tried_msg_threshold = start + (stop - start) // 100
    while pid < stop:
        if pid > tried_msg_threshold:
            print(f"Tried {tried_msg}% of all PIDs for worker #{worker_id}")
            tried_msg += 1
            tried_msg_threshold = start + int((stop - start) / 100 * tried_msg)

        if gender == GENDER_FEMALE and not pid & 0xFF < 31:
            pid += 0x100 - pid & 0xFF
            continue
        elif gender == GENDER_MALE and not pid & 0xFF >= 191:
            pid += 191 - pid & 0xFF
            continue

        if ability == ABILITY_ONE and not pid & 1 == 0:
            pid += 1
            continue
        elif ability == ABILITY_TWO and not pid & 1 == 1:
            pid += 1
            continue

        if nature is not None and not pid % 25 == nature:
            pid += (25 + nature - pid % 25) % 25
            continue

        if is_shiny and not (pid >> 16 & 0xFFFF) ^ (pid & 0xFFFF) ^ tid ^ sid < 8:
            pid += 1
            continue

        if unown_letter is not None and not \
                (pid >> 18 & 0xC0
                 | pid >> 12 & 0x30
                 | pid >> 6 & 0x0C
                 | pid & 0x03) % 28 == unown_letter:
            pid += 1
            continue

        if wurmple_evo == WURMPLE_EVO_SILCOON and not (pid >> 16) % 10 <= 4:
            pid += 1
            continue
        elif wurmple_evo == WURMPLE_EVO_CASCOON and not (pid >> 16) % 10 > 4:
            pid += 1
            continue

        return pid

    return None


def get_str_input(prompt, valid_responses=None):
    while True:
        response = input(prompt)
        if response and (valid_responses is None or response in valid_responses):
            return response
        print(f'"{response}" is an invalid response')


def get_int_input(prompt, min_val=None, max_val=None):
    while True:
        response = input(prompt)
        if response.isdigit():
            response = int(response)
            if min_val is None and max_val is None:
                return response
            if min_val is None and response <= max_val:
                return response
            if max_val is None and response >= min_val:
                return response
            if min_val <= response <= max_val:
                return response
        print(f'"{response}" is an invalid response')


def main():
    print("Generation 3 Personality Value (PID) Generator")
    print("##############################################")
    print()

    # Gender
    should_gender = get_str_input('Restrict by gender? ("y" or "n"): ', "yn")
    if should_gender == "y":
        gender = get_int_input(f'Enter the desired Gender (\"{GENDER_MALE}\" for Male, \"{GENDER_FEMALE}\" for '
                               f'Female): ', GENDER_MALE, GENDER_FEMALE)
    else:  # should_gender == "n":
        gender = None

    # Ability
    should_ability = get_str_input('Restrict by ability? ("y" or "n"): ', "yn")
    if should_ability == "y":
        ability = get_int_input(f'Enter the desired ability (\"{ABILITY_ONE}\" for first ability, '
                                f'\"{ABILITY_TWO}\" for second ability): ', ABILITY_ONE, ABILITY_TWO)
    else:  # should_ability == "n":
        ability = None

    # Nature
    should_nature = get_str_input('Restrict by nature? ("y", or "n"): ', "yn")
    if should_nature == "y":
        print("Natures:\n"
              "        |-ATK       |-DEF       |-SPD        |-SpA        |-SpD\n"
              "    +ATK|Hardy  (0) |Lonely (1) |Brave   (2) |Adamant (3) |Naughty (4)\n"
              "    +DEF|Bold   (5) |Docile (6) |Relaxed (7) |Impish  (8) |Lax     (9)\n"
              "    +SPD|Timid  (10)|Hasty  (11)|Serious (12)|Jolly   (13)|Naive   (14)\n"
              "    +SpA|Modest (15)|Mild   (16)|Quiet   (17)|Bashful (18)|Rash    (19)\n"
              "    +SpD|Calm   (20)|Gentle (21)|Sassy   (22)|Careful (23)|Quirky  (24)")
        nature = get_int_input(f'Enter the desired nature by its number: ', 0, 24)
    else:  # should_nature == "n":
        nature = random.randint(0, 24)  # because if you just skip everything you may well just always get the first hit

    # Shiny
    should_shiny = get_str_input('Restrict by shiny? ("y" or "n"): ', "yn")
    if should_shiny == "y":
        shiny = True
        tid = get_int_input(f'Enter your Trainer ID (0 to {0xFFFF}): ', 0, 0xFFFF)
        sid = get_int_input(f'Enter your Secret ID (0 to {0xFFFF}): ', 0, 0xFFFF)
    else:  # should_shiny == "n":
        shiny = False
        tid = None
        sid = None

    # Unown letter
    should_unown = get_str_input('Restrict by Unown letter? ("y" or "n"): ', "yn")
    if should_unown == "y":
        unown_letter = get_str_input(f'Enter the desired Unown letter (\"a\" to \"z\", \"!\", or \"?\"): ',
                                     "abcdefghijklmnopqrstuvwxyz!?")
    else:  # should_unown == "n":
        unown_letter = None

    # Wurmple evolution
    should_wurmple = get_str_input('Restrict by Wurmple evolution? ("y" or "n"): ', "yn")
    if should_wurmple == "y":
        wurmple_evo = get_int_input(f'Enter the desired Wurmple evolution (\"{WURMPLE_EVO_SILCOON}\" for '
                                    f'Silcoon, \"{WURMPLE_EVO_CASCOON}\" for Cascoon): ',
                                    WURMPLE_EVO_SILCOON, WURMPLE_EVO_CASCOON)
    else:  # should_wurmple == "n":
        wurmple_evo = None

    pid = brute_force_pid(gender, ability, nature, shiny, tid, sid, unown_letter, wurmple_evo)
    if pid is not None:
        print(f"Found matching PID: 0x{pid:08X}")
    else:
        print("Failed to find a matching PID")


if __name__ == "__main__":
    main()