import streamlit as st
from sentence_transformers import SentenceTransformer, util
import pandas as pd
import numpy as np
from ast import literal_eval


st.title("Mining Potential Legitimate Domains from a Typosquatted Domain")
# model_choice = st.selectbox("Select the embedding model:", ["", "Embedder-typosquat-detect-Canine", "Embedder-typosquat-detect"], index=0)
model_choice= "Embedder-typosquat-detect-Canine"
if model_choice:
    model = SentenceTransformer(f"./{model_choice}")

    domains_df = pd.read_csv(f'./{model_choice}/domains_embs.csv')
    domains_df.embedding = domains_df.embedding.apply(literal_eval)
    corpus_domains = domains_df.domain.to_list()
    corpus_embeddings = np.stack(domains_df.embedding.values).astype(np.float32)  # Ensure embeddings are float32

    st.write("Enter a potential typosquatted domain and select the number of top results to retrieve.")

    domain = st.text_input("Potential Typosquatted Domain")
    top_k = st.number_input("Top K Results", min_value=1, max_value=50, value=5, step=1)

    if st.button("Search for Legitimate Domains"):
        if domain:
            # Perform Semantic Search
            query_emb = model.encode(domain).astype(np.float32)  # Ensure query embedding is also float32
            semantic_res = util.semantic_search(query_emb, corpus_embeddings, top_k=top_k)[0]
            ids = [r['corpus_id'] for r in semantic_res]
            scores = [r['score'] for r in semantic_res]

            res_df = domains_df.loc[ids, ['domain']].copy()
            res_df['score'] = scores
            
            st.write("Mined Domains:")
            st.dataframe(res_df)
        else:
            st.warning("Please enter a domain to perform the search.")