import streamlit as st
from langchain.prompts import PromptTemplate
import requests
from langchain.llms import HuggingFaceHub
from langchain.chains import LLMChain
import io
import os
from PIL import Image
import json
from model import model,tokenizer

API = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
# 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)

# Save image to a file
def save_image(image, image_path):
    image.save(image_path)



# content generation
def generate_content(topic):
  keyword=topic
  prompt = [{'role': 'user', 'content': f'''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 2000 words long. Use SEO best practices for content optimization.
          Add ## before section headers
  '''}]
  inputs = tokenizer.apply_chat_template(
      prompt,
      add_generation_prompt=True,
      return_tensors='pt'
  )

  tokens = model.generate(
      inputs.to(model.device),
      max_new_tokens=10024,
      temperature=0.8,
      do_sample=True
  )

  content = tokenizer.decode(tokens[0], skip_special_tokens=False)
  # print(content)
  return content

def divide_content(text):
    sections = {}
    lines = text.split('\n')

    current_section = None

    for line in lines:
        line = line.strip()  # Remove leading and trailing whitespaces
        if line.startswith("##"):
            # Found a new section marker
            current_section = line[2:]
            sections[current_section] = ""
        elif current_section is not None and line:
            # Append the line to the current section if it's not empty
            sections[current_section] += line + " "

    # Remove trailing whitespaces from each section
    for section_name, section_content in sections.items():
        sections[section_name] = section_content.rstrip()

    return sections


# Image Generation
API_URL = "https://api-inference.huggingface.co/models/goofyai/3d_render_style_xl"
headers = {"Authorization": "Bearer API"}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.content

def generat_image(image_prompt,name):
      image_bytes = query({
          "inputs": image_prompt,
      })
      image = Image.open(io.BytesIO(image_bytes))
      image.save(f"{name}.png")
      return image

def display_content_with_images(blog):
    blog_images = [key for key in list(blog.keys()) if "_image" in key]
    # Streamlit Display
    st.header(blog['title'])
    i = 0
    # Introduction
    col1, col2 = st.columns(2, gap='medium')
    with col1:
      st.header('Introduction')
      st.write(blog['Introduction'])
    with col2:
      st.image(blog[blog_images[i]], use_column_width=True)
      i+=1

    # History
    st.header('History and Background')
    st.write(blog['History and Background'])
    st.image(blog[blog_images[i]], use_column_width=True)
    i+=1
    # Content
    col1, col2 = st.columns(2, gap='medium')
    with col1:
      st.header('Key Concepts and Terminology')
      st.write(blog['Key Concepts and Terminology'])
    with col2:
      st.image(blog[blog_images[i]], use_column_width=True)
      i+=1
    # Use Cases and Applications
    st.header('Use Cases and Applications')
    st.write(blog['Use Cases and Applications'])

    # Benefits and Drawbacks
    st.header('Benefits and Drawbacks')
    st.write(blog['Benefits and Drawbacks'])

    # Future Outlook
    st.header('Future Outlook')
    st.write(blog['Future Outlook'])

    # Conclusion
    col1, col2 = st.columns(2, gap='medium')
    with col1:
      st.header('Conclusion')
      st.write(blog['Conclusion'])
    with col2:
      st.image(blog[blog_images[i]], use_column_width=True)
      i+=1


# Streamlit App

# Title
st.sidebar.title('📝 Previous Ideas')
st.title("AI Blog Content Generator 😊")
# Main Page
col1, col2, col3 = st.columns((1, 3, 1), gap='large')


existing_ideas = load_ideas()

# Input and button
topic = st.text_input("Enter Title for the blog")
button_clicked = st.button("Create blog!❤️")


# 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(topic)
    display_content_with_images(selected_idea_from_list[selected_idea])
else:
    index = 0

# Check if the topic exists in previous ideas before generating
if button_clicked and topic not in keys:
    st.write('Generating blog post about', topic, '...')
    st.write('This may take a few minutes.')

    topic_query = topic

    content = generate_content(topic)
    # st.write(content)
    blog = divide_content(content)
    st.write(blog)
    st.header(topic)
    keyss = list(blog.keys())
    image_prompts = []
    i=0
    while len(image_prompts)<4:
      try:
        image_prompts.append((keyss[i],blog[keyss[i]].splitlines()[0]))
        i+=1
      except Exception as e:
        print(e)
        i+=1

    # Blog Data
    blog_data = {
        'title': topic,
        'Introduction': blog[' Introduction'],
        'History and Background': blog[' History and Background'],
        'Key Concepts and Terminology': blog[' Key Concepts and Terminology'],
        'Use Cases and Applications': blog[' Use Cases and Applications'],
        'Benefits and Drawbacks': blog[' Benefits and Drawbacks'],
        'Future Outlook': blog[' Future Outlook'],
        'Conclusion': blog[' Conclusion'],
    }
    for k,image in image_prompts:
      img = generat_image(image,f" {k}{topic}")
      blog_data[f'{k}_image'] = f" {k}{topic}.png"
    
    display_content_with_images(blog_data)

    # Save blog with images
    existing_ideas.append({topic: blog_data})
    # 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))

    save_ideas(existing_ideas)