import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import io
import cv2
import numpy as np
import easyocr
import os
from dotenv import load_dotenv
import requests
import logging

# Load environment variables
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Hugging Face API setup
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
HF_TOKEN = os.getenv("HF_TOKEN")
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

def load_image(image_file):
    img = Image.open(image_file)
    return img

def detect_text(image):
    reader = easyocr.Reader(['en'])
    img_array = np.array(image)
    results = reader.readtext(img_array)
    return [(text, box) for (box, text, _) in results]

def replace_text_in_image(image, text_to_replace, new_text):
    img_array = np.array(image)
    for (text, box) in detect_text(image):
        if text == text_to_replace:
            x, y, w, h = int(box[0][0]), int(box[0][1]), int(box[2][0] - box[0][0]), int(box[2][1] - box[0][1])
            mask = np.zeros(img_array.shape[:2], dtype=np.uint8)
            cv2.rectangle(mask, (x, y), (x+w, y+h), 255, -1)
            img_array = cv2.inpaint(img_array, mask, 3, cv2.INPAINT_TELEA)
            image = Image.fromarray(img_array)
            draw = ImageDraw.Draw(image)
            font = ImageFont.truetype("arial.ttf", 40)
            draw.text((x, y), new_text, font=font, fill="#000000")
            return image
    return Image.fromarray(img_array)

def query(payload):
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        
        logger.debug(f"API response status code: {response.status_code}")
        logger.debug(f"API response headers: {response.headers}")
        
        content_type = response.headers.get('Content-Type', '')
        if 'application/json' in content_type:
            return response.json()
        elif 'image' in content_type:
            return response.content
        else:
            logger.error(f"Unexpected content type: {content_type}")
            st.error(f"Unexpected content type: {content_type}")
            return None
    except requests.exceptions.RequestException as e:
        logger.error(f"Request failed: {str(e)}")
        st.error(f"Request failed: {str(e)}")
        return None

def increase_image_quality(image, scale_factor):
    width, height = image.size
    new_size = (width * scale_factor, height * scale_factor)
    return image.resize(new_size, Image.LANCZOS)

def image_text_replacer():
    st.header("Image Text Replacer")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = load_image(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        text_results = detect_text(image)
        
        st.subheader("Detected Text:")
        
        edited_image = image.copy()
        
        for i, (text, box) in enumerate(text_results):
            if text.strip():  # Only process non-empty text
                st.text(f"{i+1}. {text}")
                
                new_text = st.text_input(f"Enter new text for '{text}':", value=text, key=f"new_text_{i}")
                
                if st.button(f"Replace '{text}'", key=f"replace_{i}"):
                    edited_image = replace_text_in_image(edited_image, text, new_text)
                    st.image(edited_image, caption='Edited Image', use_column_width=True)
        
        # Provide download option for the edited image
        buf = io.BytesIO()
        edited_image.save(buf, format="PNG")
        byte_im = buf.getvalue()
        st.download_button(
            label="Download edited image",
            data=byte_im,
            file_name="edited_image.png",
            mime="image/png"
        )

def poster_generator():
    st.header("Generate Poster")
    poster_type = st.selectbox("Poster Type", ["Fashion", "Movie", "Event", "Advertisement", "Other"])
    prompt = st.text_area("Prompt")
    num_images = st.number_input("Number of Images", min_value=1, max_value=5, value=1)
    quality_factor = st.number_input("Quality Factor", min_value=1, max_value=4, value=1)

    if st.button("Generate Images"):
        if poster_type == "Other":
            full_prompt = f"A colorful poster with the following elements: {prompt}"
        else:
            full_prompt = f"A colorful {poster_type.lower()} poster with the following elements: {prompt}"

        generated_images = []
        for i in range(num_images):
            with st.spinner(f"Generating image {i+1}..."):
                logger.info(f"Generating image {i+1} with prompt: {full_prompt}")
                response = query({"inputs": full_prompt})
                
                if isinstance(response, bytes):
                    image = Image.open(io.BytesIO(response))
                    if quality_factor > 1:
                        image = increase_image_quality(image, quality_factor)
                    generated_images.append(image)
                else:
                    st.error("Failed to generate image")

        # Display generated images
        for i, img in enumerate(generated_images):
            st.image(img, caption=f"Generated Poster {i+1}", use_column_width=True)
            
            # Provide download option for the generated image
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            byte_im = buf.getvalue()
            st.download_button(
                label=f"Download generated poster {i+1}",
                data=byte_im,
                file_name=f"generated_poster_{i+1}.png",
                mime="image/png"
            )

def main():
    st.title("Image Processing App")

    app_mode = st.sidebar.selectbox("Choose the app mode",
        ["Image Text Replacer", "Poster Generator"])

    if app_mode == "Image Text Replacer":
        image_text_replacer()
    elif app_mode == "Poster Generator":
        poster_generator()

if __name__ == "__main__":
    main()