import os
from queue import Queue
import gradio as gr
import argilla as rg
from argilla.webhooks import webhook_listener
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize Argilla client
client = rg.Argilla(
    api_url=os.getenv("API_URL"),
    api_key=os.getenv("ARGILLA_API_KEY"),
)

# Get the webhook server
server = rg.get_webhook_server()

# Queue to store events for display
incoming_events = Queue()

# Set up the webhook listener for response creation
@webhook_listener(events=["response.created"])
async def update_validation_space_on_answer(response, type, timestamp):
    """
    Webhook listener that triggers when a new response is added to an answering space.
    It will automatically update the corresponding validation space with the new response.
    """
    # Store the event for display in the UI
    incoming_events.put({"event": type, "data": str(response)})
    
    # Get the record from the response
    record = response.record
    
    # Check if this is from an answering space
    dataset_name = record.dataset.name
    if not dataset_name.endswith("_responder_preguntas"):
        print(f"Ignoring event from non-answering dataset: {dataset_name}")
        return  # Not an answering space, ignore
    
    # Extract the country from the dataset name
    country = dataset_name.replace("_responder_preguntas", "")
    print(f"Processing response for country: {country}")
    
    # Connect to the validation space
    validation_dataset_name = f"{country}_validar_respuestas"
    try:
        validation_dataset = client.datasets(validation_dataset_name)
        print(f"Found validation dataset: {validation_dataset_name}")
    except Exception as e:
        print(f"Error connecting to validation dataset: {e}")
        # You would need to import the create_validation_space function
        from your_module import create_validation_space
        validation_dataset = create_validation_space(country)
    
    # Create a validation record
    validation_record = {
        "question": record.fields["question"],
        "answer": response.value,
    }
    
    # Add the record to the validation space
    validation_dataset.records.log(records=[validation_record])
    print(f"Added new response to validation space for {country}")

# Function to read the next event from the queue
def read_next_event():
    if not incoming_events.empty():
        return incoming_events.get()
    return {}

# Create Gradio interface
with gr.Blocks() as demo:
    argilla_server = client.http_client.base_url
    gr.Markdown("## Argilla Webhooks - Validation Space Updater")
    gr.Markdown(f"""
        This application listens for new responses in Argilla answering spaces and automatically 
        adds them to the corresponding validation spaces.
        
        Connected to Argilla server: {argilla_server}
        
        The webhook listens for:
        - `response.created` events from datasets ending with `_responder_preguntas`
        
        When a new response is created, it is added to the corresponding validation space
        named `[country]_validar_respuestas`.
        
        You can view the incoming events in the JSON component below.
    """)
    json_component = gr.JSON(label="Incoming response events:", value={})
    gr.Timer(1, active=True).tick(read_next_event, outputs=json_component)

# Mount the Gradio app to the FastAPI server
gr.mount_gradio_app(server, demo, path="/")

# Start the FastAPI server
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(server, host="0.0.0.0", port=7860)