import os
import tempfile
import requests
import streamlit as st
from src.pipeline.main import LearnableAI
from src.services.use_case_three.story_generator import StoryGenerator
from src.services.image_generation.image_gen import image_generation


def main():
    st.title("LearnableAI")

    learnable_ai = LearnableAI()

    st.sidebar.header("Configuration")
    use_case = st.sidebar.radio(
        "Select Use Case",
        ["Use Case: 1(Word-Sentences)", "Use Case: 2(Image Describer)", "Use Case: 3(Story Generator From Image)"]
    )

    difficulty = st.sidebar.selectbox(
        "Select Difficulty Level",
        ["Simple", "Easy", "Challenging", "Advanced"]
    )

    if use_case == "Use Case: 1(Word-Sentences)":
        st.header("Word to Sentence Generation")
        st.write(
            "Enter a list of words and the child interest to generate the content for eg.Words: Apple, Banana interest: Ironman")

        words = st.text_area(
            "Words:",
            "",
            help="Enter multiple words separated by commas"
        )

        interest = st.text_input("Child Interest:", "",
                                 help="Enter the interest area for the content ie. Superheroes, Animals, etc.")

        if st.button("Generate Content", key="text_process"):
            if words and interest:
                try:
                    word_list = [w.strip() for w in words.split(",")] if "," in words else [words.strip()]
                    result = learnable_ai.word_to_sentence(
                        words=word_list,
                        interest=interest,
                        difficulties=difficulty
                    )

                    st.success("Content Generated Successfully!")
                    st.write("### Results")
                    st.write(result.split("</think>")[1].split("Image Generation Prompt")[0])
                    image_prompt = result.split("Image Generation Prompt")[1]
                    input = {
                        "prompt": image_prompt,
                        "go_fast": True,
                        "megapixels": "1",
                        "num_outputs": 1,
                        "aspect_ratio": "1:1",
                        "output_format": "webp",
                        "output_quality": 80,
                        "num_inference_steps": 4
                    }
                    output_image_url = image_generation(input)
                    image_bytes = requests.get(output_image_url).content
                    st.write("### Generated Image")
                    st.image(image_bytes, caption="Generated Image")


                except Exception as e:
                    st.error(f"An error occurred: {str(e)}")
            else:
                st.warning("Please fill in both words and interest fields!")

    elif use_case == "Use Case: 2(Image Describer)":
        st.header("Image Description")
        st.write("Upload an image or capture from camera to generate to generate the description of the image")

        upload_method = st.radio(
            "Choose Upload Method",
            ["Upload File", "Capture from Camera"]
        )

        if upload_method == "Upload File":
            image_file = st.file_uploader(
                "Upload an image",
                type=["jpg", "jpeg", "png"]
            )
        else:
            image_file = st.camera_input("Take a picture")

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

            if st.button("Generate Description", key="image_process"):
                try:  # temp image file
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
                        temp_file.write(image_file.read())
                        temp_image_path = temp_file.name

                    result = learnable_ai.image_describer(
                        image_path=temp_image_path,
                        difficulties=difficulty
                    )

                    st.success("Image Processed Successfully!")
                    st.write("### Image Description Results")
                    st.write(result)

                    os.remove(temp_image_path)  # remove temp image file

                except Exception as e:
                    st.error(f"An error occurred while processing the image: {str(e)}")


    elif use_case == "Use Case: 3(Story Generator From Image)":
        st.title("Story Generator from Images")
        st.write("Story writer from the images. Upload up to 3 images to generate a children's story.")
        st.subheader("Upload up to 3 images and generate a children's story")

        uploaded_files = st.file_uploader(
            "Upload up to 3 images",
            type=["jpg", "jpeg", "png"],
            accept_multiple_files=True
        )

        if uploaded_files:
            if len(uploaded_files) > 3:
                st.error("You can upload a maximum of 3 images.")
            else:
                st.write("Uploaded Images:")
                for uploaded_file in uploaded_files:
                    st.image(uploaded_file, caption=uploaded_file.name, use_column_width=True)

                if st.button("Generate Story"):
                    story_generator = StoryGenerator()
                    story_parts = []

                    try:
                        for uploaded_file in uploaded_files:
                            with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
                                temp_file.write(uploaded_file.read())
                                temp_image_path = temp_file.name

                            story_part = story_generator.process_image(temp_image_path)
                            story_parts.append(story_part)

                            os.remove(temp_image_path)

                        full_story = "\n\n".join(story_parts)
                        st.success("Story Generated Successfully!")
                        st.write("### Your Story")
                        st.write(full_story)

                    except Exception as e:
                        st.error(f"An error occurred: {str(e)}")


if __name__ == "__main__":
    main()