import gradio as gr
import json

class JSONLViewer:
    def __init__(self, file_path):
        self.file_path = file_path
        self.current_index = 0
        self.data = []
        self.load_data()

    def load_data(self):
        with open(self.file_path, 'r', encoding='utf-8') as file:
            self.data = [json.loads(line) for line in file]

    def save_data(self):
        with open(self.file_path, 'w', encoding='utf-8') as file:
            for item in self.data:
                json.dump(item, file, ensure_ascii=False)
                file.write('\n')

    def get_current_record(self):
        if 0 <= self.current_index < len(self.data):
            return self.data[self.current_index]
        return None

    def update_status(self, status):
        current_record = self.get_current_record()
        if current_record:
            current_record['Status'] = status
            self.save_data()

    def move_prev(self):
        if self.current_index > 0:
            self.current_index -= 1
        return self.get_current_record()

    def move_next(self):
        if self.current_index < len(self.data) - 1:
            self.current_index += 1
        return self.get_current_record()

viewer = JSONLViewer('plik.jsonl')

def update_ui(record):
    if record:
        instruction = record['conversations'][0]['value'] if record['conversations'] else ""
        chosen = record['chosen']['value'] if 'chosen' in record else ""
        rejected = record['rejected']['value'] if 'rejected' in record else ""
        chosen_score = record['chosen_score'] if 'chosen_score' in record else ""
        rejected_score = record['rejected_score'] if 'rejected_score' in record else ""
        status = record.get('Status', "")
        return instruction, chosen, rejected, chosen_score, rejected_score, status
    return "", "", "", "", "", ""

def on_prev():
    record = viewer.move_prev()
    return update_ui(record)

def on_next():
    record = viewer.move_next()
    return update_ui(record)

def on_ok():
    viewer.update_status("ok")
    record = viewer.get_current_record()
    return update_ui(record)

def on_rejected():
    viewer.update_status("rejected")
    record = viewer.get_current_record()
    return update_ui(record)

with gr.Blocks() as demo:
    instruction = gr.Markdown(label="Instruction")
    chosen = gr.Markdown(label="CHOSEN")
    rejected = gr.Markdown(label="REJECTED")
    chosen_score = gr.Number(label="CHOSEN SCORE")
    rejected_score = gr.Number(label="REJECTED SCORE")
    status = gr.Textbox(label="STATUS")

    with gr.Row():
        prev_btn = gr.Button("PREV")
        next_btn = gr.Button("NEXT")
        ok_btn = gr.Button("OK")
        rejected_btn = gr.Button("REJECTED")

    prev_btn.click(on_prev, outputs=[instruction, chosen, rejected, chosen_score, rejected_score, status])
    next_btn.click(on_next, outputs=[instruction, chosen, rejected, chosen_score, rejected_score, status])
    ok_btn.click(on_ok, outputs=[instruction, chosen, rejected, chosen_score, rejected_score, status])
    rejected_btn.click(on_rejected, outputs=[instruction, chosen, rejected, chosen_score, rejected_score, status])

    # Initialize UI with first record
    initial_record = viewer.get_current_record()
    if initial_record:
        init_instruction, init_chosen, init_rejected, init_chosen_score, init_rejected_score, init_status = update_ui(initial_record)
        instruction.value = init_instruction
        chosen.value = init_chosen
        rejected.value = init_rejected
        chosen_score.value = init_chosen_score
        rejected_score.value = init_rejected_score
        status.value = init_status

demo.launch()