import os
import streamlit as st

from dotenv import dotenv_values

from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA, LLMChain
from langchain.document_loaders import TextLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFaceEndpoint
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma

env_vars = dotenv_values(".env")
LOCAL=False
if not LOCAL:
    env_vars = {
        "HUGGINGFACE_API_TOKEN": st.secrets["HUGGINGFACE_API_TOKEN"],
        "HUGGINGFACE_MODEL": st.secrets["HUGGINGFACE_MODEL"]
    }

# prepare Falcon Huggingface API
llm = HuggingFaceEndpoint(
            endpoint_url= f"https://api-inference.huggingface.co/models/{env_vars['HUGGINGFACE_MODEL']}",
            huggingfacehub_api_token=env_vars["HUGGINGFACE_API_TOKEN"],
            task="text-generation",
            model_kwargs = {
                "min_length": 8192,
                "max_length":8192,
                "temperature":0.9,
                "max_new_tokens":4000,
                "num_return_sequences":1
            }
        )


keyword_prompt_template = "What are 10 important keywords related too: {word}? Only return a list of words and do not include any duplicates."
keyword_chain = LLMChain(
    llm=llm,
    prompt=PromptTemplate.from_template(keyword_prompt_template),
)

outline_prompt_template = "Create an outline for an article about: {word}."
outline_chain = LLMChain(
    llm=llm,
    prompt=PromptTemplate.from_template(outline_prompt_template),
)

article_prompt_template = """
Act as an expert SEO Writer.
Write a well crafted article using the given outline as a guide.
Use the relavant keywords you are provided.
Apply EEAT principles and SEO best practices.
It is important that the content is at least 1500 words.
Be sure to include section headers.
OUTLINE: {outline}
KEYWORDS: {keywords}
"""
article_chain = LLMChain(
    llm=llm,
    prompt=PromptTemplate.from_template(article_prompt_template),
)   

st.title("Blog Writer")
keyword = st.text_input("Input the keyword you wish to write about")

if keyword:
    with st.spinner("Writing..."):
        st.write(f"Writing about: {keyword}")
        with st.expander("Keywords"):
            keywords = keyword_chain.run(keyword)
            st.write(keywords)
        with st.expander("Outline"):
            outline = outline_chain.run(keyword)
            st.write(outline)
        with st.expander("Article"):
            article = article_chain.run(outline=outline, keywords=keywords)
            st.write(article)