app.py

import gradio as gr
import pandas as pd
import os
import logging
from collections import defaultdict
import matplotlib
matplotlib.use('Agg') # Set backend for Matplotlib

# --- Module Imports ---
from utils.gradio_utils import get_url_user_token

# Functions from newly created/refactored modules
from config import (
    PLOT_ID_TO_FORMULA_KEY_MAP,
    LINKEDIN_CLIENT_ID_ENV_VAR,
    BUBBLE_APP_NAME_ENV_VAR,
    BUBBLE_API_KEY_PRIVATE_ENV_VAR,
    BUBBLE_API_ENDPOINT_ENV_VAR
)

from services.analytics_tab_module import AnalyticsTab

# UPDATED: Using the new data loading function from the refactored state manager
from services.state_manager import load_data_from_bubble
from ui.ui_generators import (
    build_analytics_tab_plot_area,
    build_home_tab_ui, # NEW: Import the function to build the Home tab UI
    create_enhanced_report_tab, # NEW: Import the function to build the enhanced Report tab UI
    BOMB_ICON, EXPLORE_ICON, FORMULA_ICON, ACTIVE_ICON
)

#ui import
from ui.config import custom_title_css

# NEW: Import the new OKR UI functions
from ui.okr_ui_generator import create_enhanced_okr_tab, format_okrs_for_enhanced_display, get_initial_okr_display
from ui.analytics_plot_generator import update_analytics_plots_figures, create_placeholder_plot
from formulas import PLOT_FORMULAS

# --- CHATBOT MODULE IMPORTS ---
from features.chatbot.chatbot_prompts import get_initial_insight_prompt_and_suggestions
from features.chatbot.chatbot_handler import generate_llm_response

# --- AGENTIC PIPELINE (DISPLAY ONLY) IMPORTS ---
try:
    # This is the main function called on initial load to populate the agentic tabs
    from run_agentic_pipeline import load_and_display_agentic_results
    # This function is now called when a new report is selected from the dropdown
    from services.report_data_handler import fetch_and_reconstruct_data_from_bubble
    # UI formatting functions
    from ui.insights_ui_generator import (
        format_report_for_display, # This will now return header HTML and body Markdown
        # REMOVED: extract_key_results_for_selection, - Moved to okr_ui_generator (implicitly)
        # REMOVED: format_single_okr_for_display - Moved to okr_ui_generator (implicitly)
    )
    AGENTIC_MODULES_LOADED = True
except ImportError as e:
    logging.error(f"Could not import agentic pipeline display modules: {e}. Tabs 3 and 4 will be disabled.")
    AGENTIC_MODULES_LOADED = False
    # Placeholder functions to prevent app from crashing if imports fail
    def load_and_display_agentic_results(*args, **kwargs):
        # NOTE: This return signature MUST match agentic_display_outputs
        # Adjusted return values for the new split report display components and the new OKR HTML
        empty_header_html = """
        <div class="report-title">📝 Comprehensive Analysis Report</div>
        <div class="report-subtitle">AI-Generated Insights from Your LinkedIn Data</div>
        <div class="status-badge">Generated from Bubble.io</div>
        """
        empty_body_markdown = """
        <div class="empty-state">
            <div class="empty-state-icon">📄</div>
            <div class="empty-state-title">No Report Selected</div>
            <div class="empty-state-description">
                Please select a report from the library above to view its detailed analysis and insights.
            </div>
        </div>
        """
        # The load_and_display_agentic_results function returns many values.
        # Ensure the placeholder returns the correct number of gr.update components
        # matching the `outputs` in the .then() call later.
        return (
            gr.update(value="Modules not loaded."), # agentic_pipeline_status_md (0)
            gr.update(choices=[], value=None),       # report_selector_dd (1)
            gr.update(choices=[], value=[]),         # key_results_cbg (2) - KEPT HIDDEN for compatibility
            gr.update(value="Modules not loaded."), # okr_detail_display_md (3) - KEPT HIDDEN for compatibility
            None,                                    # orchestration_raw_results_st (4)
            [],                                      # selected_key_result_ids_st (5) - KEPT HIDDEN for compatibility
            [],                                      # key_results_for_selection_st (6) - KEPT HIDDEN for compatibility
            gr.update(value=empty_header_html),      # report_header_html_display (7)
            gr.update(value=empty_body_markdown),   # report_body_markdown_display (8)
            {},                                      # reconstruction_cache_st (9)
            gr.update(value=get_initial_okr_display()), # NEW: enhanced_okr_display_html (10)
            gr.update(value={}) # NEW: actionable_okrs_data_st (11)
        )
    def fetch_and_reconstruct_data_from_bubble(*args, **kwargs):
        return None, {}
    def format_report_for_display(report_data):
        # Placeholder for when modules are not loaded, returns structure matching the new design
        return {'header_html': '<h1>Agentic modules not loaded.</h1>', 'body_markdown': 'Report display unavailable.'}


