import requests
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont
from transformers import BlipProcessor, BlipForConditionalGeneration
import streamlit as st
import torch
import os
from io import BytesIO
import barcode
from barcode.writer import ImageWriter

# Set device to CPU
device = torch.device("cpu")

# Load the BLIP model and processor
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to(device)

# Your Gemini API key
GEMINI_API_KEY = os.getenv('gemini_api')

def generate_detailed_description(image):
    inputs = processor(image, return_tensors="pt")
    out = model.generate(**inputs)
    description = processor.decode(out[0], skip_special_tokens=True)
    return description

def suggest_enhancements(image):
    suggestions = [
        "Consider adjusting the brightness for a more vivid image.",
        "Increase contrast to make the image details stand out more.",
        "Apply sharpening to enhance image details.",
        "Use a filter to create a specific mood or style."
    ]
    return suggestions

def enhance_description_with_gemini(description):
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key={GEMINI_API_KEY}"
    headers = {"Content-Type": "application/json"}
    data = {
        "contents": [
            {
                "parts": [
                    {
                        "text": (
                            f"Based on the following description of an image, provide a more detailed description, "
                            f"generate possible captions, and suggest logo ideas. \n\n"
                            f"Description: {description}\n\n"
                            f"Detailed Description:\n"
                            f"Captions:\n"
                            f"Logo Suggestions:\n"
                        )
                    }
                ]
            }
        ]
    }
    
    try:
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        response_json = response.json()
        if "candidates" in response_json and len(response_json["candidates"]) > 0:
            candidate = response_json["candidates"][0]
            if "content" in candidate and "parts" in candidate["content"]:
                parts = candidate["content"]["parts"]
                if len(parts) > 0:
                    return parts[0].get("text", "No response text found")
        return "No contents or parts in response"
    except requests.exceptions.RequestException as e:
        return f"Request failed: {e}"
def create_label(image, ingredients, usage, expiry_date, barcode_data):
    try:
        # Create an image for the label
        label_image = Image.new('RGB', (800, 1200), color='white')
        draw = ImageDraw.Draw(label_image)

        # Define font and size
        try:
            font = ImageFont.truetype("arial.ttf", 24)
        except IOError:
            font = ImageFont.load_default()
        
        # Draw the main image
        image.thumbnail((600, 400))
        label_image.paste(image, (100, 50))

        # Draw text
        text_position = (100, 500)
        draw.text(text_position, f"Ingredients:\n{ingredients}", font=font, fill="black")
        text_position = (100, 600)
        draw.text(text_position, f"Usage:\n{usage}", font=font, fill="black")
        text_position = (100, 700)
        draw.text(text_position, f"Expiry Date: {expiry_date}", font=font, fill="black")

        # Draw barcode
        if len(barcode_data) == 12:  # Check length for EAN-13
            barcode_image = generate_barcode(barcode_data)
            label_image.paste(barcode_image, (100, 800))
        else:
            draw.text((100, 800), "Invalid Barcode Data", font=font, fill="red")
        
        return label_image
    except Exception as e:
        raise RuntimeError(f"Error generating label: {e}")

# Streamlit interface continuation
if st.button('Generate Label'):
    if ingredients and usage and expiry_date and barcode_data:
        try:
            label_image = create_label(image, ingredients, usage, expiry_date, barcode_data)
            st.image(label_image, caption="Generated Label")
        except Exception as e:
            st.error(f"Error generating label: {e}")
    else:
        st.error("Please provide all required details for the label.")

def generate_barcode(data):
    code = barcode.get('ean13', data, writer=ImageWriter())
    barcode_image = code.render()
    return barcode_image

# Streamlit interface
st.title("Image Detailed Description, Captions, Logo Suggestions Generator with Image Enhancement Options")

