import pandas as pd
import requests
import gradio as gr
from sentence_transformers import SentenceTransformer, util
from sklearn.linear_model import LogisticRegression
import os

# Load semantic similarity model
model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")

# Initialize classifier
classifier = LogisticRegression()
feedback_file = "feedback.csv"

# Load feedback if available
if os.path.exists(feedback_file):
    feedback_data = pd.read_csv(feedback_file)
    if len(feedback_data) > 10:
        X = feedback_data[["similarity", "resp_length"]]
        y = feedback_data["label"]
        classifier.fit(X, y)
else:
    feedback_data = pd.DataFrame(columns=["similarity","resp_length","label"])

def run_tests(file, bearer_token=None):
    df = pd.read_excel(file)
    results = []
    headers = {}
    if bearer_token:
        headers["Authorization"] = f"Bearer {bearer_token}"
    
    for _, row in df.iterrows():
        test_name = row["TestName"]
        endpoint = row["Endpoint"]
        method = row["Method"]
        payload = row.get("Payload", None)
        expected_status = row["ExpectedStatus"]
        expected_body = str(row["ExpectedBody"])
        
        try:
            response = requests.request(method, endpoint, headers=headers, data=payload)
            status_match = response.status_code == expected_status
            body = response.text
        except Exception as e:
            results.append((test_name, "Error", str(e)))
            continue
        
        if not status_match:
            result = "Fail"
        elif body == expected_body:
            result = "Pass"
        else:
            sim = util.cos_sim(model.encode(expected_body), model.encode(body)).item()
            resp_length = len(body)
            
            # If classifier trained, use it
            if len(feedback_data) > 10:
                pred = classifier.predict([[sim, resp_length]])[0]
                result = "Pass" if pred == 1 else "Fail"
            else:
                result = "Pass" if sim > 0.75 else "Uncertain"
            
            # Store similarity for feedback loop
            results.append((test_name, result, sim))
            continue
        
        results.append((test_name, result, None))
    
    return pd.DataFrame(results, columns=["TestName", "Result", "Similarity"])

def feedback_handler(similarity, resp_length, label):
    global feedback_data
    feedback_data = feedback_data.append({"similarity": similarity,
                                          "resp_length": resp_length,
                                          "label": label}, ignore_index=True)
    feedback_data.to_csv(feedback_file, index=False)
    return "Feedback saved!"

demo = gr.Interface(
    fn=run_tests,
    inputs=[gr.File(label="Upload Excel"), gr.Textbox(label="Bearer Token")],
    outputs="dataframe",
    title="Adaptive API Test Runner",
    description="Upload test cases in Excel, run API tests, and classify results with semantic similarity + ML feedback loop."
)

demo.launch()