import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain


# Inicializa o pipeline do Hugging Face
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
llm = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Define os templates de prompt
prompt_template_br = (
    "Você é um analista de teste, responsável pela criação de testes unitários, com base na história de usuário {us} "
    "e nos critérios de aceites {ca} baseados no gherking, crie os testes unitários para a linguagem de programação {lp} "
    "baseados no framework {fw} "
    "e explique como usar o código"
)

prompt_template = (
    "You are a test analyst, responsible for creating unit tests based on the user story {us} "
    "and the acceptance criteria {ca} based on Gherkin. Create the unit tests for the programming language {lp} "
    "based on the {fw} framework and explain how to use the code."
)

# Interface do Streamlit
st.title("AutoDevSuite : BDD-based Test Generator")
us = st.text_area("Enter US:")
ca = st.text_area("Enter your Gherkin-based acceptance criteria:")
lp = st.text_input("Enter the programming language:")
fw = st.text_input("Enter the framework:")

prompt = PromptTemplate(
    input_variables=["us", "ca", "lp", "fw"],
    template=prompt_template
)
chain_1 = LLMChain(llm=llm, prompt=prompt)

if st.button("Generate Unit Tests"):
    # Gera o resultado usando os valores de entrada
    inputs = {"us": us, "ca": ca, "lp": lp, "fw": fw}
    result = chain_1.run(inputs)
    
    # Exibe o resultado
    st.write("Results:")
    st.write(result)