import gradio as gr
import json
from agent import build_graph
from langchain_core.messages import HumanMessage

# Load questions from metadata.jsonl
def load_questions(jsonl_path):
	questions = []
	with open(jsonl_path, 'r', encoding='utf-8') as f:
		for line in f:
			try:
				obj = json.loads(line)
				if 'Question' in obj:
					questions.append(obj['Question'])
			except Exception:
				continue
	return questions

questions_list = load_questions('metadata.jsonl')
# Basic filter for violence/intimate (simple keyword check)
def is_safe_question(q):
	unsafe_keywords = ['kill', 'murder', 'sex', 'intimate', 'violence', 'abuse']
	q_lower = q.lower()
	return not any(word in q_lower for word in unsafe_keywords)

# Backend logic: send question to agent
def get_answer(question, history):
	if not is_safe_question(question):
		return "Sorry, this question is not allowed.", history
	try:
		graph = build_graph(provider="groq")
		messages = [HumanMessage(content=question)]
		result = graph.invoke({"messages": messages})
		# Get last message as answer
		answer = result["messages"][-1].content if result["messages"] else "No answer."
		history = history + [(question, answer)]
		return answer, history
	except Exception as e:
		return f"Error: {str(e)}", history

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="purple")) as demo:
	gr.Markdown("# AI Agent Q&A", elem_id="title")
	gr.Markdown("Ask your own question or select one from the list below.")
	with gr.Row():
		question_box = gr.Textbox(label="Type your question", lines=3)
		question_list = gr.Dropdown(choices=questions_list, label="Or choose a question", interactive=True)
	submit_btn = gr.Button("Submit", elem_id="submit-btn")
	reset_btn = gr.Button("Reset", elem_id="reset-btn")
	answer_box = gr.Textbox(label="Answer", interactive=False)
	with gr.Accordion("Show previous Q&A", open=False):
		history_box = gr.Dataframe(headers=["Question", "Answer"], datatype=["str", "str"], interactive=False)

	state = gr.State([])
	def submit_fn(q_text, q_list, history):
		question = q_text if q_text else q_list
		if not question:
			return "Please enter or select a question.", history
		return get_answer(question, history)

	def reset_fn():
		return "", []

	submit_btn.click(
	submit_fn,
	inputs=[question_box, question_list, state],
	outputs=[answer_box, state],
		api_name="submit",
	)
	reset_btn.click(
	reset_fn,
	inputs=[],
	outputs=[answer_box, state],
		api_name="reset",
	)
	state.change(lambda h: h, inputs=state, outputs=history_box)

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