import streamlit as st
from similarity import PokemonSimilarity
from PIL import Image
import io
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Set page config first
st.set_page_config(
    page_title="Pokemon Similarity Finder",
    page_icon="🎮",
    layout="centered"
)

# Initialize the similarity engine
@st.cache_resource
def get_similarity_engine():
    logger.info("Initializing similarity engine...")
    engine = PokemonSimilarity()
    logger.info("Similarity engine initialized successfully")
    return engine

similarity_engine = get_similarity_engine()

# Title and description
st.title("🎮 Pokemon Similarity Finder")
st.markdown("""
Upload an image of a Pokemon or provide an image URL and we'll find the closest match in our database!
""")

# File uploader and URL input
uploaded_file = st.file_uploader("Choose a Pokemon image...", type=["jpg", "jpeg", "png"])
image_url = st.text_input("Or enter an image URL:")

# Determine which input to use
image_input = None
input_type = None
if uploaded_file is not None:
    image_input = uploaded_file.getvalue()
    input_type = 'file'
elif image_url:
    image_input = image_url
    input_type = 'url'

if image_input is not None:
    if input_type == 'file':
        logger.info(f"File uploaded: {uploaded_file.name}")
        try:
            # Display the uploaded image
            image = Image.open(io.BytesIO(image_input))
            st.image(image, caption="Uploaded Image", use_column_width=True)
            logger.info("Successfully displayed uploaded image")
        except Exception as e:
            logger.error(f"Error loading image: {str(e)}")
            st.error(f"❌ Error loading image: {str(e)}")
            st.info("Please make sure you've uploaded a valid image file.")
    elif input_type == 'url':
        logger.info(f"Image URL provided: {image_url}")
        st.image(image_url, caption="Image from URL", use_column_width=True)

    # Add a button to trigger the similarity search
    if st.button("Find Similar Pokemon"):
        logger.info("Find Similar Pokemon button clicked")
        with st.spinner("Analyzing image..."):
            try:
                logger.info(f"Finding closest Pokemon match using {input_type} input...")
                pokemon_name = similarity_engine.find_closest_pokemon(image_input)
                logger.info(f"Found closest Pokemon: {pokemon_name}")
                st.success(f"🎯 The closest Pokemon is: **{pokemon_name}**")
                st.balloons()
            except Exception as e:
                logger.error(f"Error during Pokemon matching: {str(e)}")
                st.error(f"❌ Error: {str(e)}")
                st.info("Please try uploading a different image, using a different URL, or try again later.")