import sqlite3
from collections import defaultdict

class OriginDexPlanner:
    def __init__(self, db_path):
        self.conn = sqlite3.connect(db_path)
        self.conn.row_factory = sqlite3.Row

    def get_game_data(self):
        """Fetch game and pokemon information from the database."""
        cursor = self.conn.cursor()
        cursor.execute("SELECT * FROM games")
        games = cursor.fetchall()
        
        cursor.execute("SELECT * FROM pokemon_forms WHERE name IN ('Elekid', 'Electabuzz', 'Electivire')")
        pokemon_forms = cursor.fetchall()
        
        cursor.execute("SELECT * FROM encounters")
        encounters = cursor.fetchall()
        
        exclusive_encounters = []
        cursor.execute("SELECT * FROM evolution_chains WHERE from_pfic IN (SELECT PFIC FROM pokemon_forms WHERE name IN ('Elekid', 'Electabuzz', 'Electivire'))")
        evolutions = cursor.fetchall()
        
        return games, pokemon_forms, encounters, exclusive_encounters, evolutions

    def build_pokemon_data(self):
        """Construct data dictionaries to store relationships between the Pokemon and encounters."""
        games, pokemon_forms, encounters, exclusive_encounters, evolutions = self.get_game_data()
        
        pokemon_by_gen = defaultdict(list)
        game_by_gen = defaultdict(list)
        encounter_data = defaultdict(list)
        
        for game in games:
            game_by_gen[game['generation']].append(game)

        for pokemon in pokemon_forms:
            pokemon_by_gen[pokemon['generation']].append(pokemon)
        
        # Build a dict of encounters, mapping game_id to Pokemon encountered in it
        for encounter in encounters:
            encounter_data[encounter['game_id']].append(encounter)
        
        evolution_map = defaultdict(list)
        for evolution in evolutions:
            from_pfic = evolution['from_pfic']
            to_pfic = evolution['to_pfic']
            # Lookup generation from pokemon_forms table
            cursor = self.conn.cursor()
            cursor.execute("SELECT generation FROM pokemon_forms WHERE PFIC = ?", (to_pfic,))
            to_generation = cursor.fetchone()['generation']
            evolution_map[from_pfic].append((to_pfic, to_generation))
        
        exclusive_group_map = defaultdict(list)
        for exclusive_encounter in exclusive_encounters:
            exclusive_group_map[exclusive_encounter['group_id']].append(exclusive_encounter['encounter_id'])
        
        return pokemon_by_gen, game_by_gen, encounter_data, evolution_map, exclusive_group_map

    def generate_plan(self, generations_to_group_list=None):
        """Generate a game plan to catch all Pokemon in their origin generation."""
        if generations_to_group_list is None:
            generations_to_group_list = []
        
        pokemon_by_gen, game_by_gen, encounter_data, evolution_map, exclusive_group_map = self.build_pokemon_data()
        
        plan = []
        pokemon_to_catch = defaultdict(list)
        caught_pokemon = set()
        pokemon_names = {}
        
        # Handle multiple sets of grouped generations
        grouped_generations_sets = [set(generations) for generations in generations_to_group_list]
        for grouped_generations in grouped_generations_sets:
            grouped_games = []
            grouped_pokemon_list = []
            for gen in grouped_generations:
                if gen in game_by_gen:
                    grouped_games.extend(game_by_gen[gen])
                    grouped_pokemon_list.extend(pokemon_by_gen[gen])
            
            if grouped_games:
                plan.extend(self.generate_plan_for_generation(grouped_games, grouped_pokemon_list, encounter_data, evolution_map, exclusive_group_map, pokemon_names, caught_pokemon, game_by_gen))
        
        # Loop through each generation, generating a plan for each game
        for gen in sorted(game_by_gen.keys()):
            # Skip generations that have been grouped
            if any(gen in grouped_generations for grouped_generations in grouped_generations_sets):
                continue
            games = game_by_gen[gen]
            pokemon_list = pokemon_by_gen[gen]
            
            plan.extend(self.generate_plan_for_generation(games, pokemon_list, encounter_data, evolution_map, exclusive_group_map, pokemon_names, caught_pokemon, game_by_gen))
        
        return plan
    
    def generate_plan_for_generation(self, games, pokemon_list, encounter_data, evolution_map, exclusive_group_map, pokemon_names, caught_pokemon, game_by_gen):
        """Generate a game plan for a specific generation or group of generations."""
        plan = []
        pokemon_to_catch = defaultdict(list)
        pokemon_to_breed = defaultdict(list)
        catchable_pokemon = set()
        
        # Create a set of baby Pokémon PFICs
        baby_pokemon = set(pokemon['PFIC'] for pokemon in pokemon_list if pokemon['is_baby_form'])
        
        # Identify unique evolution chains and the starting point
        for pokemon in pokemon_list:
            pfic = pokemon['PFIC']
            pokemon_names[pfic] = f'{pokemon["name"]} ({pokemon["form_name"] if pokemon["form_name"] is not None else ""})'
            if pfic not in evolution_map:  # Not evolved from anything
                catchable_pokemon.add(pfic)
            else:
                catchable_pokemon.add(pfic)
                # Add all evolution stages to catchable list if they belong to the current or previous generations
                current = pfic
                current_generation = pokemon['generation']
                while current in evolution_map:
                    next_pfic, next_generation = evolution_map[current][0]
                    if any(p['PFIC'] == next_pfic for p in pokemon_list):
                        if next_generation <= current_generation:
                            catchable_pokemon.add(next_pfic)
                    current = next_pfic
        
        # Determine the best game to catch the most Pokemon
        best_game = None
        max_catchable = 0
        game_catchable_pokemon = defaultdict(set)
        for game in games:
            game_id = game['id']
            encounters = encounter_data[game_id]
            
            # Track how many unique Pokemon are available in this game
            for encounter in encounters:
                pfic = encounter['pfic']
                if pfic in catchable_pokemon:
                    game_catchable_pokemon[game_id].add(pfic)
            
            # Update the best game if this game has more catchable Pokemon
            if len(game_catchable_pokemon[game_id]) > max_catchable:
                max_catchable = len(game_catchable_pokemon[game_id])
                best_game = game
        
        # Use the best game to catch as many unique Pokemon as possible
        remaining_games = [game for game in games if game != best_game]
        if best_game:
            game_id = best_game['id']
            game_name = best_game['name']
            encounters = encounter_data[game_id]
            
            for encounter in encounters:
                pfic = encounter['pfic']
                if pfic in catchable_pokemon and pfic not in caught_pokemon:
                    evolution_chain = self.get_full_evolution_chain(pfic, evolution_map)
                    base_pokemon = evolution_chain[0]
                    if base_pokemon not in caught_pokemon:
                        # Calculate the number of Pokemon needed for the full evolution chain
                        count = len(evolution_chain)
                        if encounter['static_encounter']:
                            count = min(count, encounter['static_encounter_count'] or 1)
                        
                        if base_pokemon in baby_pokemon:
                            # For baby Pokémon, catch the first non-baby evolution and breed for the baby
                            non_baby_pokemon = next((p for p in evolution_chain if p not in baby_pokemon), None)
                            if non_baby_pokemon:
                                pokemon_to_catch[game_name].append({'pfic': non_baby_pokemon, 'count': 1})
                                pokemon_to_breed[game_name].append({'pfic': base_pokemon, 'count': 1})
                        else:
                            pokemon_to_catch[game_name].append({'pfic': base_pokemon, 'count': count})
                        
                        caught_pokemon.update(evolution_chain)
                    
                    # Account for exclusive encounters
                    exclusive_group_ids = [group['group_id'] for group in exclusive_group_map if group['encounter_id'] == encounter['id']]
                    if exclusive_group_ids:
                        other_options = exclusive_group_map[exclusive_group_ids[0]]
                        if len(other_options) > 1:
                            # Pick one of the options, mark the rest to catch in another game
                            catchable_pokemon.remove(pfic)
                        else:
                            continue
        
        # Use remaining games to catch any remaining Pokemon
        for game in remaining_games:
            game_id = game['id']
            game_name = game['name']
            encounters = encounter_data[game_id]
            
            for encounter in encounters:
                pfic = encounter['pfic']
                if pfic in catchable_pokemon and pfic not in caught_pokemon:
                    evolution_chain = self.get_full_evolution_chain(pfic, evolution_map)
                    base_pokemon = evolution_chain[0]
                    if base_pokemon not in caught_pokemon:
                        # Calculate the number of Pokemon needed for the full evolution chain
                        count = len(evolution_chain)
                        if encounter['static_encounter']:
                            count = min(count, encounter['static_encounter_count'] or 1)
                        
                        if base_pokemon in baby_pokemon:
                            # For baby Pokémon, catch the first non-baby evolution and breed for the baby
                            non_baby_pokemon = next((p for p in evolution_chain if p not in baby_pokemon), None)
                            if non_baby_pokemon:
                                pokemon_to_catch[game_name].append({'pfic': non_baby_pokemon, 'count': 1})
                                pokemon_to_breed[game_name].append({'pfic': base_pokemon, 'count': 1})
                        else:
                            pokemon_to_catch[game_name].append({'pfic': base_pokemon, 'count': count})
                        
                        caught_pokemon.update(evolution_chain)
        
        # Handle new evolutions or forms added in later generations
        for pfic in list(caught_pokemon):
            if pfic in evolution_map:
                for new_evolution, new_generation in evolution_map[pfic]:
                    # Check if the new evolution is from a later generation
                    full_evolution_chain = self.get_full_evolution_chain(pfic, evolution_map)
                    current_generation = min((pokemon['generation'] for pokemon in pokemon_list if pokemon['PFIC'] in full_evolution_chain), default=None)
                    if new_generation and current_generation and new_generation > current_generation:
                        # Add the baby/base form to be caught in the new generation if required
                        base_pokemon = full_evolution_chain[0]
                        new_gen_games = game_by_gen[new_generation]
                        for game in new_gen_games:
                            game_id = game['id']
                            game_name = game['name']
                            encounters = encounter_data[game_id]
                            for encounter in encounters:
                                if encounter['pfic'] == base_pokemon and base_pokemon not in caught_pokemon:
                                    if base_pokemon in baby_pokemon:
                                        non_baby_pokemon = next((p for p in full_evolution_chain if p not in baby_pokemon), None)
                                        if non_baby_pokemon:
                                            pokemon_to_catch[game_name].append({'pfic': non_baby_pokemon, 'count': 1})
                                            pokemon_to_breed[game_name].append({'pfic': base_pokemon, 'count': 1})
                                    else:
                                        pokemon_to_catch[game_name].append({'pfic': base_pokemon, 'count': 1})
                                    caught_pokemon.add(base_pokemon)
                                    break
        
        # Generate the plan for this generation or group of generations
        for game_name in set(list(pokemon_to_catch.keys()) + list(pokemon_to_breed.keys())):
            plan.append(f"Play {game_name}:")
            if game_name in pokemon_to_catch:
                plan.append(f"  Catch {len(pokemon_to_catch[game_name])} unique Pokemon:")
                for pokemon in pokemon_to_catch[game_name]:
                    plan.append(f"    - {pokemon_names[pokemon['pfic']]}: {pokemon['count']} times")
            
            if game_name in pokemon_to_breed:
                plan.append(f"  Breed {len(pokemon_to_breed[game_name])} Pokemon:")
                for pokemon in pokemon_to_breed[game_name]:
                    plan.append(f"    - {pokemon_names[pokemon['pfic']]}: {pokemon['count']} times")
        
        return plan
    
    def get_full_evolution_chain(self, pfic, evolution_map):
        """Get the full evolution chain for a given PFIC."""
        chain = [pfic]
        while pfic in evolution_map:
            pfic = evolution_map[pfic][0][0]
            chain.append(pfic)
        return chain
    
    def display_plan(self, generations_to_group_list=None):
        plan = self.generate_plan(generations_to_group_list)
        for step in plan:
            print(step)

# Example usage
planner = OriginDexPlanner('pokemon_forms.db')
planner.display_plan(generations_to_group_list=[[1, 2], [3, 4, 5, 6]])