import os
import gradio as gr
from huggingface_hub import InferenceClient
from retriever import find_similar_foundations

# -------------------------------------------------------------------
# 1. Setup client for chatbot
# -------------------------------------------------------------------
# Use my token stored as a Space secret for inference
client = InferenceClient(
    provider="featherless-ai",
    api_key=os.environ["HF_TOKEN_inf"]
)


# -------------------------------------------------------------------
# 2. Chatbot function
# -------------------------------------------------------------------
def chat_with_model(message, history):
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": message})

    response = client.chat.completions.create(
        model="mistralai/Mistral-7B-Instruct-v0.2",
        messages=messages,
        max_tokens=512,
    )

    reply = response.choices[0].message.content
    history.append((message, reply))
    return history, history


# -------------------------------------------------------------------
# 3. Foundations Retriever function (for UI)
# -------------------------------------------------------------------
def retrieve_foundations(query, top_k):
    results = find_similar_foundations(query, top_k=top_k)
    output = "\n\n".join(
        [f"**{r['rank']}. {r['title']}**\n{r['purpose']}\n(similarity: {r['similarity']:.4f})"
         for r in results]
    )
    return output


# -------------------------------------------------------------------
# 4. Gradio Interface
# -------------------------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Mistral Perspective Chatbot & Foundation Finder")

    with gr.Tab("💬 Chatbot"):
        chatbot = gr.Chatbot()
        msg = gr.Textbox(placeholder="Ask me anything...", show_label=False)
        state = gr.State([])  # keeps conversation history
        msg.submit(chat_with_model, [msg, state], [chatbot, state])


    with gr.Tab("🔎 Find Aligned Foundations"):
        perspective = gr.Textbox(
            label="Enter your philanthropic perspective",
            placeholder="e.g. Environmental philanthropist emphasizing animal protection while fostering children's education"
        )
        top_k = gr.Slider(1, 5, value=2, step=1, label="Number of results")
        output = gr.Dataframe(headers=["Title", "Purpose", "similarity"], wrap=True)

        btn = gr.Button("Find Foundations")
        btn.click(fn=retrieve_foundations, inputs=[perspective, top_k], outputs=output)

demo.launch(server_name="0.0.0.0", server_port=7860)