import sqlite3
import os
from collections import defaultdict

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

    def get_all_data(self):
        cursor = self.conn.cursor()
        cursor.execute("SELECT * FROM games ORDER BY generation, id")
        games = cursor.fetchall()
        
        cursor.execute("SELECT * FROM pokemon_forms ORDER BY generation, national_dex")
        pokemon_forms = cursor.fetchall()
        
        cursor.execute("SELECT * FROM encounters")
        encounters = cursor.fetchall()
        
        cursor.execute("SELECT * FROM evolution_chains")
        evolutions = cursor.fetchall()
        
        return games, pokemon_forms, encounters, evolutions

    def build_data_structures(self):
        games, pokemon_forms, encounters, evolutions = self.get_all_data()
        
        pokemon_by_gen = defaultdict(list)
        game_by_gen = defaultdict(list)
        encounter_data = defaultdict(list)
        evolution_map = 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)
        
        for encounter in encounters:
            encounter_data[encounter['game_id']].append(encounter)
        
        for evolution in evolutions:
            evolution_map[evolution['from_pfic']].append(evolution['to_pfic'])
        
        return pokemon_by_gen, game_by_gen, encounter_data, evolution_map

    def generate_efficient_plan(self, generation_groups):
        pokemon_by_gen, game_by_gen, encounter_data, evolution_map = self.build_data_structures()
        plan = []
        caught_pokemon = set()

        for group in generation_groups:
            group_plan = self.plan_for_group(group, pokemon_by_gen, game_by_gen, encounter_data, evolution_map, caught_pokemon)
            plan.extend(group_plan)

        return plan

    def plan_for_group(self, generations, pokemon_by_gen, game_by_gen, encounter_data, evolution_map, caught_pokemon):
        group_plan = []
        needed_pokemon = set()
        games_in_group = []

        for gen in generations:
            needed_pokemon.update(pokemon['PFIC'] for pokemon in pokemon_by_gen[gen])
            games_in_group.extend(game_by_gen[gen])

        # Check for new evolutions of already caught Pokémon
        new_evolutions = set()
        for pfic in caught_pokemon:
            if pfic in evolution_map:
                for evolved_pfic in evolution_map[pfic]:
                    if self.get_pokemon_generation(evolved_pfic) in generations and evolved_pfic not in caught_pokemon:
                        new_evolutions.add(pfic)
                        needed_pokemon.add(evolved_pfic)
                        print(f"New evolution: {self.get_pokemon_name(pfic)} into {self.get_pokemon_name(evolved_pfic)}")

        games_in_group.sort(key=lambda g: len([e for e in encounter_data[g['id']] if e['pfic'] in needed_pokemon]), reverse=True)

        for game in games_in_group:
            game_plan = self.plan_for_game(game, needed_pokemon, encounter_data, evolution_map, caught_pokemon, set(generations), new_evolutions)
            if game_plan:
                group_plan.extend(game_plan)

        return group_plan

    def plan_for_game(self, game, needed_pokemon, encounter_data, evolution_map, caught_pokemon, target_generations, new_evolutions):
        game_plan = []
        to_catch = defaultdict(int)
        to_evolve = defaultdict(list)
        to_breed = defaultdict(list)
        planned_evolutions = defaultdict(int)
        processed_chains = set()

        # First, handle new evolutions for already caught Pokémon
        for pfic in new_evolutions:
            evolution_chains, visit_counts, baby_forms = self.get_evolution_chains(pfic, evolution_map, target_generations)
            
            new_evolutions_count = sum(1 for chain in evolution_chains for pokemon in chain if pokemon not in caught_pokemon)
            
            if new_evolutions_count > 0:
                if any(encounter['pfic'] == pfic for encounter in encounter_data[game['id']]):
                    to_catch[pfic] += new_evolutions_count - planned_evolutions[pfic]
                    planned_evolutions[pfic] += new_evolutions_count
                
                    for chain in evolution_chains:
                        if len(chain) > 1:
                            for i in range(len(chain) - 1):
                                from_pfic, to_pfic = chain[i], chain[i+1]
                                if to_pfic not in caught_pokemon:
                                    to_evolve[from_pfic].append(to_pfic)
                                    caught_pokemon.add(to_pfic)
                                    needed_pokemon.discard(to_pfic)

        # Then proceed with the regular planning
        for encounter in encounter_data[game['id']]:
            if encounter['pfic'] in needed_pokemon and encounter['pfic'] not in caught_pokemon:
                evolution_chains, visit_counts, baby_forms = self.get_evolution_chains(encounter['pfic'], evolution_map, target_generations)
                base_form = evolution_chains[0][0]
                
                if base_form not in caught_pokemon:
                    chain_key = tuple(sorted(set(pokemon for chain in evolution_chains for pokemon in chain)))
                    if chain_key not in processed_chains:
                        processed_chains.add(chain_key)
                        
                        num_to_catch = max(0, visit_counts[base_form] - planned_evolutions[base_form])
                        if num_to_catch > 0:
                            to_catch[encounter['pfic']] += num_to_catch
                            planned_evolutions[base_form] += num_to_catch

                        for chain in evolution_chains:
                            if len(chain) > 1:
                                for i in range(len(chain) - 1):
                                    from_pfic, to_pfic = chain[i], chain[i+1]
                                    if from_pfic not in baby_forms:
                                        if not self.is_final_evolution(to_pfic, evolution_map) or (self.is_final_evolution(to_pfic, evolution_map) and not self.is_in_to_evolve(to_pfic, to_evolve)):
                                            to_evolve[from_pfic].append(to_pfic)
                                        else:
                                            to_catch[encounter['pfic']] += -1
                                            planned_evolutions[base_form] += -1

                                    if from_pfic in baby_forms and from_pfic not in encounter_data[game['id']]:
                                        if to_pfic not in to_breed or from_pfic not in to_breed[to_pfic]:
                                            to_breed[to_pfic].append(from_pfic)

                        caught_pokemon.update(pfic for chain in evolution_chains for pfic in chain)
                        needed_pokemon.difference_update(pfic for chain in evolution_chains for pfic in chain)

        if to_catch or to_evolve or to_breed:
            game_plan.append(f"Play {game['name']}:")
            if to_catch:
                game_plan.append("  Catch:")
                for pfic, count in to_catch.items():
                    game_plan.append(f"    - {self.get_pokemon_name(pfic)}: {count} time(s)")
            if to_evolve:
                game_plan.append("  Evolve:")
                for from_pfic, to_pfics in to_evolve.items():
                    for to_pfic in to_pfics:
                        game_plan.append(f"    - {self.get_pokemon_name(from_pfic)} into {self.get_pokemon_name(to_pfic)}")
            if to_breed:
                game_plan.append("  Breed:")
                for parent_pfic, baby_pfics in to_breed.items():
                    for baby_pfic in baby_pfics:
                        game_plan.append(f"    - {self.get_pokemon_name(parent_pfic)} to get {self.get_pokemon_name(baby_pfic)}")

        return game_plan

    def plan_evolutions(self, evolution_chains, evolution_map, caught_pokemon):
        chains, visit_counts = self.get_evolution_chains(evolution_chains[0][0], evolution_map)
        evolution_plan = []
        needed_counts = defaultdict(int)
        
        # Count how many of each Pokémon we need based on visit counts
        for pfic, count in visit_counts.items():
            if pfic not in caught_pokemon:
                needed_counts[pfic] = count
        
        # Plan evolutions
        for chain in chains:
            for i in range(len(chain) - 1):
                from_pfic, to_pfic = chain[i], chain[i+1]
                if needed_counts[to_pfic] > 0:
                    evolution_plan.append((from_pfic, to_pfic))
                    needed_counts[to_pfic] -= 1
                    needed_counts[from_pfic] -= 1
        
        return evolution_plan, needed_counts

    def get_evolution_chains(self, pfic, evolution_map, target_generations):
        visited = defaultdict(int)
        unique_paths = set()
        baby_forms = set()
        
        def dfs(current_pfic, path):
            path_tuple = tuple(path)
            if path_tuple in unique_paths:
                return
            
            unique_paths.add(path_tuple)
            for pfic in path:
                visited[pfic] += 1
                if self.is_baby_pokemon(pfic):
                    baby_forms.add(pfic)
            
            if current_pfic in evolution_map:
                for next_pfic in evolution_map[current_pfic]:
                    if self.get_pokemon_generation(next_pfic) in target_generations:
                        dfs(next_pfic, path + [next_pfic])

        dfs(pfic, [pfic])
        chains = [list(path) for path in unique_paths]
        
        # Adjust visit counts for baby Pokémon
        for baby in baby_forms:
            visited[baby] = 1
        
        return chains, visited, baby_forms

    def get_pokemon_name(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT name, form_name FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        if result['form_name']:
            return f"{result['name']} ({result['form_name']})"
        return result['name']

    def display_plan(self, generation_groups, output_file='efficiency_plan.txt'):
        plan = self.generate_efficient_plan(generation_groups)
        
        # Print to console
        for step in plan:
            print(step)
        
        # Write to file
        with open(output_file, 'w', encoding='utf-8') as f:
            for step in plan:
                f.write(f"{step}\n")
        
        print(f"\nPlan has been written to {os.path.abspath(output_file)}")

    def is_baby_pokemon(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT is_baby_form FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        return result['is_baby_form'] if result else False

    def get_pokemon_generation(self, pfic):
        cursor = self.conn.cursor()
        cursor.execute("SELECT generation FROM pokemon_forms WHERE PFIC = ?", (pfic,))
        result = cursor.fetchone()
        return result['generation'] if result else None

    def is_final_evolution(self, pfic, evolution_map):
        return pfic not in evolution_map

    def is_in_to_evolve(self, pfic, to_evolve):
        return any(pfic in to_pfics for to_pfics in to_evolve.values())

# Example usage
planner = EfficiencyOriginDexPlanner('pokemon_forms.db')
planner.display_plan([[1, 2], [3, 4, 5, 6], [7], [8], [9]])