from transformers import pipeline
import gradio as gr

# Load the model
model_name = "google/flan-t5-base"  # Choose your desired model
generator = pipeline("text2text-generation", model=model_name)

# Function to generate test cases
def generate_test_cases(method, url, headers, payload):
    prompt = f"""
    Generate comprehensive API test cases:
    Method: {method}
    URL: {url}
    Headers: {headers}
    Payload: {payload}

    Include:
    - Happy path
    - Negative test cases
    - Performance scenarios
    - Boundary conditions
    """
    try:
        response = generator(prompt, max_length=500)
        return response[0]['generated_text']
    except Exception as e:
        return f"Error generating test cases: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_test_cases,
    inputs=[
        gr.Textbox(label="HTTP Method (GET, POST, etc.)"),
        gr.Textbox(label="API URL"),
        gr.Textbox(label="Headers (JSON format)"),
        gr.Textbox(label="Payload (JSON format)"),
    ],
    outputs="text",
    title="API Test Case Generator",
    description="Generate structured test cases for APIs using Hugging Face models."
)

if __name__ == "__main__":
    iface.launch()