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

client = rg.Argilla(
    api_url=os.getenv("ARGILLA_API_URL"),
    api_key=os.getenv("ARGILLA_API_KEY"),
)
server = rg.get_webhook_server()
incoming_events = Queue()

@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.
    """
    try:
        incoming_events.put({"event": type, "timestamp": str(timestamp)})
        
        record = response.record
        original_user_id = str(response.user_id)
        
        dataset_name = record.dataset.name
        if not dataset_name.endswith("Responder"):
            print(f"Ignoring event from non-answering dataset: {dataset_name}")
            return
    
        print(f"Processing response for dataset: {dataset_name}")

        validation_dataset_name = dataset_name.replace("Responder", "Validar")
        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}")

        response_dict = response.to_dict()
        answers = [response_dict["values"]['answer_1']['value']]
        if "answer_2" in response_dict["values"]:
            answers.append(response_dict["values"]['answer_2']['value'])
        if "answer_3" in response_dict["values"]:
            answers.append(response_dict["values"]['answer_3']['value'])

        new_records = []

        for answer in answers:
            if answer:
                validation_record = {
                    "question": record.fields["question"],
                    "answer": answer,
                    "metadata": {
                        "original_responder_id": original_user_id,
                        "original_dataset": dataset_name
                    }
                }
                new_records.append(validation_record)
        
        validation_dataset.records.log(records=new_records)
        print(f"Added new response to {update_validation_space_on_answer}")
        
    except Exception as e:
        print(f"Error in webhook handler: {e}")
        incoming_events.put({"event": "error", "error": str(e)})

def read_next_event():
    if not incoming_events.empty():
        return incoming_events.get()
    return {}

with gr.Blocks() as demo:
    argilla_server = client.http_client.base_url
    gr.Markdown("# Argilla Webhooks - Top Secret")
    gr.Timer(1, active=True).tick(read_next_event)

gr.mount_gradio_app(server, demo, path="/")

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