import streamlit as st
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import io

def main():
    # Sidebar logo and title
    st.sidebar.image("logo.png", use_container_width=True)  # Display the logo at the top
    st.sidebar.title("Label Generator")
    st.sidebar.markdown("For [AI Energy Score Leaderboard](https://huggingface.co/spaces/AIEnergyScore/Leaderboard)")


    # Initialize data_df
    data_df = None

    # Read Data from CSV
    try:
        data_df = pd.read_csv("data.csv")
    except FileNotFoundError:
        st.sidebar.error("Could not find 'data.csv'! Please make sure it's present.")
        return
    except Exception as e:
        st.sidebar.error(f"Error reading 'data.csv': {e}")
        return

    # Ensure the CSV has required columns
    required_columns = ["model", "provider", "date", "task", "hardware", "energy", "score"]
    for col in required_columns:
        if col not in data_df.columns:
            st.sidebar.error(f"The CSV file must contain a column named '{col}'.")
            return

    # Dropdown for selecting a model
    st.sidebar.write("### Instructions:")
    st.sidebar.write("#### 1. Select a model below")
    model_options = data_df["model"].unique().tolist()  # Get model options
    selected_model = st.sidebar.selectbox(
        "Scored Models", model_options, help="Start typing to search for a model"
    )  # Searchable dropdown

    # Add step 2 instructions and move the Download button
    st.sidebar.write("#### 2. Download the label")

    # Filter the data for the selected model
    model_data = data_df[data_df["model"] == selected_model].iloc[0]

    # Dynamically select the background image based on the score
    try:
        score = int(model_data["score"])  # Convert to int
        background_path = f"{score}.png"  # E.g., "1.png", "2.png"
        background = Image.open(background_path).convert("RGBA")

        # Proportional scaling to fit within the target size
        target_size = (800, 600)  # Maximum width and height
        background.thumbnail(target_size, Image.Resampling.LANCZOS)

    except FileNotFoundError:
        st.sidebar.error(f"Could not find background image '{score}.png'. Using default background.")
        background = Image.open("default_background.png").convert("RGBA")
        background.thumbnail(target_size, Image.Resampling.LANCZOS)  # Resize default image proportionally
    except ValueError:
        st.sidebar.error(f"Invalid score '{model_data['score']}'. Score must be an integer.")
        return

    # Generate the label with text
    generated_label = create_label(background, model_data)

    # Display the label
    st.image(generated_label, caption="Generated Label Preview")

    # Download button for the label
    img_buffer = io.BytesIO()
    generated_label.save(img_buffer, format="PNG")
    img_buffer.seek(0)

    st.sidebar.download_button(
        label="Download",
        data=img_buffer,
        file_name="AIEnergyScore.png",
        mime="image/png"
    )

    # Step 3 instructions
    st.sidebar.write("#### 3. Share your label in technical reports, announcements, etc.")

def create_label(background_image, model_data):
    """
    Create the label image by adding text from model_data to the background image.
    Render at high resolution and scale down for sharper output.
    """
    # Scale factor for high resolution rendering
    scale_factor = 2
    original_size = background_image.size
    high_res_size = (original_size[0] * scale_factor, original_size[1] * scale_factor)

    # Resize background proportionally for high resolution
    high_res_image = background_image.resize(high_res_size, Image.Resampling.LANCZOS)
    draw = ImageDraw.Draw(high_res_image)

    # Load fonts with scaled sizes
    try:
        title_font = ImageFont.truetype("Inter_24pt-Bold.ttf", size=22 * scale_factor)
        details_font = ImageFont.truetype("Inter_18pt-Regular.ttf", size=18 * scale_factor)
        energy_font = ImageFont.truetype("Inter_18pt-Medium.ttf", size=20 * scale_factor)
    except Exception as e:
        st.error(f"Font loading failed: {e}")
        return high_res_image

    # Define scaled positions for text
    title_x, title_y = 28 * scale_factor, 128 * scale_factor
    details_x, details_y = 375 * scale_factor, 210 * scale_factor
    energy_x, energy_y = 350 * scale_factor, 390 * scale_factor

    # Group 1: Title (Bold)
    draw.text((title_x, title_y), str(model_data['model']), font=title_font, fill="black")
    draw.text((title_x, title_y + 30 * scale_factor), str(model_data['provider']), font=title_font, fill="black")

    # Group 2: Details (Right-Justified)
    details_lines = [
        str(model_data['date']),
        str(model_data['task']),
        str(model_data['hardware'])
    ]
    for i, line in enumerate(details_lines):
        bbox = draw.textbbox((0, 0), line, font=details_font)
        text_width = bbox[2] - bbox[0]
        draw.text((details_x - text_width, details_y + i * 40 * scale_factor), line, font=details_font, fill="black")

    # Group 3: Energy (Bottom-Center)
    energy_text = str(model_data['energy'])
    bbox = draw.textbbox((0, 0), energy_text, font=energy_font)
    energy_text_width = bbox[2] - bbox[0]
    draw.text((energy_x - energy_text_width // 2, energy_y), energy_text, font=energy_font, fill="black")

    # Downscale to original size
    final_image = high_res_image.resize(original_size, Image.Resampling.LANCZOS)
    return final_image

if __name__ == "__main__":
    main()