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

def main():
    # 1. Page Title
    st.title("AI Energy Score Label Generator")

    # 2. Read Data from CSV
    # Using pandas to read the 'data.csv' file
    # Make sure 'data.csv' is in the same folder as 'app.py'
    try:
        data_df = pd.read_csv("data.csv")
    except FileNotFoundError:
        st.error("Could not find 'data.csv'! Please make sure it's present.")
        return

    # 3. Ensure the CSV has a "Model" column
    if "Model" not in data_df.columns:
        st.error("The CSV file must contain a column named 'Model'.")
        return

    # 4. Create a dropdown list based on unique values in the 'Model' column
    model_options = data_df["Model"].unique().tolist()
    selected_model = st.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. Load the background image for the label
    # Make sure 'background.png' is in the same folder as 'app.py'
    try:
        background = Image.open("background.png")
    except FileNotFoundError:
        st.error("Could not find 'background.png'! Please make sure it's present.")
        return

    # 7. Overlay the data on the image
    #    We'll create a function to do this cleanly.
    generated_label = create_label(background, model_data)

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

    # 9. Provide a download button
    #    We'll create an in-memory file to let user download the image.
    img_buffer = io.BytesIO()
    generated_label.save(img_buffer, format="PNG")
    img_buffer.seek(0)

    st.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.
    """
    # Convert background to a format that can be edited (RGBA mode).
    label_img = background_image.convert("RGBA")

    # Create a Drawing context
    draw = ImageDraw.Draw(label_img)

    # Choose a font and size. Change path/size as needed.
    # If you don't have a custom font file, you can use a PIL built-in font.
    try:
        font = ImageFont.truetype("Roboto-SemiBold.ttf", 30)
        print("Font loaded successfully!")
    except Exception as e:
        font = ImageFont.load_default()
        print(f"Font loading failed: {e}")

    # You can customize the positions, colors, etc. For instance:
    # We'll just place the text in a simple stacked format at a fixed position.
    # Position variables (x, y) - Adjust as needed.
    x_position = 75
    y_position = 350
    line_spacing = 50

    # Extracting data from the row. Customize these lines based on your CSV columns.
    model_name = f"Model: {model_data['Model']}"
    if 'Task' in model_data:
        task = f"Task: {model_data['Task']}"
    else:
        task = "Task: N/A"
    if 'date' in model_data:
        date = f"Date: {model_data['Date']}"
    else:
        date = "Date: N/A"
    if 'Energy' in model_data:
        energy = f"Energy: {model_data['Energy']}"
    else:
        energy = "Energy: N/A"

    text_lines = [model_name, task, date, energy]

    # Draw each line on the image
    for line in text_lines:
        draw.text((x_position, y_position), line, fill="black", font=font)
        y_position += line_spacing

    return label_img

if __name__ == "__main__":
    main()