# Custom CSS for the button style and alignment
st.markdown("""
    <style>
    .download-button-container {
        display: flex;
        justify-content: flex-end;
        margin-top: -45px; /* Adjust this value to align with the image */
    }
    .stDownloadButton button {
        background-color: #4CAF50; /* Green background */
        color: white; /* White text */
        padding: 10px 20px; /* Padding */
        border: none; /* No border */
        border-radius: 5px; /* Rounded corners */
        text-align: center; /* Centered text */
        text-decoration: none; /* No underline */
        display: inline-block; /* Make the link behave like a button */
        font-size: 16px; /* Increase font size */
        margin: 4px 2px; /* Margin */
        cursor: pointer; /* Pointer cursor on hover */
        transition-duration: 0.4s; /* Transition for hover effect */
    }
    .stDownloadButton button:hover {
        background-color: #45a049; /* Darker green on hover */
    }
    </style>
""", unsafe_allow_html=True)

# Session state to store the descriptions
if 'initial_description' not in st.session_state:
    st.session_state['initial_description'] = None
if 'enhanced_text' not in st.session_state:
    st.session_state['enhanced_text'] = None

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
img_url = st.text_input("Or enter image URL...")

if uploaded_file or img_url:
    if uploaded_file:
        image = Image.open(uploaded_file).convert('RGB')
    else:
        response = requests.get(img_url, stream=True)
        image = Image.open(response.raw).convert('RGB')

    st.image(image, caption="Uploaded Image", use_column_width=True)

    if st.session_state['initial_description'] is None:
        with st.spinner("Generating caption..."):
            initial_description = generate_detailed_description(image)
            enhanced_text = enhance_description_with_gemini(initial_description)
            st.session_state['initial_description'] = initial_description
            st.session_state['enhanced_text'] = enhanced_text
    else:
        initial_description = st.session_state['initial_description']
        enhanced_text = st.session_state['enhanced_text']

    st.subheader("Initial Description")
    st.write(initial_description)

    st.subheader("Enhanced Description, Captions, and Logo Suggestions with Image Enhancement Options")
    st.write(enhanced_text)

    # Function to generate and show enhancement suggestions
    def show_enhancements(image, effect_name):
        buffer = BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)

        st.markdown('<div class="download-button-container">', unsafe_allow_html=True)
        st.download_button(
            label=f"Download {effect_name} Image",
            data=buffer,
            file_name=f"{effect_name}_image.png",
            mime="image/png"
        )
        st.markdown('</div>', unsafe_allow_html=True)

    # Display enhancement options
    st.subheader("Enhancement Options")

    if st.button('Increase Brightness'):
        enhancer = ImageEnhance.Brightness(image)
        enhanced_image = enhancer.enhance(1.5)
        st.image(enhanced_image, caption="Enhanced Brightness")
        show_enhancements(enhanced_image, "Brightness")

    if st.button('Increase Contrast'):
        enhancer = ImageEnhance.Contrast(image)
        enhanced_image = enhancer.enhance(1.5)
        st.image(enhanced_image, caption="Enhanced Contrast")
        show_enhancements(enhanced_image, "Contrast")

    if st.button('Sharpen Image'):
        enhanced_image = image.filter(ImageFilter.SHARPEN)
        st.image(enhanced_image, caption="Sharpened Image")
        show_enhancements(enhanced_image, "Sharpen")

    if st.button('Apply Gaussian Blur'):
        enhanced_image = image.filter(ImageFilter.GaussianBlur(radius=2))
        st.image(enhanced_image, caption="Blurred Image")
        show_enhancements(enhanced_image, "Blur")

    if st.button('Apply Edge Enhancement'):
        enhanced_image = image.filter(ImageFilter.EDGE_ENHANCE)
        st.image(enhanced_image, caption="Edge Enhanced Image")
        show_enhancements(enhanced_image, "Edge Enhancement")

    # Inputs for label details
    st.subheader("Generate Label")
    ingredients = st.text_area("Ingredients")
    usage = st.text_area("Usage Instructions")
    expiry_date = st.text_input("Expiry Date (e.g., 2024-12-31)")
    barcode_data = st.text_input("Barcode Data")

    if st.button('Generate Label'):
        if ingredients and usage and expiry_date and barcode_data:
            try:
                label_image = create_label(image, ingredients, usage, expiry_date, barcode_data)
                st.image(label_image)
            except Exception as e:
                st.error(f"Error generating label: {e}")