import json

import gradio as gr
from distilabel.llms import TransformersLLM
from distilabel.steps.tasks.argillalabeller import ArgillaLabeller

llm = TransformersLLM(model="microsoft/Phi-3-mini-4k-instruct")
task = ArgillaLabeller(llm=llm)
task.load()

@spaces.GPU
def process_records_gradio(records, example_records, field, question):
    try:
        # Convert string inputs to dictionaries
        records = json.loads(records)
        example_records = json.loads(example_records) if example_records else None
        field = json.loads(field) if field else None
        question = json.loads(question) if question else None

        if not field and not question:
            return "Error: Either field or question must be provided"

        task.set_runtime_parameters(
            {
                "fields": [field] if field else None,
                "question": question,
                "example_records": example_records,
            }
        )

        results = []
        for record in records:
            output = next(task.process(inputs=[{"records": record}]))
            results.append(output[0]["suggestions"])

        return json.dumps({"results": results}, indent=2)
    except Exception as e:
        return f"Error: {str(e)}"


interface = gr.Interface(
    fn=process_records_gradio,
    inputs=[
        gr.Code(label="Records (JSON)", language="json", lines=5),
        gr.Code(label="Example Records (JSON, optional)", language="json", lines=5),
        gr.Code(label="Field (JSON, optional)", language="json"),
        gr.Code(label="Question (JSON, optional)", language="json"),
    ],
    outputs=gr.Code(label="Suggestions", language="json", lines=10),
    title="Record Processing Interface",
    description="Enter JSON data for records, example records, field, and question. At least one of field or question must be provided.",
)

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