import gradio as gr
from memory_manager import embed_and_store, retrieve_relevant
from agent_engine import step_turn
from analysis_tools import analyze_sentiment_topics, plot_participation, generate_knowledge_graph
from exporter import export_txt, export_json, export_pdf, send_webhook

# Default agent parameters
DEFAULT_PARAMS = {
    agent: {"creativity": 0.7, "criticality": 0.7}
    for agent in ["Initiator", "Responder", "Guardian", "Provocateur", "Cultural", "Judge"]
}

# Agent color mapping
COLOR_MAP = {
    "Initiator": "#e6f7ff",
    "Responder": "#f6ffed",
    "Guardian": "#fff7e6",
    "Provocateur": "#f9e6ff",
    "Cultural": "#e6ffed",
    "Judge": "#f0f0f0",
    "System": "#d9d9d9",
    "User": "#ffffff"
}


def main():
    with gr.Blocks(css="""
        .convo-container {
            max-height: 400px;
            overflow-y: auto;
            padding: 10px;
            border: 1px solid #ccc;
            border-radius: 8px;
            background: #fafafa;
        }
        .message-card {
            padding: 8px;
            margin-bottom: 6px;
            border-radius: 6px;
        }
        .agent-panel {
            border: 1px solid #ddd;
            padding: 6px;
            border-radius: 4px;
            margin: 4px;
        }
    """, title="Hexa-Agent Discussion System") as demo:
        # States
        conversation_state = gr.State([])
        turn_state = gr.State(0)
        topic_state = gr.State("")
        params_state = gr.State(DEFAULT_PARAMS)

        # Header
        gr.Markdown("# 🧠 Modular Multi-Agent Discussion Platform")

        # Controls
        with gr.Row():
            topic_input = gr.Textbox(label="Topic", placeholder="Enter topic...", value="Ethical AI")
            set_topic_btn = gr.Button("Set Topic")
            clear_btn = gr.Button("Clear Discussion")
        with gr.Row():
            step_btn = gr.Button("▶️ Next Turn")
        gr.Markdown("---")

        # Conversation display
        convo_display = gr.HTML("<div class='convo-container' id='convo'></div>", label="Conversation")

        # Agent Panels
        with gr.Accordion("Agent Panels", open=False):
            initiator_panel = gr.Textbox(label="Initiator Latest", interactive=False)
            responder_panel = gr.Textbox(label="Responder Latest", interactive=False)
            guardian_panel = gr.Textbox(label="Guardian Latest", interactive=False)
            provocateur_panel = gr.Textbox(label="Provocateur Latest", interactive=False)
            cultural_panel = gr.Textbox(label="Cultural Latest", interactive=False)
            judge_panel = gr.Textbox(label="Judge Latest", interactive=False)

        # Analysis Tab
        with gr.Tab("Analysis"):
            sentiment = gr.Textbox(label="Sentiment")
            topics = gr.Textbox(label="Key Topics")
            part_plot = gr.Image(label="Participation Chart")
            graph_plot = gr.Image(label="Knowledge Graph")
            analyze_btn = gr.Button("Run Analysis")
            graph_btn = gr.Button("Generate Graph")

        # Configuration Tab: Parameter Sliders
        with gr.Tab("Configuration"):
            sliders = {}
            for agent in ["Initiator", "Responder", "Guardian", "Provocateur", "Cultural", "Judge"]:
                with gr.Row():
                    sliders[f"{agent}_creativity"] = gr.Slider(0.0, 1.0, value=DEFAULT_PARAMS[agent]['creativity'], label=f"{agent} Creativity")
                    sliders[f"{agent}_criticality"] = gr.Slider(0.0, 1.0, value=DEFAULT_PARAMS[agent]['criticality'], label=f"{agent} Criticality")

        # Export Tab
        with gr.Tab("Export"):
            fmt = gr.Radio(choices=["txt","json","pdf"], label="Format", value="txt")
            export_btn = gr.Button("Export")
            export_out = gr.File(label="Download")
            webhook_url = gr.Textbox(label="Webhook URL")
            send_btn = gr.Button("Send to Webhook")
            send_status = gr.Textbox(label="Status")

        # Event handlers
        set_topic_btn.click(lambda t: ([], 0, t), inputs=[topic_input], outputs=[convo_display, turn_state, topic_state])
        clear_btn.click(lambda: ([], 0, ""), outputs=[convo_display, turn_state, topic_state])

        def on_step(convo, turn, topic, params, *slider_vals):
            # Update params from sliders
            agents = ["Initiator","Responder","Guardian","Provocateur","Cultural","Judge"]
            new_params = {}
            idx = 0
            for agent in agents:
                new_params[agent] = {
                    'creativity': slider_vals[idx],
                    'criticality': slider_vals[idx+1]
                }
                idx += 2
            params = new_params
            if turn == 0 and topic:
                convo = [{"agent":"System","text":f"Topic: {topic}"}]
            convo = step_turn(convo, turn, topic or topic_input.value, params)
            # Build HTML
            html = ''
            for msg in convo:
                color = COLOR_MAP.get(msg['agent'], '#ffffff')
                html += f"<div class='message-card' style='background:{color};'><b>{msg['agent']}:</b> {msg['text']}</div>"
            html += "<script>var c=document.getElementById('convo'); c.scrollTop=c.scrollHeight;</script>"
            # Update panels
            panels = []
            for agent in agents:
                panels.append(next((m['text'] for m in reversed(convo) if m['agent']==agent), ''))
            return (html, convo, turn+1) + tuple(panels) + (params,)

        # Connect step with sliders
        step_btn.click(
            on_step,
            inputs=[conversation_state, turn_state, topic_state, params_state] + list(sliders.values()),
            outputs=[convo_display, conversation_state, turn_state,
                     initiator_panel, responder_panel, guardian_panel,
                     provocateur_panel, cultural_panel, judge_panel,
                     params_state]
        )

        analyze_btn.click(
            lambda convo: (
                analyze_sentiment_topics(convo)['sentiment'],
                ", ".join(analyze_sentiment_topics(convo)['topics']),
                plot_participation(convo, 'participation.png')
            ),
            inputs=[conversation_state],
            outputs=[sentiment, topics, part_plot]
        )

        graph_btn.click(
            lambda convo: generate_knowledge_graph(convo, 'graph.png'),
            inputs=[conversation_state],
            outputs=[graph_plot]
        )

        export_btn.click(
            lambda fmt, convo, topic, turn: {
                'txt': export_txt(convo, topic, turn),
                'json': export_json(convo, topic, turn),
                'pdf': export_pdf(convo, topic, turn)
            }[fmt],
            inputs=[fmt, conversation_state, topic_state, turn_state],
            outputs=[export_out]
        )

        send_btn.click(send_webhook, inputs=[webhook_url, conversation_state, topic_state, turn_state], outputs=[send_status])

    demo.launch()


if __name__ == '__main__':
    main()