import streamlit as st
import json
from datetime import datetime
from pathlib import Path
from uuid import uuid4
import os
from huggingface_hub import CommitScheduler

# Setup the directory for saving data
JSON_DATASET_DIR = Path("feedback_dataset")
JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)

# Define a unique file name
JSON_DATASET_PATH = JSON_DATASET_DIR / f"feedback-{uuid4()}.json"

# Scheduler configuration for your dataset repo
scheduler = CommitScheduler(
    repo_id="Utsav2001/Feedback",  # Replace with your dataset repo
    repo_type="dataset",
    folder_path=JSON_DATASET_DIR,  # Local directory to sync
    path_in_repo="data",  # Path in the dataset repository
    token=os.getenv('hf_write')  # Ensure your HF write token is set
)

# Streamlit UI
st.title("OpenROAD-Assistant")

if "prompts" not in st.session_state:
    st.session_state.prompts = []
if "responses" not in st.session_state:
    st.session_state.responses = []
if "show_feedback" not in st.session_state:
    st.session_state.show_feedback = False
if "feedbacks" not in st.session_state:
    st.session_state.feedbacks = []
if "feedback_type" not in st.session_state:
    st.session_state.feedback_type = None

# Function to save data using Hugging Face CommitScheduler
def save_feedback_to_hub(prompt, ai_response, feedback_type, feedback_comment):
    try:
        with scheduler.lock:
            with JSON_DATASET_PATH.open("a") as f:
                json.dump(
                    {
                        "prompt": prompt,
                        "ai_response": ai_response,
                        "feedback_type": feedback_type,
                        "feedback_comment": feedback_comment,
                        "datetime": datetime.now().isoformat(),
                    },
                    f,
                )
                f.write("\n")
        # Push the changes to Hugging Face Hub
        scheduler.push_to_hub()
        return "Feedback saved successfully to Hugging Face Hub!"
    except Exception as e:
        return f"Error saving feedback: {e}"

# User input section
user_input = st.text_input("Enter your message:")
if st.button("Get Response"):
    # Simulate an AI response (replace with actual AI model logic)
    ai_response = f"This is a response to: {user_input}"
    st.session_state.prompts.append(user_input)
    st.session_state.responses.append(ai_response)
    st.session_state.show_feedback = True
    st.write("AI Response:", ai_response)

# Feedback section
if st.session_state.show_feedback and len(st.session_state.prompts) > 0 and len(st.session_state.responses) > 0:
    st.write("Provide feedback for the last interaction:")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("👍 Thumbs Up"):
            st.session_state.feedback_type = "Positive"
            st.session_state.show_feedback_box = True
    with col2:
        if st.button("👎 Thumbs Down"):
            st.session_state.feedback_type = "Negative"
            st.session_state.show_feedback_box = True

    if "show_feedback_box" in st.session_state and st.session_state.show_feedback_box:
        feedback_comment = st.text_area("Your Feedback:")
        if st.button("Submit Feedback"):
            save_status = save_feedback_to_hub(
                st.session_state.prompts[-1],
                st.session_state.responses[-1],
                st.session_state.feedback_type,
                feedback_comment
            )
            st.success(save_status)
            st.session_state.show_feedback_box = False
            st.session_state.show_feedback = False