from machine import I2C, Pin
import ssd1306
import time

# Initialize I2C for OLED display
i2c = I2C(1, scl=Pin(22), sda=Pin(21), freq=400000)

# Initialize OLED Display
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# Button setup
up_button = Pin(0, Pin.IN, Pin.PULL_UP)
down_button = Pin(4, Pin.IN, Pin.PULL_UP)
select_button = Pin(16, Pin.IN, Pin.PULL_UP)
back_button = Pin(2, Pin.IN, Pin.PULL_UP)

# Soldiers, current party, menu variables, and state tracking
soldiers = [
    {"name": "Soldier 1", "HP": 100, "Def": 30, "Atk": 20, "Alignment": "Good", "Acquired": False},
    {"name": "Soldier 2", "HP": 90, "Def": 40, "Atk": 25, "Alignment": "Neutral", "Acquired": False},
    {"name": "Soldier 3", "HP": 110, "Def": 20, "Atk": 30, "Alignment": "Evil", "Acquired": True},
]
current_party = []
menu_items = ["DefDex", "Team"]
team_menu_items = ["Add Defmon", "Remove Defmon", "View Team"]
current_menu_selection = 1
current_team_menu_selection = 1
defdex_selection = 1
menu_active = True
in_defdex = False
in_team_menu = False
adding_defmon = False
removing_defmon = False

# Load and save functions (placeholder, implement based on your storage capabilities)
def load_acquired_soldiers():
    try:
        with open('acquired_soldiers.txt', 'r') as file:
            acquired_names = [line.strip() for line in file]
        for soldier in soldiers:
            soldier["Acquired"] = soldier["name"] in acquired_names
    except OSError:
        pass

def save_acquired_soldiers():
    with open('acquired_soldiers.txt', 'w') as file:
        for soldier in soldiers:
            if soldier["Acquired"]:
                file.write(soldier["name"] + '\n')

def load_current_party():
    global current_party
    try:
        with open('current_party.txt', 'r') as file:
            current_party = [line.strip() for line in file.readlines()]
    except OSError:
        current_party = []  # File doesn't exist; start with an empty list


def save_current_party():
    with open('current_party.txt', 'w') as file:
        for soldier_name in current_party:
            file.write(soldier_name + '\n')

load_acquired_soldiers()
load_current_party()

# Display functions
def show_menu():
    oled.fill(0)
    for index, item in enumerate(menu_items, start=1):
        if index == current_menu_selection:
            oled.text("> " + item, 0, (index - 1) * 10)
        else:
            oled.text("  " + item, 0, (index - 1) * 10)
    oled.show()

def show_team_menu():
    oled.fill(0)
    for index, item in enumerate(team_menu_items, start=1):
        if index == current_team_menu_selection:
            oled.text("> " + item, 0, (index - 1) * 10)
        else:
            oled.text("  " + item, 0, (index - 1) * 10)
    oled.show()

def show_defdex():
    oled.fill(0)
    for index, soldier in enumerate(soldiers, start=1):
        display_name = "?????" if not soldier["Acquired"] else soldier["name"]
        if soldier["Acquired"]:
            display_name += " *"
        if index == defdex_selection:
            oled.text("> " + display_name, 0, (index - 1) * 10)
        else:
            oled.text("  " + display_name, 0, (index - 1) * 10)
    oled.show()
	
def view_current_team():
    oled.fill(0)
    oled.text("Current Team:", 0, 0)
    for index, soldier_name in enumerate(current_party, start=1):
        oled.text(f"{index}. {soldier_name}", 0, index * 10)
    oled.show()

# Update and button check functions
def update_menu_selection(delta):
    global current_menu_selection
    current_menu_selection += delta
    if current_menu_selection < 1:
        current_menu_selection = len(menu_items)
    elif current_menu_selection > len(menu_items):
        current_menu_selection = 1

def update_team_menu_selection(delta):
    global current_team_menu_selection
    current_team_menu_selection += delta
    if current_team_menu_selection < 1:
        current_team_menu_selection = len(team_menu_items)
    elif current_team_menu_selection > len(team_menu_items):
        current_team_menu_selection = 1

def update_defdex_selection(delta):
    global defdex_selection
    defdex_selection += delta
    if defdex_selection < 1:
        defdex_selection = len(soldiers)
    elif defdex_selection > len(soldiers):
        defdex_selection = 1

# Add and remove Defmon functions
def add_defmon_to_team(selection):
    if len(current_party) < 3 and soldiers[selection - 1]["Acquired"]:
        current_party.append(soldiers[selection - 1]["name"])
        save_current_party()
        return True
    return False

def remove_defmon_from_team(selection):
    if soldiers[selection - 1]["name"] in current_party:
        current_party.remove(soldiers[selection - 1]["name"])
        save_current_party()
        return True
    return False

def check_buttons():
    global menu_active, in_defdex, in_team_menu, adding_defmon, removing_defmon
    global current_menu_selection, defdex_selection, current_team_menu_selection
    debounce_time = 0.2

    if not up_button.value():
        if menu_active:
            update_menu_selection(-1)
        elif in_defdex:
            update_defdex_selection(-1)
        elif in_team_menu:
            update_team_menu_selection(-1)
        time.sleep(debounce_time)
    elif not down_button.value():
        if menu_active:
            update_menu_selection(1)
        elif in_defdex:
            update_defdex_selection(1)
        elif in_team_menu:
            update_team_menu_selection(1)
        time.sleep(debounce_time)
    elif not select_button.value():
        if menu_active:
            if current_menu_selection == 1:  # DefDex selected
                menu_active = False
                in_defdex = True
            elif current_menu_selection == 2:  # Team selected
                menu_active = False
                in_team_menu = True
        elif in_defdex:
            # Assuming action for viewing or selecting within DefDex
            in_defdex = False
            menu_active = True  # Adjust based on desired action
        elif in_team_menu:
            if current_team_menu_selection == 1:  # View Current Team
                # Logic to display the current team
                print("Current Team: ", current_party)  # Example action
            elif current_team_menu_selection == 2:  # Add Defmon
                in_team_menu = False
                adding_defmon = True
            elif current_team_menu_selection == 3:  # Remove Defmon
                in_team_menu = False
                removing_defmon = True
        elif adding_defmon:
            if add_defmon_to_team(defdex_selection):
                adding_defmon = False
                in_team_menu = True  # Return to team menu after adding
        elif removing_defmon:
            if remove_defmon_from_team(defdex_selection):
                removing_defmon = False
                in_team_menu = True  # Return to team menu after removing
        time.sleep(debounce_time)
    elif not back_button.value():
        if in_defdex or adding_defmon or removing_defmon:
            in_defdex = False
            adding_defmon = False
            removing_defmon = False
            in_team_menu = True  # Return to team menu or main menu
        elif in_team_menu:
            in_team_menu = False
            menu_active = True
        time.sleep(debounce_time)

# Main loop
while True:
    if menu_active:
        show_menu()
    elif in_defdex:
        show_defdex()
    elif in_team_menu:
        show_team_menu()
    elif adding_defmon or removing_defmon:
        show_defdex()  # Use Defdex view for adding/removing Defmons
    check_buttons()