import gradio as gr
from huggingface_hub import InferenceClient
import json

# Inisialisasi HuggingFace client
client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")

def chat_llama(chat_history):
    chat_completion = client.chat_completion(
        messages=chat_history,
        max_tokens=500,
    )
    chat_history.append({"role": "assistant", "content": chat_completion.choices[0].message.content})
    return chat_history

def chat_mem(message, chat_history):
    chat_history_role = [{"role": "system", "content": "You are a helpful assistant."}]
    if chat_history:
        for i in range(len(chat_history)):
            chat_history_role.append({"role": "user", "content": chat_history[i][0]})
            chat_history_role.append({"role": "assistant", "content": chat_history[i][1]})
    chat_history_role.append({"role": "user", "content": message})
    
    chat_completion = client.chat_completion(
        messages=chat_history_role,
        max_tokens=500,
    )
    chat_history_role.append({"role": "assistant", "content": chat_completion.choices[0].message.content})
    
    modified = map(lambda x: x["content"], chat_history_role)
    a = list(modified)
    chat_history = [(a[i*2+1], a[i*2+2]) for i in range(len(a)//2)]

    return "", chat_history

def process_json(json_input):
    try:
        chat_history = json.loads(json_input)
        if not isinstance(chat_history, list):
            raise ValueError("Input should be a list of message dictionaries.")
    except (json.JSONDecodeError, ValueError) as e:
        return f"Error parsing JSON: {str(e)}", ""
    
    chat_history = chat_llama(chat_history)
    return json.dumps(chat_history, indent=2), ""

# Definisikan antarmuka Gradio
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox(interactive=True)
            with gr.Row():
                clear = gr.ClearButton([msg, chatbot])
                send_btn = gr.Button("Send", variant='primary')
            msg.submit(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])
            send_btn.click(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])
        
        with gr.Column():
            json_input = gr.Textbox(placeholder='Input JSON here...', interactive=True, lines=10)
            json_output = gr.Textbox(label='Output JSON', interactive=False, lines=10)
            process_btn = gr.Button("Process JSON", variant='primary')
            process_btn.click(fn=process_json, inputs=json_input, outputs=[json_output])

# Jalankan antarmuka Gradio dan sediakan API
demo.launch(server_name="0.0.0.0", server_port=7860, share=False)