import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import io
import base64
from datetime import datetime

def calculate_font_size(img_width, img_height, text):
    max_font_size = int(img_width * 0.1)
    min_font_size = int(img_width * 0.02)

    if len(text) < 5:
        return max_font_size
    elif len(text) < 20:
        return int((max_font_size + min_font_size) / 2)
    else:
        return min_font_size

def generate_meme(img_path, top_text, middle_text, bottom_text):
    img = Image.open(img_path)
    draw = ImageDraw.Draw(img)
    
    font_path = "Nasa21-l23X.ttf"
    font_size_top = calculate_font_size(img.width, img.height, top_text)
    font_size_middle = calculate_font_size(img.width, img.height, middle_text)
    font_size_bottom = calculate_font_size(img.width, img.height, bottom_text)

    font_top = ImageFont.truetype(font_path, font_size_top)
    font_middle = ImageFont.truetype(font_path, font_size_middle)
    font_bottom = ImageFont.truetype(font_path, font_size_bottom)

    draw.text((10,10), top_text, font=font_top, fill="white")
    draw.text((10,img.height // 2 - font_size_middle // 2), middle_text, font=font_middle, fill="white")
    draw.text((10,img.height - 10 - font_size_bottom), bottom_text, font=font_bottom, fill="white")

    return img

def get_image_download_link(img, filename="meme.png"):
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    href = f'<a href="data:image/png;base64,{img_str}" download="{filename}">Download Image</a>'
    return href

st.title("Meme Generator")

uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png"])

if uploaded_file is not None:
    st.image(uploaded_file, caption='Uploaded Image.', use_column_width=True)

    top_text = st.text_input("Enter top text")
    middle_text = st.text_input("Enter middle text")
    bottom_text = st.text_input("Enter bottom text")

    if st.button("Generate Meme"):
        result_img = generate_meme(uploaded_file, top_text, middle_text, bottom_text)

        filename = datetime.now().strftime('%Y%m%d%H%M%S') + ".png"
        st.image(result_img, caption='Generated Meme', use_column_width=True)

        st.markdown(get_image_download_link(result_img, filename), unsafe_allow_html=True)

# References
st.write("### References:")
st.write(f"- [DOI](https://doi.org/10.57967/hf/0428)")
#st.write(f"- [arXiv](https://arxiv.org/abs/1910.09700)")