# main.py
import streamlit as st
import transformers
import langchain
import agents
from streamlit.script_runner import StopException

# Define function to reverse prompt engineer code
def reverse_prompt_engineer(code):
    # Use natural language processing to analyze code
    nlp_analysis = langchain.analyze(code)

    # Choose the best free pretrained model for this task
    model_name = "microsoft/CodeGPT-small-py-adaptedGPT2"
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
    model = transformers.AutoModelForCausalLM.from_pretrained(model_name)

    # Generate perfect prompt using analyzed code
    perfect_prompt = agents.generate_prompt(nlp_analysis)

    # Chat with user to make additional changes to prompt
    chatbot = agents.ChatGPT(model=model, tokenizer=tokenizer)
    final_prompt = chatbot.chat(perfect_prompt)

    # Use final prompt to generate similar code using ChatGPT
    generated_code = chatbot.generate_code(final_prompt)

    return generated_code

# Streamlit UI
st.set_page_config(page_title="Code Generator", layout="wide", initial_sidebar_state="expanded")
st.title("Code Generator")

st.sidebar.title("Input")
code_input = st.sidebar.text_area("Enter your code here:", '''
def greet(name):
    print("Hello, " + name + ". How are you doing today?")

greet("John")
''')

if st.sidebar.button("Generate Code"):
    if code_input.strip() == "":
        st.error("Please enter some code in the input field.")
    else:
        try:
            generated_code = reverse_prompt_engineer(code_input)
            st.code(generated_code)
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            raise StopException