import gradio as gr
from comp import generate_response
import re

# --- Constants ---
WORKFLOW_SYSTEM_PROMPT = """You are an expert in analyzing conversations and extracting user workflows.
Based on the provided chat history, identify the user's main goal or intent.
Then, break down the conversation into a series of actionable steps that represent the workflow to achieve that goal.
The output should be in two parts, clearly separated:
**Intent**: [A concise description of the user's goal]
**Steps**:
[A numbered list of steps]
"""

# --- Helper Functions ---
def parse_workflow_response(response):
    intent_match = re.search(r"\*\*Intent\*\*:\s*(.*)", response, re.IGNORECASE)
    steps_match = re.search(r"\*\*Steps\*\*:\s*(.*)", response, re.DOTALL | re.IGNORECASE)

    intent = intent_match.group(1).strip() if intent_match else "Could not determine intent."
    steps = steps_match.group(1).strip() if steps_match else "Could not determine steps."
    
    return intent, steps

# --- Gradio UI ---

with gr.Blocks() as demo:
    gr.Markdown("# Ling Playground")

    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("## Chat")
            chat_chatbot = gr.Chatbot(label="Chat", bubble_full_width=False)
            with gr.Row():
                chat_msg = gr.Textbox(
                    label="Your Message",
                    scale=4,
                )
                send_btn = gr.Button("Send", scale=1)
            
        with gr.Column(scale=1):
            gr.Markdown("## Workflow Extraction")
            intent_textbox = gr.Textbox(label="Task Intent", interactive=False)
            steps_textbox = gr.Textbox(
                label="Extracted Steps", interactive=False, lines=15
            )
    
    chat_clear = gr.ClearButton([chat_msg, chat_chatbot, intent_textbox, steps_textbox])

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        user_message = history[-1][0]
        history[-1][1] = ""
        # Main chat model call (uses default system prompt)
        for response in generate_response(user_message, history[:-1]):
            if "</think>" in response:
                parts = response.split("</think>", 1)
                thinking_text = parts[0].replace("<think>", "")
                body_text = parts[1]

                md_output = f"**Thinking...**\n```\n{thinking_text}\n```\n\n{body_text}"
                history[-1][1] = md_output
            else:
                history[-1][1] = response
            yield history

    def update_workflow(history):
        if not history or not history[-1][0]:
            return "", ""

        # The last user message is the main prompt for the workflow agent
        user_message = history[-1][0]
        # The rest of the conversation is the history
        chat_history_for_workflow = history[:-1]

        # Call the model with the workflow system prompt
        full_response = ""
        for response in generate_response(
            user_message,
            chat_history_for_workflow,
            system_prompt=WORKFLOW_SYSTEM_PROMPT
        ):
            full_response = response
        
        intent, steps = parse_workflow_response(full_response)
        return intent, steps

    # Handler for pressing Enter in the textbox
    (   chat_msg.submit(user, [chat_msg, chat_chatbot], [chat_msg, chat_chatbot], queue=False)
        .then(bot, chat_chatbot, chat_chatbot)
        .then(update_workflow, chat_chatbot, [intent_textbox, steps_textbox])
    )

    # Handler for clicking the Send button
    (   send_btn.click(user, [chat_msg, chat_chatbot], [chat_msg, chat_chatbot], queue=False)
        .then(bot, chat_chatbot, chat_chatbot)
        .then(update_workflow, chat_chatbot, [intent_textbox, steps_textbox])
    )

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