import gradio as gr
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import tempfile
import os

# Load embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

def create_chatbot(role, context, info, conv_starter, file):
    knowledge_chunks = []
    chunk_embeddings = None

    if file:
        with tempfile.NamedTemporaryFile(delete=False, mode="wb") as temp:
            temp.write(file)
            temp_path = temp.name


        with open(temp_path, 'r', encoding='utf-8') as f:
            text = f.read()
        os.unlink(temp_path)

        knowledge_chunks = [chunk.strip() for chunk in text.split('\n\n') if chunk.strip()]

        if knowledge_chunks:
            chunk_embeddings = embedding_model.encode(knowledge_chunks)
            status = f"✅ Loaded {len(knowledge_chunks)} knowledge chunks"
        else:
            status = "❌ File is empty"
    else:
        status = "⚠️ No file uploaded"

    # Store all chatbot settings and knowledge in a dict (state)
    return status, {
        "role": role,
        "context": context,
        "info": info,
        "conv_starter": conv_starter,
        "knowledge": knowledge_chunks,
        "embeddings": chunk_embeddings
    }

def respond(message, history, state):
    # Special info queries
    if any(keyword in message.lower() for keyword in ["more info", "contact", "information", "email", "details"]):
        return state["info"]

    # No knowledge base loaded
    if not state.get("knowledge"):
        return "⚠️ Please upload knowledge base first"

    # Embed user query
    query_embedding = embedding_model.encode([message])
    similarities = cosine_similarity(query_embedding, state["embeddings"])[0]
    max_index = np.argmax(similarities)
    max_similarity = similarities[max_index]

    # If similar enough, return the best chunk
    if max_similarity > 0.45:
        return state["knowledge"][max_index]

    # Fallback
    return f"{state['role']}\n{state['context']}\nI can't help with that specific question."

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🤖 Custom Chatbot Creator")
    gr.Markdown("Configure every aspect of your chatbot below")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Configuration Panel")

            with gr.Group():
                role = gr.Textbox(label="Role", value="AI Assistant specialized in technical queries")
                context = gr.Textbox(label="Context", value="Focus on providing concise, accurate answers based on the knowledge base")
                info = gr.Textbox(label="Contact Info", value="For more information, contact support@example.com")
                conv_starter = gr.Textbox(label="Conversation Starter", value="Ask me about topics in the knowledge base")

            with gr.Group():
                file = gr.File(label="Knowledge Base (.txt only)", file_types=[".txt"], type="binary")
                create_btn = gr.Button("Create Chatbot", variant="primary")

            status = gr.Textbox(label="Status", interactive=False)

            gr.Markdown("### Instructions")
            gr.Markdown("1. Configure all fields\n2. Upload knowledge file\n3. Click 'Create Chatbot'\n4. Chat in the right panel")

        with gr.Column(scale=2):
            gr.Markdown("## Chat Interface")
            state = gr.State({})
            chatbot = gr.ChatInterface(
                respond,
                chatbot=gr.Chatbot(
                    height=500,
                    type="messages",  # Use OpenAI-style messages for future compatibility
                    avatar_images=(None, (None, "https://i.imgur.com/7kQEsHU.png"))
                ),
                textbox=gr.Textbox(placeholder="Type your message...", container=False, autofocus=True),
                submit_btn="Ask"
            )

    create_btn.click(
        create_chatbot,
        inputs=[role, context, info, conv_starter, file],
        outputs=[status, state]
    )

if __name__ == "__main__":
    app.launch()