import gradio as gr
import openai  # Can be replaced with another AI model API
import os
import subprocess

# AI Model Selection
MODEL_MAP = {
    "gpt": "gpt-4",
    "gemini": "gemini-1",
    "custom": "http://custom-llm-api"
}

def generate_test(model, prompt):
    """Generate test script using the selected AI model."""
    model_api = MODEL_MAP.get(model, "gpt-4")
    try:
        response = openai.ChatCompletion.create(
            model=model_api,
            messages=[{"role": "system", "content": prompt}]
        )
        test_script = response["choices"][0]["message"]["content"]

        # Save the generated test script
        with open("generated_test.py", "w") as f:
            f.write(test_script)

        return test_script
    except Exception as e:
        return f"Error: {str(e)}"

def run_tests():
    """Execute AI-generated test scripts and return results."""
    try:
        result = subprocess.run(["pytest", "generated_test.py", "--json-report"], capture_output=True, text=True)
        return result.stdout + "\n" + result.stderr
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# AI Test Framework")
    
    with gr.Row():
        model = gr.Dropdown(["gpt", "gemini", "custom"], value="gpt", label="Choose AI Model")
        prompt = gr.Textbox(label="Test Prompt", placeholder="Describe the test case...")

    generate_button = gr.Button("Generate Test Script")
    test_script_output = gr.Textbox(label="Generated Test Script", interactive=False)

    generate_button.click(generate_test, inputs=[model, prompt], outputs=test_script_output)

    run_tests_button = gr.Button("Run Tests")
    test_results_output = gr.Textbox(label="Test Execution Results", interactive=False)

    run_tests_button.click(run_tests, inputs=[], outputs=test_results_output)

app.launch()