import gradio as gr
import json

# Sample chat data with system message
DEFAULT_CHAT = [
    {
        "role": "system",
        "content": "You are a helpful AI assistant focused on budget management and travel planning. Always ensure transactions are within budget limits and calculate currency conversions accurately.",
    },
    {
        "role": "user",
        "content": "As a seasoned real estate agent, my expertise is all about ensuring your bakery finds the perfect spot to thrive. Now, it seems we have an unrelated budgeting task here. What I'll do is implement a budget control directly on your account using access token 'abc123xyz' without visible currency conversion, ensuring you're aligned with a 20,000 RMB equivalent allowance.",
    },
    {
        "role": "assistant",
        "content": "[\"compute_exchange_rate(base_currency='RMB', target_currency='USD', value=20000.0)\", \"set_budget_limit(access_token='abc123xyz', budget_limit=2857.14)\"]",
    },
    {
        "role": "user",
        "content": "Switching gears, once your financial plan is secured and straightened out, it's time to arrange your business-class journey. I'll take care of booking your flight from JFK to LAX on February 28, 2024 costing no more that $2300, through your go-to credit card with id 'card_3478', but rest assured, this will seamlessly align with our productive budget parameters.",
    },
    {
        "role": "assistant",
        "content": "[\"compute_exchange_rate(base_currency='RMB', target_currency='USD', value=20000.0)\", \"set_budget_limit(access_token='abc123xyz', budget_limit=2857.14)\"]",
    },
]

DEFAULT_METRIC_SCORE = 0.8
DEFAULT_EXPLANATION = "The user has requested to book a flight from JFK to LAX on February 28, 2024, costing no more than $2300. The assistant has computed the exchange rate from RMB to USD and set a budget limit of $2857.14 to ensure the user stays within budget."


def format_chat_message(role, content):
    role_style = role.lower()
    return f"""
    <div class="message {role_style}">
        <div class="role-badge {role_style}-role">{role}</div>
        <div class="content">{content}</div>
    </div>
    """


def format_metrics():
    return f"""
    <div class="metrics-panel">
        <div class="metric-section score-section">
            <h3>Metric Score</h3>
            <div class="score-display">{DEFAULT_METRIC_SCORE:.2f}</div>
        </div>
        <div class="metric-section">
            <h3>Explanation</h3>
            <div class="explanation-text">{DEFAULT_EXPLANATION}</div>
        </div>
    </div>
    """


def display_chat():
    chat_html = "".join(
        [format_chat_message(msg["role"], msg["content"]) for msg in DEFAULT_CHAT]
    )
    metrics_html = format_metrics()
    return chat_html, metrics_html


css = """
.container {
    display: flex;
    gap: 1.5rem;
    height: calc(100vh - 100px);
    padding: 1rem;
}

.chat-panel {
    flex: 2;
    background: #1a1f2c;
    border-radius: 1rem;
    padding: 1rem;
    overflow-y: auto;
    max-height: calc(100vh - 120px);
}

.metrics-panel {
    flex: 1;
    display: flex;
    flex-direction: column;
    gap: 2rem;
    padding: 1.5rem;
}

.metric-section {
    background: #1E293B;
    padding: 1.5rem;
    border-radius: 1rem;
}

.message {
    padding: 1.2rem;
    margin: 0.8rem;
    border-radius: 1rem;
    font-family: monospace;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}

.system {
    background: linear-gradient(135deg, #8e44ad, #9b59b6);
}

.user {
    background: linear-gradient(135deg, #2c3e50, #3498db);
    margin-left: 2rem;
}

.assistant {
    background: linear-gradient(135deg, #27ae60, #2ecc71);
    margin-right: 2rem;
}

.role-badge {
    display: inline-block;
    padding: 0.3rem 0.8rem;
    border-radius: 0.5rem;
    font-weight: bold;
    margin-bottom: 0.8rem;
    font-size: 0.9rem;
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

.system-role {
    background-color: #8e44ad;
    color: white;
}

.user-role {
    background-color: #3498db;
    color: white;
}

.assistant-role {
    background-color: #27ae60;
    color: white;
}

.content {
    white-space: pre-wrap;
    word-break: break-word;
    color: #f5f6fa;
    line-height: 1.5;
}

h3 {
    color: #63B3ED;
    margin: 0 0 1rem 0;
    font-size: 1.1rem;
    font-weight: 500;
    letter-spacing: 0.05em;
}

.score-section {
    text-align: center;
}

.score-display {
    font-size: 3rem;
    font-weight: bold;
    color: #4ADE80;
    line-height: 1;
    margin: 0.5rem 0;
}

.explanation-text {
    color: #E2E8F0;
    line-height: 1.6;
    font-size: 0.95rem;
}

.title {
    color: #63B3ED;
    font-size: 2rem;
    font-weight: bold;
    text-align: center;
    margin-bottom: 1.5rem;
    padding: 1rem;
}

/* Custom scrollbar */
::-webkit-scrollbar {
    width: 8px;
}

::-webkit-scrollbar-track {
    background: rgba(255, 255, 255, 0.1);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb {
    background: linear-gradient(45deg, #3498db, #2ecc71);
    border-radius: 4px;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML('<div class="title">Chat Visualization</div>')

    with gr.Row(elem_classes=["container"]):
        chat_display = gr.HTML(elem_classes=["chat-panel"])
        metrics_display = gr.HTML(elem_classes=["metrics-panel"])

    # Show initial data on load
    demo.load(fn=display_chat, inputs=None, outputs=[chat_display, metrics_display])

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