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

def main():
    # 1. Sidebar for Dropdown, Buttons, and Instructions
    st.sidebar.title("AI Energy Score Label Generator")
    st.sidebar.write("### Instructions:")
    st.sidebar.write("1. Select a model from the dropdown.")
    st.sidebar.write("2. Review the label preview.")
    st.sidebar.write("3. Download the label as a PNG.")
    st.sidebar.markdown("[Learn more about AI Energy Scores](https://example.com)")

    # 2. 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

    # 3. 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

    # 4. Create a dropdown list based on unique values in the 'model' column
    model_options = data_df["model"].unique().tolist()
    selected_model = st.sidebar.selectbox("Select a Model:", model_options)

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

    # 6. Dynamically select the background image based on the score
    try:
        score = int(model_data["score"])  # Convert to int to avoid issues
        background_path = f"{score}.png"  # E.g., "1.png", "2.png"
        background = Image.open(background_path).convert("RGBA")
    except FileNotFoundError:
        st.sidebar.error(f"Could not find background image '{score}.png'. Using default background.")
        background = Image.open("default_background.png").convert("RGBA")
    except ValueError:
        st.sidebar.error(f"Invalid score '{model_data['score']}'. Score must be an integer.")
        return

    # 7. Overlay the data on the image
    generated_label = create_label(background, model_data)

    # 8. Display the generated label in the main area
    st.image(generated_label, caption="Generated Label Preview")

    # 9. Provide a download button in the sidebar
    img_buffer = io.BytesIO()
    generated_label.save(img_buffer, format="PNG")
    img_buffer.seek(0)

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

def create_label(background_image, model_data):
    """
    This function takes a background image and a row (model_data) from the CSV,
    then draws text on the image. Finally, returns the modified image object.
    """
    label_img = background_image.convert("RGBA")
    draw = ImageDraw.Draw(label_img)

    # Load the Inter variable font
    try:
        inter_font_path = "Inter-VariableFont_opsz,wght.ttf"
        title_font = ImageFont.truetype(inter_font_path, 16, layout_engine=ImageFont.LAYOUT_RAQM)
        details_font = ImageFont.truetype(inter_font_path, 12, layout_engine=ImageFont.LAYOUT_RAQM)
        energy_font = ImageFont.truetype(inter_font_path, 14, layout_engine=ImageFont.LAYOUT_RAQM)
    except Exception as e:
        st.error(f"Font loading failed: {e}")
        return label_img

    # Define positions for the text groups (easy to tweak!)
    title_x, title_y = 20, 20  # Position for title group
    details_x, details_y = 300, 20  # Position for details group
    energy_x, energy_y = 150, 400  # Position for energy group

    # Group 1: Title (Left-Justified)
    draw.text((title_x, title_y), f"Model: {model_data['model']}", font=title_font, fill="black")
    draw.text((title_x, title_y + 25), f"Provider: {model_data['provider']}", font=title_font, fill="black")

    # Group 2: Details (Right-Justified)
    details_lines = [
        f"Date: {model_data['date']}",
        f"Task: {model_data['task']}",
        f"Hardware: {model_data['hardware']}"
    ]
    for i, line in enumerate(details_lines):
        draw.text((details_x, details_y + i * 20), line, font=details_font, fill="black", anchor="ra")

    # Group 3: Energy (Bottom-Center)
    energy_text = f"Energy: {model_data['energy']}"
    draw.text((energy_x, energy_y), energy_text, font=energy_font, fill="black")

    return label_img

if __name__ == "__main__":
    main()