import streamlit as st
import cv2
from google.generativeai import Client
import os
from rubik_solver import Solver

# API Key setup
os.environ["GOOGLE_API_KEY"] = "AIzaSyCttG4s8Tue8RpSb2mK1tvkJSp17YETGGk"  # Replace with your actual API key
client = Client()

# App title and instructions
st.title("Rubik's Cube Solver")
st.write("Upload images of each cube face for solution.")

# Uploaders for each face
face_images = [st.file_uploader("Face {}".format(i+1)) for i in range(6)]

# Submit button
submit_button = st.button("Solve Cube!")

if submit_button:
    try:
        if all(face_images):
            cube_state = []
            for face in face_images:
                # Read image
                image = cv2.imdecode(np.fromstring(face.read(), np.uint8), cv2.IMREAD_UNCHANGED)

                # Use Gemini Pro API for color detection
                response = client.generate_text(
                    prompt="Identify the colors on this Rubik's cube face",
                    image=image,
                    model="gemini-pro-vision"
                )
                colors = response.text.splitlines()  # Parse response for colors

                # Add face colors to cube state
                cube_state.append(colors)

            # Check for successful color detection
            if all(cube_state):
                st.write("Colors detected! Solving...")
            else:
                st.error("Some faces had errors! Retry with clear images.")
        else:
            st.error("Please upload all 6 face images.")

        if all(cube_state):
            # Create solver object
            solver = Solver(cube_state)

            # Solve the cube and get moves
            solution = solver.solve()

            # Display solution steps
            st.write("Solution:")
            for step, move in enumerate(solution, start=1):
                st.write(f"Step {step}: {move}")

                # Use animation or visuals to guide user (optional)

                # Allow user to proceed step-by-step
                if not st.button(f"Next Step ({step}/{len(solution)})"):
                    break

            st.success("Cube solved!")

    except Exception as e:
        st.error("An error occurred: {}".format(e))