import gradio as gr
from huggingface_hub import InferenceClient
import difflib
import os
from sentence_transformers import SentenceTransformer
import torch
import numpy as np

# load knowledge text
with open("knowledge.txt", "r", encoding="utf-8") as f:
    knowledge_text = f.read()

chunks = [chunk.strip() for chunk in knowledge_text.split("\n\n") if chunk.strip()]
embedder = SentenceTransformer('all-MiniLM-L6-v2')
chunk_embeddings = embedder.encode(chunks, convert_to_tensor=True)

def get_relevant_context(query, top_k=3):
    query_embedding = embedder.encode(query, convert_to_tensor=True)
    query_embedding = query_embedding / query_embedding.norm()
    norm_chunk_embeddings = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    similarities = torch.matmul(norm_chunk_embeddings, query_embedding)
    top_k_indices = torch.topk(similarities, k=top_k).indices.cpu().numpy()
    return "\n\n".join([chunks[i] for i in top_k_indices])

custom_theme = gr.themes.Soft(
    primary_hue="teal",
    secondary_hue="stone",
    neutral_hue="teal",
    spacing_size="lg",
    radius_size="lg",
    text_size="lg"
)

client = InferenceClient("google/gemma-2-2b-it")

def is_driving_related(message):
    driving_keywords = [
        "drive", "driving", "permit", "car", "road", "lane", "traffic",
        "license", "parallel", "park", "stop sign", "brake", "accelerate",
        "merge", "intersection", "seatbelt", "speed limit", "turn signal",
        "pulled over", "parking", "roundabout"
    ]
    message_words = message.lower().split()
    for word in message_words:
        for keyword in driving_keywords:
            if difflib.SequenceMatcher(None, word, keyword).ratio() > 0.8:
                return True
    return False

def respond(message, history):
    # if not history:
    #     if not is_driving_related(message):
    #         yield "Hey there! I’m Drive Wise 🚗 — I can only help with driving topics like road rules, parking tips, or permit prep. What driving question do you have for me?"
    #         return

    messages = [{"role": "system", "content": """You are Drive Wise, a friendly and supportive AI chatbot designed to help new drivers learn essential driving skills and traffic laws. Your goal is to make learning to drive simple, confident, and stress-free.
Speak in a warm, approachable tone. Explain road rules, parking steps, permit test prep, and basic legal info clearly. Use step-by-step instructions and short, digestible explanations. Avoid giving complex legal advice or sounding condescending. Make responses concise and conversational for easier flow of the conversation and only have about a minimum of 2 bullet points per main point."""}]

    context = get_relevant_context(message, top_k=3)
    
    if history:
        for turn in history:
            messages.append({"role": turn["role"], "content": turn["content"]})

    messages.append({"role": "user", "content": message})
    response = ""
    for msg in client.chat_completion(messages, max_tokens=500, temperature=0.1, stream=True):
        token = msg.choices[0].delta.content        
        response += token
        yield response

with gr.Blocks(theme=custom_theme) as chatbot:
    gr.Image(
    value="drive_wise_alt_logo.png",
    show_label = False,
    show_share_button = False,
    show_download_button = False
    )

# # About text
# about_text = """
# # **Drive Wise**: Your partner for the road 🚗  
# This chatbot will help you learn more about driving!  
# Ask me about road rules, parking, your learner’s permit, or how to handle situations like being pulled over.  
# **Please note that these laws are mostly, if not, only applicable for New York!**  
# (It is recommended that you use this chat in light mode for better usability! 💡)
# """

gr.ChatInterface(respond, type="messages")

chatbot.launch()