import gradio as gr
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Pre-existing emails
pre_existing_emails = [
    "Build system failed due to missing dependencies.",
    "Agent deployment was successful on the new cluster.",
    "Cluster performance metrics are above expectations.",
    "General inquiry about the project timeline.",
    "Build logs indicate a configuration error.",
    "Agent is not responding after the last update.",
    "Cluster nodes need to be restarted.",
    "General feedback on the new UI.",
    "Build completed successfully.",
    "Agent registration failed due to timeout."
]

# Fit TF-IDF vectorizer on pre-existing emails
vectorizer = TfidfVectorizer().fit(pre_existing_emails)
email_vectors = vectorizer.transform(pre_existing_emails)

def semantic_search(query: str) -> dict:
    query_vec = vectorizer.transform([query])
    similarities = cosine_similarity(query_vec, email_vectors).flatten()
    top_indices = similarities.argsort()[-5:][::-1]
    results = [
        {
            "email": pre_existing_emails[i],
            "similarity": float(similarities[i])
        }
        for i in top_indices
    ]
    return {"results": results}

def keyword_search(query: str) -> dict:
    matches = [
        {
            "email": email,
            "match": query.lower() in email.lower()
        }
        for email in pre_existing_emails if query.lower() in email.lower()
    ]
    return {"results": matches[:5]}

with gr.Blocks() as demo:
    gr.Markdown("# Email Tools")
    gr.Markdown("Semantic and keyword search tools for pre-existing emails.")
    
    with gr.Tab("Semantic Search"):
        gr.Interface(fn=semantic_search,
                     api_name="semantic_search",
                     description="Search through pre-existing emails using semantic search.",
                     inputs=gr.Textbox(placeholder="Enter your email or query..."),
                     outputs=gr.JSON())
    
    with gr.Tab("Keyword Search"):
        gr.Interface(fn=keyword_search,
                     api_name="keyword_search",
                     description="Search through pre-existing emails using keyword search.",
                     inputs=gr.Textbox(placeholder="Enter keyword..."),
                     outputs=gr.JSON())

if __name__ == "__main__":
    # demo.launch()
    demo.launch(mcp_server=True)