import os
import streamlit as st  # Front End / UI
from keyfile import secret_value
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SequentialChain

os.environ["OPENAI_API_KEY"] = secret_value

def main():
    st.title("Chain Based Story Generator")

    input1 = st.text_input("An inspirational object to start a story")
    input2 = st.text_input("A tone or a writer's style to shape the story")
    input3 = st.text_input("The final language for translating the text")

    if st.button("Generate"):
        result = process_input(input1, input2, input3)
        st.markdown(f"<div>{result}</div>", unsafe_allow_html = True)
        
def process_input(a, b, c):
    model = ChatOpenAI(model = "gpt-4o-mini")
    
    # Chain 1
    prompt_one = PromptTemplate(
        input_variables = ["object"],
        template = "A story about {object}",
    )

    chain_one = LLMChain(
        llm = model, 
        prompt = prompt_one,
        output_key = "story"
    )

    # Chain 2
    prompt_two = PromptTemplate(
        input_variables = ["style", "story"],
        template = "You are a professional writer. You will be given a short story and you will inspire from this original story and create a new story in {style} style. The original story: {story}.",
    )

    chain_two = LLMChain(
        llm = model, 
        prompt = prompt_two,
        output_key = "expanded_story"
        
    )

    # Chain 3
    prompt_three = PromptTemplate(
        input_variables = ["expanded_story", "language"],
        template = "Translate the story: {expanded_story} into {language}?",
    )

    chain_three = LLMChain(
        llm = model, 
        prompt = prompt_three,
        output_key = "translated"
        
    )

    # SEQUENCE
    final_chain = SequentialChain(
        chains = [chain_one, chain_two, chain_three],
        input_variables = ["object", "style", "language"],
        output_variables = ["translated"],
        verbose = True
    )

    resp = final_chain({
        "object": a,
        "style":b,
        "language": c
    })

    return resp["translated"]

if __name__ == "__main__":
    main()