import gradio as gr
from analyzer import analyze_code

# System prompt for the chatbot
CHATBOT_SYSTEM_PROMPT = (
    "You are a helpful assistant. Your goal is to help the user describe their ideal open-source repo. "
    "Ask questions to clarify what they want, their use case, preferred language, features, etc. "
    "When the user clicks 'End Chat', analyze the conversation and return about 5 keywords for repo search. "
    "Return only the keywords as a comma-separated list."
)

# Store the conversation
conversation_history = []

# Function to handle chat
def chat_with_user(user_message, history):
    from openai import OpenAI
    client = OpenAI()
    # Build the message list for the LLM
    messages = [
        {"role": "system", "content": CHATBOT_SYSTEM_PROMPT}
    ]
    for msg in history:
        messages.append({"role": "user", "content": msg[0]})
        if msg[1]:
            messages.append({"role": "assistant", "content": msg[1]})
    messages.append({"role": "user", "content": user_message})
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        max_tokens=256,
        temperature=0.7
    )
    assistant_reply = response.choices[0].message.content
    return assistant_reply

# Function to end chat and extract keywords
def extract_keywords_from_conversation(history):
    # Combine all user and assistant messages into a single string
    conversation = "\n".join([f"User: {msg[0]}\nAssistant: {msg[1]}" for msg in history if msg[1]])
    prompt = (
        "Given the following conversation between a user and an assistant about finding an ideal open-source repo, "
        "extract about 5 keywords that would be most useful for searching Hugging Face repos to find the most relevant results for the user. "
        "Return only the keywords as a comma-separated list.\n\nConversation:\n" + conversation
    )
    keywords = analyze_code(prompt)
    return keywords

with gr.Blocks() as chatbot_demo:
    gr.Markdown("## Repo Recommendation Chatbot")
    chatbot = gr.Chatbot()
    state = gr.State([])  # conversation history
    user_input = gr.Textbox(label="Your message", placeholder="Describe your ideal repo or answer the assistant's questions...")
    send_btn = gr.Button("Send")
    end_btn = gr.Button("End Chat and Extract Keywords")
    keywords_output = gr.Textbox(label="Extracted Keywords for Repo Search", interactive=False)

    def user_send(user_message, history):
        assistant_reply = chat_with_user(user_message, history)
        history = history + [[user_message, assistant_reply]]
        return history, history, ""

    def end_chat(history):
        keywords = extract_keywords_from_conversation(history)
        return keywords

    send_btn.click(user_send, inputs=[user_input, state], outputs=[chatbot, state, user_input])
    end_btn.click(end_chat, inputs=state, outputs=keywords_output)

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