import gradio as gr
import openai
import os
import docx
from docx.shared import Pt

# OpenRouter API Key
OPENROUTER_API_KEY = "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"
OPENROUTER_MODEL = "sophosympatheia/rogue-rose-103b-v0.2:free"

# Initialize OpenAI client with OpenRouter base URL
print(f"Using API Key: {OPENROUTER_API_KEY}")
openai_client = openai.OpenAI(
    api_key=OPENROUTER_API_KEY,
    base_url="https://openrouter.ai/api/v1"  # OpenRouter API endpoint
)

# Few-shot examples for test case generation
FEW_SHOT_EXAMPLES = """
Example 1:
Project: E-commerce Website
Requirement: Users should be able to add items to their cart.
Test Case:
1. Navigate to the product page.
2. Click the "Add to Cart" button.
3. Verify that the item appears in the cart.

Example 2:
Project: Banking App
Requirement: Users should be able to transfer funds between accounts.
Test Case:
1. Log in to the banking app.
2. Navigate to the "Transfer Funds" section.
3. Select the source and destination accounts.
4. Enter the amount and confirm the transfer.
5. Verify that the transaction is reflected in both accounts.
"""

# Function to generate test cases using OpenRouter API
def generate_test_cases(project_details):
    print("Generating test cases...")
    try:
        response = openai_client.chat.completions.create(
            model=OPENROUTER_MODEL,
            messages=[
                {
                    "role": "system",
                    "content": "You are a software testing expert. Your task is to generate test cases based on the provided project details. "
                               "Use the following examples as a guide:\n" + FEW_SHOT_EXAMPLES
                },
                {
                    "role": "user",
                    "content": f"Project Details:\n{project_details}\n\nGenerate test cases for the above project."
                }
            ]
        )
        print("Test cases generated.")
        return response.choices[0].message.content
    except Exception as e:
        print(f"Error generating test cases: {e}")
        return f"Error: {e}"

# Function to create a Word document
def create_word_document(content, filename="test_cases.docx"):
    print("Creating Word document...")
    doc = docx.Document()
    doc.add_heading("Generated Test Cases", level=1)
    
    # Add content to the document
    for line in content.split("\n"):
        if line.strip():
            doc.add_paragraph(line)
    
    # Save the document
    doc.save(filename)
    print(f"Word document saved as {filename}.")
    return filename

# Gradio UI for software testing automation
def gradio_ui():
    with gr.Blocks() as demo:
        gr.Markdown("## Software Testing Automation Tool")
        gr.Markdown("Upload project details or enter them as text to generate test cases and documentation.")
        
        # Inputs
        project_input = gr.Textbox(label="Enter Project Details", lines=5, placeholder="Describe the project requirements...")
        file_upload = gr.File(label="Or Upload Project Document (TXT, PDF, DOCX)", file_types=[".txt", ".pdf", ".docx"])
        
        # Outputs
        output_text = gr.Textbox(label="Generated Test Cases", lines=10, interactive=False)
        download_link = gr.File(label="Download Test Cases as Word Document")
        
        # Buttons
        generate_btn = gr.Button("Generate Test Cases")
        
        # Event handling
        generate_btn.click(
            fn=lambda project_details: generate_test_cases(project_details),
            inputs=[project_input],
            outputs=[output_text]
        )
        generate_btn.click(
            fn=lambda project_details: create_word_document(generate_test_cases(project_details)),
            inputs=[project_input],
            outputs=[download_link]
        )
    
    return demo

# Launch Gradio UI
demo = gradio_ui()
print("Launching Gradio UI...")
demo.launch()