import gradio as gr
from transformers import pipeline
import json
import signal

# Load a small, free, instruction-following model
generator = pipeline("text2text-generation", model="google/flan-t5-large")

# Timeout handling
class TimeoutException(Exception):
    pass

def timeout_handler(signum, frame):
    raise TimeoutException("Processing took too long. Try a simpler input.")

signal.signal(signal.SIGALRM, timeout_handler)

def generate_test_cases(user_story):
    try:
        signal.alarm(180)  # Set a 3-minute timeout

        # Structured prompt for better response
        prompt = (
            f"Generate structured test cases for the following user story:\n"
            f"User Story: {user_story}\n"
            f"Provide output in a clear, structured way with a test case title, steps, and expected result."
        )

        output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]

        # Simple manual post-processing to force JSON format
        test_cases = []
        cases = output.split("\n\n")  # Split into test cases

        for i, case in enumerate(cases, start=1):
            lines = case.split("\n")
            if len(lines) >= 3:
                title = lines[0].strip()
                steps = [line.strip() for line in lines[1:-1]]
                expected_result = lines[-1].strip()

                test_cases.append({
                    "id": i,
                    "title": title,
                    "steps": steps,
                    "expected_result": expected_result
                })

        if not test_cases:
            return "Error: Model did not return structured test cases. Try again."

        formatted_output = json.dumps({"test_cases": test_cases}, indent=4)

        signal.alarm(0)  # Disable timeout if successful
        return formatted_output

    except TimeoutException:
        return "Processing timed out. Please try again with a simpler input."

# Gradio UI
iface = gr.Interface(
    fn=generate_test_cases,
    inputs=gr.Textbox(lines=5, placeholder="Enter your user story here..."),
    outputs="text",
    title="AI Test Case Generator",
    description="Enter a user story and get structured test cases in JSON format.",
)

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