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=500
    )
    
    return response.choices[0].message.content.strip()

def display_image():
    return "ByteShield_New.png"


theme = gr.themes.Soft().set(
    body_background_fill="#20235c",
    body_text_color="#c7a50e",
    block_background_fill="#b9d3eb",
    block_border_color="#FFD700",
    button_primary_background_fill="#FFD700",
    button_primary_text_color="#dae9f7"
)

with gr.Blocks(theme=theme) as demo:
    with gr.TabItem("Resources"):
            gr.Markdown("## Resources")

            gr.Button(
                "🗓️ Period Tracker",
                link="https://drive.google.com/file/d/1_KNELAUDLLidwAT3fs2JBuO1yPgMGoDv/view"
            )

            gr.Button(
                "🤰 New Moms Support Group",
                link="https://www.instagram.com/firsttimemomsacademy/"
            )

css = """
button[data-testid="example"] {
    color: #FFD700 !important;
}
"""

with gr.Blocks(theme=theme, css = css) as demo:
    gr.Image(display_image(), show_label=False, interactive=False)
    chatbot = gr.ChatInterface(
    fn = respond,
    cache_examples = False,
    textbox=gr.Textbox(placeholder="Ask me anything!", container=False, scale=7),
    title = "ByteShield - Your AI Guardian 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?"]
)
#title_hotline= "# Select To Get Hotline Number"

#with gr.Tabs():
        #with gr.TabItem("Resources"): 
            #gr.Markdown("### Resources")
            #open_google = gr.Button(value="🗓️ Period Tracker", link="https://drive.google.com/file/d/1_KNELAUDLLidwAT3fs2JBuO1yPgMGoDv/view")
            #open_google = gr.Button(value="👩🏻‍🍼 New Moms Support Group", link="https://www.instagram.com/firsttimemomsacademy/")
            
        
   # with gr.TabItem("Call a Hotline"): 
        #gr.Markdown(title_hotline)
        #gr.Markdown(hotline_text)
        #dropdown = gr.Dropdown(choices=["General Health", "Maternal Mental Health", "Domestic Violence", "Postpartum Support"],
        #label="Choose Your Hotline"
                                  
        #output = gr.Textbox(label="Hotline Info", interactive=False)
    
    #dropdown.change(fn=show_info, inputs=dropdown, outputs=output)
     #with gr.Tab("Educational PDFs"):
       # gr.Markdown("### 📘 Helpful Resources")


demo.launch()