"""
Fintech Multi-Agent Orchestrator - HuggingFace Spaces Demo
Uses Gemma 3 via HuggingFace Inference API
"""

import gradio as gr
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from io import BytesIO
import base64
from PIL import Image

from graph import run_orchestrator


def create_chart(chart_data: dict) -> Image.Image:
    """Generate chart from chart_data dict."""
    chart_type = chart_data.get("type", "bar")
    title = chart_data.get("title", "Chart")
    data = chart_data.get("data", {})
    
    fig, ax = plt.subplots(figsize=(10, 6))
    
    if chart_type == "pie":
        labels = list(data.keys())
        values = list(data.values())
        colors = plt.cm.Set3(np.linspace(0, 1, len(labels)))
        ax.pie(values, labels=labels, autopct='%1.1f%%', colors=colors, startangle=90)
        ax.set_title(title, fontsize=14, fontweight='bold')
    
    elif chart_type == "bar":
        labels = list(data.keys())
        values = list(data.values())
        colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(labels)))
        bars = ax.bar(labels, values, color=colors)
        ax.set_title(title, fontsize=14, fontweight='bold')
        ax.set_ylabel('Amount ($)')
        plt.xticks(rotation=45, ha='right')
        for bar, val in zip(bars, values):
            ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(values)*0.02, 
                   f'${val:,.0f}', ha='center', va='bottom', fontsize=9)
    
    elif chart_type == "line":
        x = list(range(len(data)))
        y = list(data.values())
        ax.plot(x, y, marker='o', linewidth=2, markersize=8, color='#2E86AB')
        ax.fill_between(x, y, alpha=0.3, color='#2E86AB')
        ax.set_xticks(x)
        ax.set_xticklabels(list(data.keys()))
        ax.set_title(title, fontsize=14, fontweight='bold')
        ax.set_ylabel('Amount ($)')
        ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'${x:,.0f}'))
    
    plt.tight_layout()
    
    # Convert to PIL Image
    buf = BytesIO()
    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
    buf.seek(0)
    plt.close()
    
    return Image.open(buf)


def respond(query: str, history: list):
    """Main chat handler."""
    if not query.strip():
        return history, None
    
    # Run orchestrator
    response, chart_data = run_orchestrator(query)
    
    # Create chart if needed
    chart_image = None
    if chart_data:
        chart_image = create_chart(chart_data)
    
    # Update history (Gradio 5.x messages format)
    history = history or []
    history.append({"role": "user", "content": query})
    history.append({"role": "assistant", "content": response})
    
    return history, chart_image


# Gradio Interface
with gr.Blocks(
    title="Fintech Multi-Agent Orchestrator",
) as demo:
    
    gr.Markdown("""
    # 🏦 Fintech Multi-Agent Orchestrator
    
    **Powered by Gemma 3 via HuggingFace Inference API**
    
    Ask questions about your finances, request calculations, or generate charts!
    
    ### Example queries:
    - "What is my net worth?"
    - "Show my portfolio as a pie chart"
    - "Calculate compound interest on $10000 at 8% for 5 years"
    - "Show my assets breakdown"
    - "Show net worth projection"
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Chat",
                height=400,
                type="messages"
            )
            
            with gr.Row():
                query_input = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask about your finances...",
                    scale=4,
                )
                submit_btn = gr.Button("Send", variant="primary", scale=1)
        
        with gr.Column(scale=1):
            chart_output = gr.Image(
                label="Generated Chart",
                height=400,
            )
    
    with gr.Row():
        clear_btn = gr.Button("Clear Chat")
    
    submit_btn.click(
        respond,
        inputs=[query_input, chatbot],
        outputs=[chatbot, chart_output],
    ).then(
        lambda: "",
        outputs=[query_input],
    )
    
    query_input.submit(
        respond,
        inputs=[query_input, chatbot],
        outputs=[chatbot, chart_output],
    ).then(
        lambda: "",
        outputs=[query_input],
    )
    
    clear_btn.click(
        lambda: ([], None),
        outputs=[chatbot, chart_output],
    )


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