import os
import subprocess
import time
# import tempfile
from virtualenv_mgr import VirtualenvManager

import streamlit as st
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SequentialChain
from langchain.memory import ConversationBufferMemory

os.environ['OPENAI_API_KEY'] = st.secrets["openai_api_key"]

st.title("dAIgramGen : The AI-based Diagram Generator")

daig_name = st.text_input("Diagram name")
daig_name_lower = daig_name.lower()
daig_name_underscore = daig_name_lower.replace(" ", "_")
daig_python_filename = daig_name_underscore + ".py"
daig_png_filename = daig_name_underscore + ".png"

daig_prompt = st.text_area("Prompt for diagrams")

daig_template = PromptTemplate(
    input_variables = ['topic', 'name'],
    template = 'Using the Python diagrams library generate the code for {topic} and name the diagram {name}. The Diagram class in the generated code should have the following parameters set: output_format as jpg and show=False. Just return the code. No explaination of the code needed'
)

daig_memory = ConversationBufferMemory(input_key='topic', memory_key='daig_history')

llm = OpenAI(temperature=0.9)
daig_chain = LLMChain(llm=llm, prompt=daig_template, verbose=True, output_key='code', memory=daig_memory)

if daig_prompt and daig_name:
	code = daig_chain.run({'topic':daig_prompt, 'name': daig_name})
	st.subheader('Generated Code')
	st.code(code)
	venv = VirtualenvManager()
	venv.create_env()
	venv.add_dependency("diagrams")
	
	if os.path.exists(daig_python_filename):
		os.remove(daig_python_filename)
	f = open(daig_python_filename, "w")
	f.write(code)
	f.flush()
	process = venv.run_code(f.name)

	time.sleep(5)
	if os.path.exists(daig_png_filename):
		st.image(daig_png_filename)
    
	with st.expander('Code History'):
		st.info(daig_memory.buffer)