import os
import gradio as gr
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import HuggingFaceHub
from langchain_core.documents import Document
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.tools.tavily_search.tool import TavilySearchResults
# Load environment variables
load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HF_TOKEN")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

# Set API keys
#os.environ["HUGGINGFACEHUB_API_TOKEN"] = HUGGINGFACEHUB_API_TOKEN
#os.environ["TAVILY_API_KEY"] = TAVILY_API_KEY

# Prompt Template
prompt_template = PromptTemplate(
    input_variables=["context", "user_story"],
    template="""You are a QA expert. Based on the context below and the given user story, write a detailed list of test cases.

Context:
{context}

User Story:
{user_story}

Test Cases:"""
)

# Load knowledge from RAG (local file)
def load_rag_knowledge():
    with open("rag_knowledge_base.txt", "r", encoding="utf-8") as file:
        content = file.read()
    docs = [Document(page_content=content)]
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_store = FAISS.from_documents(docs, embeddings)
    return vector_store.similarity_search("test case generation", k=1)

# Tavily search
def tavily_search(query):
    search = TavilySearchResults(k=1)
    results = search.run(query)
    return results[0]['content'] if results else "No relevant results from Tavily."

# LLM call with combined context
def call_llm_with_context(context, user_story):
    llm = HuggingFaceHub(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        model_kwargs={"temperature": 0.7, "max_new_tokens": 500}
    )
    chain = LLMChain(llm=llm, prompt=prompt_template)
    return chain.run({"context": context, "user_story": user_story})

# Generate test cases pipeline
def generate_test_cases(user_story):
    rag_docs = load_rag_knowledge()
    rag_text = "\n".join([doc.page_content for doc in rag_docs])
    tavily_text = tavily_search(user_story)
    full_context = f"{rag_text}\n\n{tavily_text}"
    test_cases = call_llm_with_context(full_context, user_story)
    return rag_text.strip(), tavily_text.strip(), test_cases.strip()

# Gradio handler
def handle_generate(user_story):
    rag, tavily, result = generate_test_cases(user_story)
    return rag, tavily, result

# ----------------- Gradio UI -----------------
with gr.Blocks() as demo:
    gr.Markdown("# 🧪 TechTales TestCaseGenerator using RAG + Tavily + Mistral + LangChain -  Developed by Pankaj Kumar")
    gr.Markdown("🚀 Enter a user story below to generate test cases using your knowledge base and Tavily search.")

    user_input = gr.Textbox(label="📝 Enter User Story", lines=4, placeholder="As a user, I want to...")

    btn = gr.Button("🔍 Generate Test Cases")

    rag_output = gr.Textbox(label="📚 Knowledge from RAG File", lines=8)
    tavily_output = gr.Textbox(label="🌐 Knowledge from Tavily Search", lines=8)
    result_output = gr.Textbox(label="✅ Final Test Cases", lines=12)

    btn.click(
        handle_generate,
        inputs=[user_input],
        outputs=[rag_output, tavily_output, result_output]
    )

demo.launch()