import gradio as gr
import os
import numpy as np

from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer

# ---------------------------
# Load Knowledge Base
# ---------------------------

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

chunks = [chunk.strip() for chunk in knowledge_base.split("\n\n") if chunk.strip()]

# ---------------------------
# Embedding Model
# ---------------------------

embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
chunk_embeddings = embedding_model.encode(chunks)

# ---------------------------
# Language Model
# ---------------------------

client = InferenceClient(
    "Qwen/Qwen2.5-7B-Instruct",
    token=os.environ.get("HF_TOKEN")
)

# ---------------------------
# Retrieval Function
# ---------------------------

def retrieve_context(query, top_k=3):
    query_embedding = embedding_model.encode([query])[0]

    similarities = np.dot(chunk_embeddings, query_embedding)

    top_indices = np.argsort(similarities)[-top_k:][::-1]

    context = "\n\n".join([chunks[i] for i in top_indices])

    return context

# ---------------------------
# Chatbot Response Function
# ---------------------------

def respond(message, history):

    context = retrieve_context(message)

    messages = [
        {
            "role": "system",
            "content": f"""
You are Pathway AI, an educational guidance assistant designed to help students discover opportunities, resources, mentorship programs, scholarships, and career pathways.

Your mission is to make educational and professional opportunities more accessible, especially for students who may not have access to strong guidance networks.

Guidelines:
- Use the provided context as your primary source of information.
- If the answer is not available in the provided context, clearly state that you do not currently have that information in your knowledge base.
- Do not invent scholarships, organizations, opportunities, or facts.
- Be friendly, encouraging, supportive, and informative.
- Keep responses concise and easy to understand.
- When appropriate, suggest actionable next steps.
- When discussing careers, recommend relevant skills, resources, and opportunities.
- When discussing scholarships or programs, summarize eligibility and benefits when available.

Context:
{context}
"""
        },
        {
            "role": "user",
            "content": message
        }
    ]

    response = client.chat_completion(
        messages=messages,
        max_tokens=500
    )

    return response.choices[0].message.content.strip()
# ---------------------------
# UI Theme
# ---------------------------

theme = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="pink"
)

css = """
footer {
    display: none;
}

.gradio-container {
    max-width: 1000px !important;
    margin: auto !important;
    background: linear-gradient(
        180deg,
        #faf5ff 0%,
        #fdf4ff 100%
    );
}

.hero-card {
    text-align: center;
    background: white;
    padding: 25px;
    border-radius: 20px;
    margin-bottom: 20px;
    box-shadow: 0 4px 12px rgba(0,0,0,0.08);
}

.feature-card {
    background: white;
    padding: 20px;
    border-radius: 20px;
    margin-bottom: 20px;
    box-shadow: 0 4px 12px rgba(0,0,0,0.08);
}


.hero-card,
.feature-card {
    color: #1f2937 !important;
}

.hero-card h1,
.hero-card h2,
.hero-card h3,
.feature-card h1,
.feature-card h2,
.feature-card h3,
.feature-card p {
    color: #1f2937 !important;
}

body {
    color: #1f2937 !important;
}

.gradio-container {
    color: #1f2937 !important;
}

"""

# ---------------------------
# Build UI
# ---------------------------

with gr.Blocks(theme=theme, css=css) as demo:

    gr.Image(
        "banner.png.png",
        show_label=False,
        container=False
    )

    gr.HTML("""
    <div class="feature-card">
        <h3>🚀 What can Pathway AI help with?</h3>

        🎓 Scholarships<br>
        🚀 Opportunities<br>
        💜 Women in STEM<br>
        🤝 Mentorship<br>
        💻 Learning Resources<br>
        🧭 Career Exploration
    </div>
    """)

    gr.Markdown(
        "### 💡 Try asking one of the example questions below to get started!"
    )

    chatbot = gr.Chatbot(
        height=500,
        show_label=False
    )

    msg = gr.Textbox(
        placeholder="Ask Pathway AI a question...",
        label=""
    )

    with gr.Row():
        send_btn = gr.Button("Send")
        clear_btn = gr.Button("Clear Chat")

    examples = gr.Examples(
        examples=[
            ["What scholarships are available for women in STEM?"],
            ["How can I find a mentor in technology?"],
            ["What opportunities are available for high school students interested in AI?"],
            ["I want to become a software engineer. Where should I start?"],
            ["What coding resources are best for beginners?"],
            ["Tell me about women leaders in STEM."],
            ["I want to learn machine learning. What resources would you recommend?"]
        ],
        inputs=msg
    )

    def chat(message, history):
        response = respond(message, history)
        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": response}
        ]
        return "", history

    send_btn.click(
        chat,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot]
    )

    msg.submit(
        chat,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot]
    )

    clear_btn.click(
        lambda: [],
        outputs=chatbot
    )

    gr.Markdown("""
---
Built by KWK '26 AI/ML Scholars - Group A3 💜
""")

demo.launch()