with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
              title="LinkedIn Organization Dashboard") as app:
    # --- STATE MANAGEMENT ---
    token_state = gr.State(value={
        "token": None, "client_id": None, "org_urn": None,
        "bubble_posts_df": pd.DataFrame(), "bubble_post_stats_df": pd.DataFrame(),
        "bubble_mentions_df": pd.DataFrame(), "bubble_follower_stats_df": pd.DataFrame(),
        "bubble_agentic_analysis_data": pd.DataFrame(), # To store agentic results from Bubble
        "url_user_token_temp_storage": None,
        "config_date_col_posts": "published_at", "config_date_col_mentions": "date",
        "config_date_col_followers": "date", "config_media_type_col": "media_type",
        "config_eb_labels_col": "li_eb_label"
    })

    # States for analytics tab chatbot
    chat_histories_st = gr.State({})
    current_chat_plot_id_st = gr.State(None)
    plot_data_for_chatbot_st = gr.State({})

    # States for agentic results display
    orchestration_raw_results_st = gr.State(None)
    # KEPT for compatibility with load_and_display_agentic_results signature
    key_results_for_selection_st = gr.State([])
    selected_key_result_ids_st = gr.State([])

    # --- NEW: Session-specific cache for reconstructed OKR data ---
    reconstruction_cache_st = gr.State({})
    # NEW: State to hold the actionable_okrs dictionary explicitly
    actionable_okrs_data_st = gr.State({})


    # --- UI LAYOUT ---
    # --- ENHANCED UI HEADER ---
    with gr.Row():
        with gr.Column():
            gr.HTML("""
                <div style="
                    background: linear-gradient(135deg, #f8fafc 0%, #f1f5f9 100%);
                    border-radius: 15px;
                    padding: 40px 20px;
                    margin: 20px 0;
                    box-shadow: 0 8px 25px rgba(0, 0, 0, 0.15);
                    border: 2px solid #e2e8f0;
                    text-align: center;
                ">
                    <div style="
                        font-size: 48px;
                        font-weight: 800;
                        color: #0077B5;
                        margin: 0 0 15px 0;
                        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
                    ">
                        <span style="font-size: 40px; margin-right: 10px;">🚀</span>LinkedIn Organization Dashboard
                    </div>
                    <div style="
                        font-size: 18px;
                        color: #666666;
                        margin: 0;
                        font-weight: 400;
                        font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                        line-height: 1.4;
                    ">
                        Advanced Analytics & AI-Powered Insights for Your LinkedIn Performance
                    </div>
                </div>
            """)

    url_user_token_display = gr.Textbox(label="User Token (Hidden)", interactive=False, visible=False)
    org_urn_display = gr.Textbox(label="Org URN (Hidden)", interactive=False, visible=False)
    # MODIFIED: Initially hide the status_box and make its value empty
    status_box = gr.Textbox(label="Status", interactive=False, value="", visible=False)

    app.load(fn=get_url_user_token, inputs=None, outputs=[url_user_token_display, org_urn_display], api_name="get_url_params", show_progress=False)

    def initial_data_load_sequence(url_token, org_urn_val, current_state):
        """
        Handles the initial data loading from Bubble and controls status_box visibility.
        """
        status_msg, new_state = load_data_from_bubble(url_token, org_urn_val, current_state)

        # Determine if the status box should be visible
        # It should be visible if the token is not available (None)
        # OR if the status message indicates an error or warning, even if a token *might* be present.
        show_status_box = False
        if new_state.get("token") is None:
            show_status_box = True
        elif "Error" in status_msg or "Warning" in status_msg:
            show_status_box = True
        
        # Return the gr.update object to control value and visibility
        return gr.update(value=status_msg, visible=show_status_box), new_state

    analytics_icons = {'bomb': BOMB_ICON, 'explore': EXPLORE_ICON, 'formula': FORMULA_ICON, 'active': ACTIVE_ICON}
    analytics_tab_instance = AnalyticsTab(
        token_state=token_state,
        chat_histories_st=chat_histories_st,
        current_chat_plot_id_st=current_chat_plot_id_st,
        plot_data_for_chatbot_st=plot_data_for_chatbot_st,
        plot_id_to_formula_map=PLOT_ID_TO_FORMULA_KEY_MAP,
        plot_formulas_data=PLOT_FORMULAS,
        icons=analytics_icons,
        fn_build_plot_area=build_analytics_tab_plot_area,
        fn_update_plot_figures=update_analytics_plots_figures,
        fn_create_placeholder_plot=create_placeholder_plot,
        fn_get_initial_insight=get_initial_insight_prompt_and_suggestions,
        fn_generate_llm_response=generate_llm_response
    )

    def update_report_display(selected_report_id: str, current_token_state: dict):
        """
        Updates the report header and body display when a new report is selected.
        This function now expects format_report_for_display to return a dict with
        'header_html' and 'body_markdown'.
        """
        # Define empty states for header and body
        empty_header_html = """
        <div class="report-title">📝 Comprehensive Analysis Report</div>
        <div class="report-subtitle">AI-Generated Insights from Your LinkedIn Data</div>
        <div class="status-badge">Generated from Bubble.io</div>
        """
        empty_body_markdown_no_selection = """
        <div class="empty-state">
            <div class="empty-state-icon">📋</div>
            <div class="empty-state-title">Select a Report</div>
            <div class="empty-state-description">
                Choose a report from the dropdown above to view its detailed analysis and insights.
            </div>
        </div>
        """
        empty_body_markdown_no_data = """
        <div class="empty-state">
            <div class="empty-state-icon">⚠️</div>
            <div class="empty-state-title">Data Not Available</div>
            <div class="empty-state-description">
                Analysis data is not loaded or is empty. Please try refreshing the page.
            </div>
        </div>
        """
        empty_body_markdown_not_found = lambda _id: f"""
        <div class="empty-state">
            <div class="empty-state-icon">❌</div>
            <div class="empty-state-title">Report Not Found</div>
            <div class="empty-state-description">
                Report with ID '{_id}' was not found in the database.
            </div>
        </div>
        """

        if not selected_report_id:
            # When no report is selected, update both header and body
            return gr.update(value=empty_header_html), gr.update(value=empty_body_markdown_no_selection)

        agentic_df = current_token_state.get("bubble_agentic_analysis_data")
        if agentic_df is None or agentic_df.empty:
            # When no data is available, update both header and body
            return gr.update(value=empty_header_html), gr.update(value=empty_body_markdown_no_data)

        selected_report_series_df = agentic_df[agentic_df['_id'] == selected_report_id]
        if selected_report_series_df.empty:
            # When report is not found, update both header and body
            return gr.update(value=empty_header_html), gr.update(value=empty_body_markdown_not_found(selected_report_id))

        selected_report_series = selected_report_series_df.iloc[0]

        # Call the format_report_for_display, which now returns a dict
        formatted_content_parts = format_report_for_display(selected_report_series)

        # Update the two separate Gradio components
        return (
            gr.update(value=formatted_content_parts['header_html']),
            gr.update(value=formatted_content_parts['body_markdown'])
        )


    with gr.Tabs() as tabs:
        # --- NEW HOME TAB ---
        with gr.TabItem("🏠 Home", id="tab_home"):

            # Call the new function from ui_generators to build the Home tab content
            btn_graphs, btn_reports, btn_okr, btn_help = build_home_tab_ui()

            # Link buttons to tab selection
            btn_graphs.click(fn=lambda: gr.update(selected="tab_analytics_module"), outputs=tabs)
            btn_reports.click(fn=lambda: gr.update(selected="tab_agentic_report"), outputs=tabs)
            btn_okr.click(fn=lambda: gr.update(selected="tab_agentic_okrs"), outputs=tabs)
            # btn_help.click(fn=lambda: gr.update(selected="tab_help"), outputs=tabs) # Uncomment if you add a help tab


        analytics_tab_instance.create_tab_ui() # This is the "Graphs" tab, assuming its ID is "tab_analytics"

        # --- REPLACED: Agentic Analysis Report Tab with enhanced UI ---
        # The create_enhanced_report_tab function now builds this entire tab's UI.
        # It also returns the relevant Gradio components needed for callbacks.
        with gr.TabItem("📝 AI Analysis Reports", id="tab_agentic_report", visible=AGENTIC_MODULES_LOADED):

            # The create_enhanced_report_tab function handles the CSS and HTML structure
            # MODIFIED: Unpacked 4 values instead of 3
            agentic_pipeline_status_md, report_selector_dd, report_header_html_display, report_body_markdown_display = \
                create_enhanced_report_tab(AGENTIC_MODULES_LOADED)


        with gr.TabItem("🎯 OKRs & Action Items", id="tab_agentic_okrs", visible=AGENTIC_MODULES_LOADED):

            gr.Markdown("## 🎯 AI Generated OKRs and Actionable Tasks (from Bubble.io)")
            gr.Markdown("Basato sull'analisi AI, l'agente ha proposto i seguenti OKR.")

            if not AGENTIC_MODULES_LOADED:
                gr.Markdown("🔴 **Error:** Agentic modules could not be loaded.")

            # Keep the old components but make them invisible to maintain load_and_display_agentic_results signature
            with gr.Column(visible=False):
                gr.Markdown("### Suggested Key Results (OLD UI - HIDDEN)")
                key_results_cbg = gr.CheckboxGroup(label="Select Key Results", choices=[], value=[], interactive=True)
                gr.Markdown("### Detailed OKRs and Tasks (OLD UI - HIDDEN)")
                okr_detail_display_md = gr.Markdown("I dettagli OKR appariranno qui.")

            # NEW: Add the enhanced OKR display HTML component
            enhanced_okr_display_html = create_enhanced_okr_tab()

            # REMOVED: The old update_okr_display_on_selection function and its change event
            # as the new UI handles display dynamically from raw_results_st

    if AGENTIC_MODULES_LOADED:
        report_selector_dd.change(
            fn=update_report_display, # This now calls the enhanced function
            # MODIFIED: Updated outputs to match the two new display components
            inputs=[report_selector_dd, token_state],
            outputs=[report_header_html_display, report_body_markdown_display],
            show_progress="minimal"
        )

    # Ensure agentic_display_outputs correctly maps to the newly created components
    # This list must match the outputs of load_and_display_agentic_results
    agentic_display_outputs = [
        agentic_pipeline_status_md,            # 0: Status Markdown (hidden)
        report_selector_dd,                    # 1: Dropdown for selecting reports
        key_results_cbg,                       # 2: Checkbox group for OKRs (kept hidden)
        okr_detail_display_md,                 # 3: Markdown for detailed OKR display (kept hidden)
        orchestration_raw_results_st,          # 4: Raw results state
        selected_key_result_ids_st,            # 5: Selected KR IDs state (kept hidden)
        key_results_for_selection_st,          # 6: All KRs for selection state (kept hidden)
        report_header_html_display,            # 7: New HTML output for header
        report_body_markdown_display,          # 8: New Markdown output for body
        reconstruction_cache_st,               # 9: Reconstruction cache state
        enhanced_okr_display_html,             # 10: NEW: The enhanced HTML display for OKRs
        actionable_okrs_data_st                # 11: NEW: The actionable_okrs dictionary state
    ]

    # MODIFIED: Chain the initial_load_event to first show "Loading...", then perform data load.
    initial_load_event = org_urn_display.change(
        fn=lambda: gr.update(value="Loading data...", visible=True), # Show loading message immediately
        inputs=[],
        outputs=[status_box],
        show_progress="full"
    ).then(
        fn=initial_data_load_sequence,
        inputs=[url_user_token_display, org_urn_display, token_state],
        outputs=[status_box, token_state], # status_box is updated here with final message/visibility
        show_progress="full"
    )

    initial_load_event.then(
        fn=analytics_tab_instance._refresh_analytics_graphs_ui,
        inputs=[token_state, analytics_tab_instance.date_filter_selector, analytics_tab_instance.custom_start_date_picker,
                analytics_tab_instance.custom_end_date_picker, chat_histories_st],
        outputs=analytics_tab_instance.graph_refresh_outputs_list,
        show_progress="full"
    ).then(
        fn=load_and_display_agentic_results,
        inputs=[token_state, reconstruction_cache_st],
        # MODIFIED: Updated outputs to match all components returned by load_and_display_agentic_results (now 12)
        outputs=agentic_display_outputs,
        show_progress="minimal"
    ).then( # NEW CHAIN: Update the enhanced OKR display after load_and_display_agentic_results runs
    fn=format_okrs_for_enhanced_display,
    inputs=[reconstruction_cache_st], # Change from actionable_okrs_data_st to reconstruction_cache_st
    outputs=[enhanced_okr_display_html],
    show_progress="minimal"
    )


if __name__ == "__main__":
    if not os.environ.get(LINKEDIN_CLIENT_ID_ENV_VAR):
        logging.warning(f"WARNING: '{LINKEDIN_CLIENT_ID_ENV_VAR}' is not set.")
    if not all(os.environ.get(var) for var in [BUBBLE_APP_NAME_ENV_VAR, BUBBLE_API_KEY_PRIVATE_ENV_VAR, BUBBLE_API_ENDPOINT_ENV_VAR]):
        logging.warning("WARNING: One or more Bubble environment variables are not set.")
    if not AGENTIC_MODULES_LOADED:
        logging.warning("CRITICAL: Agentic modules failed to load.")
    if not os.environ.get("GEMINI_API_KEY"):
        logging.warning("WARNING: 'GEMINI_API_KEY' is not set.")

    app.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), debug=True)