import gradio as gr
from root import RootSignals

client = None  # Initialize client as None
custom_judge = None  # Initialize custom judge as None

def initialize_client(api_key):
    global client
    return RootSignals(api_key=api_key)

def create_judge(api_key, judge_name, intent, judge_prompt):
    global client, custom_judge
    if not client:
        client = initialize_client(api_key)
    
    # Create custom judge
    custom_judge = client.evaluators.create(
        name=judge_name,
        predicate=judge_prompt + " {{output}}",
        intent=intent,
        model="gpt-4o",
    )
    
    # Update the visibility of the evaluation and results sections
    eval_section.visible = True
    results_section.visible = True
    evaluate_btn.visible = True
    return gr.Info("Custom LLM-Judge is created successfully!")

def evaluate_response(llm_response):
    global client, custom_judge
    if not client or not custom_judge:
        return "Please create a judge first", "Please create a judge first"
    
    # Run evaluation using custom judge
    evaluation_result = custom_judge.run(response=llm_response)
    score = evaluation_result.score
    justification = evaluation_result.justification
    return score, justification

# Create the interface with a custom layout
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue")) as demo:
    gr.HTML("""<a href="https://api.visitorbadge.io/api/visitors?path=https://huggingface.co/spaces/root-signals/RootEvaluatorsDemo">
               <img src="https://api.visitorbadge.io/api/visitors?path=https://huggingface.co/spaces/root-signals/RootEvaluatorsDemo" />
               </a>""")

    with gr.Row():
        gr.Image(value="https://app.rootsignals.ai/images/root-signals-color.svg", height=70)
        gr.Markdown("<div>&nbsp;</div>")  # Add some space below the image
    gr.Markdown("# Custom Judge Demo by Root Signals")

    gr.Markdown("[Sign-up](https://app.rootsignals.ai/register) to create your API key!")
    
    api_key = gr.Textbox(
        label="🔑 Root Signals API Key",
        placeholder="Enter your Root Signals API key...",
        type="password",
        show_label=True,
    )
    
    with gr.Row():
        # Left column - Judge Creation
        with gr.Column():
            gr.Markdown("### Create Custom Judge")
            judge_name = gr.Textbox(label="👨‍⚖️ Judge Name", placeholder="Enter a name for your custom judge...", interactive=True)
            user_intent = gr.Textbox(label="👤 Intent", placeholder="Enter the high-level intent for this judge...", interactive=True)
            judge_prompt = gr.Textbox(label="📝 Custom Judge Prompt", placeholder="Enter the custom judge prompt...", interactive=True)
            create_judge_btn = gr.Button("✨ CREATE JUDGE", variant="primary")
            info_message = gr.Info()
            
            # Evaluation section (initially hidden)
            eval_section = gr.Column(visible=False)
            with eval_section:
                gr.Markdown("### Evaluate Response")
                llm_response = gr.Textbox(label="🤖 LLM Response", placeholder="Enter the LLM response to be evaluated...", interactive=True)
                evaluate_btn = gr.Button("🧐 EVALUATE", variant="primary", visible=False)
        
        # Right column - Results
        results_section = gr.Column(visible=False)
        with results_section:
            score = gr.Textbox(label="📊 Score (between 0 and 1)", interactive=False)
            justification = gr.TextArea(label="💬 Justification", interactive=False)
    
    # Button click events
    create_judge_btn.click(
        fn=create_judge,
        inputs=[api_key, judge_name, user_intent, judge_prompt],
        outputs=info_message
    )
    
    evaluate_btn.click(
        fn=evaluate_response,
        inputs=[llm_response],
        outputs=[score, justification]
    )

    gr.Markdown("[Homepage](https://www.rootsignals.ai/) | [Python SDK Docs](https://sdk.rootsignals.ai/en/latest/)")

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