from sentence_transformers import SentenceTransformer
import gradio as gr
from huggingface_hub import InferenceClient
import numpy as np
import torch
import os
import gradio as gr
#pip install https://gradio-builds.s3.amazonaws.com/75c684efb87624bee2fb63b08122564e6538509e/gradio-6.17.3-py3-none-any.whl


def image_classifier(inp):
    return {'cat': 0.3, 'dog': 0.7}

demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
demo.launch()


with open("knowledge.txt", "r", encoding="utf-8") as file:
  knowledge_base = file.read()

def preprocess_text(text):
  cleaned_text = text.strip()

  chunks = cleaned_text.split("\n")

  cleaned_chunks = []

  for chunk in chunks:
    stripped_chunk = chunk.strip()
    if len (stripped_chunk)>0:
      cleaned_chunks.append(stripped_chunk)


  #print(cleaned_chunks)

  #print (len(cleaned_chunks))

  return cleaned_chunks
    
cleaned_chunks = preprocess_text(knowledge_base)
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
  chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True) # Replace ... with the text_chunks list

  #print(chunk_embeddings)

  #print(chunk_embeddings.shape)

  return chunk_embeddings

chunk_embeddings = create_embeddings(cleaned_chunks)# Complete this line

def get_top_chunks(query, chunk_embeddings, text_chunks):
  query_embedding = model.encode(query, convert_to_tensor=True) # Complete this line

  query_embedding_normalized = query_embedding / query_embedding.norm()

  chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)

  similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized) # Complete this line

  #print(similarities)

  top_indices = torch.topk(similarities, k=3).indices

  #print(top_indices)

  top_chunks = []

  for i in top_indices:
    chunk = text_chunks[i]
    top_chunks.append(chunk)

  return top_chunks

top_results = get_top_chunks("Your account has been compromised", chunk_embeddings, cleaned_chunks) # Complete this line

#print(top_results)



#with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.pink)) as demo:


cleaned_chunks = preprocess_text(knowledge_base) 

client = InferenceClient("Qwen/Qwen2.5-7B-Instruct", token=os.getenv("ByteShield_Token"))
def respond(message, history):
    top_chunks = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    context = "\n".join(top_chunks)
    messages = [{"role": "system","content": f"You are a friendly, tech expert chatbot. Use this context to answer:\n{context}"}]
    
    if history:
        messages.extend(history)
        
    messages.append({"role": "user", "content": message})
    
    response = client.chat_completion(
        messages,
        max_tokens=1000
    )
    
    return response.choices[0].message.content.strip()

with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    chatbot = gr.ChatInterface(
    fn = respond,
    cache_examples = False,
    textbox=gr.Textbox(placeholder="Ask me anything!", container=False, scale=7),
    title = "ByteShield - Your AI Gaurdian for Online Safety",
    description = "Ask me anything about online safety!",
    examples = ["Generate me some strong passwords to use.", "What are some security measures I can take to stay safe online?", "Explain how a data breach works.", "How do I know if a message is a scam or not?"]
)



demo.launch()