import gradio as gr
from bs4 import BeautifulSoup
from langchain.embeddings import HuggingFaceEmbeddings
import pickle
import torch
import io

class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else: return super().find_class(module, name)


def get_hugging_face_model():
    model_name = "mchochlov/codebert-base-cd-ft"
    hf = HuggingFaceEmbeddings(model_name=model_name)
    return hf


def get_db():
    with open("codesearchdb.pickle", "r") as f:
        #db = CPU_Unpickler(f).load()
        db = f.load()
    return db


def get_similar_links(query, db, embeddings):
    embedding_vector = embeddings.embed_query(query)
    docs_and_scores = db.similarity_search_by_vector(embedding_vector, k = 10)
    hrefs = []
    for docs in docs_and_scores:
        html_doc = docs.page_content
        soup = BeautifulSoup(html_doc, 'html.parser')
        href = [a['href'] for a in soup.find_all('a', href=True)]
        hrefs.append(href)
    links = []
    for href_list in hrefs:
        for link in href_list:
            links.append(link)
    return links


def find_similar_questions(text_input):
    embedding_vector = get_hugging_face_model()
    db = get_db()
    query = text_input
    answer = get_similar_links(query, db, embedding_vector)
    return "\n".join(set(answer))


iface = gr.Interface(
    fn=find_similar_questions,
    inputs=gr.inputs.Textbox(lines=20, label="Enter a Code Example From Leetcode"),
    outputs=gr.outputs.Textbox(label="Similar Questions on Leetcode"),
    title="📒 DSASearch Engine 🤖",
    description="Find similar questions on Leetcode based on a code example.",
    allow_flagging=False,
)

iface.launch()