import transformers
import streamlit as st
import requests
import json
import sqlite3
import os
from transformers import pipeline

# Load AI model for test case generation
test_case_generator = pipeline("text-generation", model="microsoft/CodeGPT-small-py")

DB_PATH = "results.db"

# Agar database file exist nahi karti toh create karo
if not os.path.exists(DB_PATH):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS results (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            api_name TEXT,
            test_status TEXT,
            run_timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()
    

def fetch_swagger_data(swagger_url):
    try:
        response = requests.get(swagger_url)
        return response.json()
    except Exception as e:
        st.error(f"Error fetching Swagger data: {e}")
        return None

def generate_bdd_test_cases(swagger_data):
    test_cases = {}
    for path, methods in swagger_data.get("paths", {}).items():
        for method, details in methods.items():
            prompt = f"Generate BDD test case for {method.upper()} {path} with {details.get('parameters', [])}"
            ai_response = test_case_generator(prompt, max_length=200, num_return_sequences=1)
            test_cases[f"{method.upper()} {path}"] = ai_response[0]['generated_text']
    return test_cases

def execute_test_script(test_case):
    # Simulate API call execution
    st.write(f"Executing test: {test_case}")
    return {"status": "Passed", "response_time": "200ms"}

def store_test_results(test_case, result):
    conn = sqlite3.connect("test_results.db")
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS results (test_case TEXT, status TEXT, response_time TEXT)
    """)
    cursor.execute("INSERT INTO results (test_case, status, response_time) VALUES (?, ?, ?)", 
                   (test_case, result["status"], result["response_time"]))
    conn.commit()
    conn.close()

# Function to get past results
def show_past_results():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM results")
    data = cursor.fetchall()
    conn.close()
    return data

# Streamlit UI
st.title("AI-Powered API Test Case Generator")
swagger_url = st.text_input("Enter Swagger URL")
if st.button("Generate Test Cases"):
    swagger_data = fetch_swagger_data(swagger_url)
    if swagger_data:
        test_cases = generate_bdd_test_cases(swagger_data)
        st.session_state["test_cases"] = test_cases
        st.success("Test Cases Generated!")
        st.json(test_cases)

if "test_cases" in st.session_state:
    st.subheader("Execute Test Cases")
    for api, test_case in st.session_state["test_cases"].items():
        if st.button(f"Run {api}"):
            result = execute_test_script(test_case)
            store_test_results(api, result)
            st.success(f"{api} - {result['status']}")

st.subheader("Past Execution Results")
past_results = show_past_results()
st.table(past_results)