from diffusers import StableDiffusionPipeline
import torch
from langchain.chains import LLMChain
from langchain.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate
import requests
import base64
import streamlit as st
import json

# Load existing ideas from a file
def load_ideas():
    try:
        with open("ideas.json", "r") as file:
            ideas = json.load(file)
    except FileNotFoundError:
        ideas = []
    return ideas

# Save ideas to a file
def save_ideas(ideas):
    with open("ideas.json", "w") as file:
        json.dump(ideas, file)

# Function to generate content
@torch.no_grad()
def generate_content(topic):
    hub_llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta")
    prompt = PromptTemplate(
        input_variables=['keyword'],
        template="""
        Write a comprehensive article about {keyword} covering the following aspects:
        Introduction, History and Background, Key Concepts and Terminology, Use Cases and Applications, Benefits and Drawbacks, Future Outlook, Conclusion
        Ensure that the article is well-structured, informative, and at least 1500 words long. Use SEO best practices for content optimization.
        """
    )
    hub_chain = LLMChain(prompt=prompt, llm=hub_llm, verbose=True)
    content = hub_chain.run(topic)

    subheadings = [
        "Introduction",
        "History and Background",
        "Key Concepts and Terminology",
        "Use Cases and Applications",
        "Benefits and Drawbacks",
        "Future Outlook",
        "Conclusion",
    ]

    for subheading in subheadings:
        if (subheading + ":") in content:
            content = content.replace(subheading + ":", "## " + subheading + "\n")
        elif subheading in content:
            content = content.replace(subheading, "## " + subheading + "\n")

    return content



# generate image

import io
from PIL import Image
# Function to generate an image using the pre-created or newly created pipeline
@torch.no_grad()
def generate_image(topic):
  API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
  headers = {"Authorization": "Bearer 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"}

  def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content
  image_bytes = query({
    "inputs": f"A blog banner about {topic}",
  })
  # You can access the image with PIL.Image for example

  image = Image.open(io.BytesIO(image_bytes))
  image.save(f"{topic}.png")
  return image

# Streamlit app
st.title("Blog Generator")

# Input and button
topic = st.text_input("Enter Title for the blog")
button_clicked = st.button("Create blog!")
# Load existing ideas
existing_ideas = load_ideas()
st.sidebar.header("Previous Ideas:")

# Display existing ideas in the sidebar
keys = list(set([key for idea in existing_ideas for key in idea.keys()]))
if topic in keys:
    index = keys.index(topic)
    selected_idea = st.sidebar.selectbox("Select Idea", keys, key=f"selectbox{topic}", index=index)
    # Display content and image for the selected idea
    selected_idea_from_list = next((idea for idea in existing_ideas if selected_idea in idea), None)
    st.subheader(selected_idea)
    st.image(selected_idea_from_list[selected_idea]["image_path"])
    st.markdown(selected_idea_from_list[selected_idea]["content"])
else:
    index = 0
# Handle button click
if button_clicked:
    # Generate content and update existing ideas
    content, image = generate_content(topic),generate_image(topic)
    if image:
      image_path = f"{topic}.png"
    existing_ideas.append({topic: {"content": content, "image_path": image_path}})
    save_ideas(existing_ideas)
    # Update keys and selected idea in the sidebar
    keys = list(set([key for idea in existing_ideas for key in idea.keys()]))
    selected_idea = st.sidebar.selectbox("Select Idea", keys, key=f"selectbox{topic}", index=keys.index(topic))
    st.image(image)
    st.markdown(content)