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
)
# 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
)
# 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)
        )
    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.'}
    # REMOVED from insights_ui_generator.py, so also remove placeholder if not needed by other direct calls
    # def extract_key_results_for_selection(okr_data):
    #     return []
    # def format_single_okr_for_display(okr_data, **kwargs):
    #     return "Agentic modules not loaded. OKR display unavailable."


# --- ANALYTICS TAB MODULE IMPORT ---
from services.analytics_tab_module import AnalyticsTab

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(message)s')

# API Key Setup
user_provided_api_key = os.environ.get("GEMINI_API_KEY")
if user_provided_api_key:
    os.environ["GOOGLE_API_KEY"] = user_provided_api_key
    logging.info("GOOGLE_API_KEY environment variable has been set from GEMINI_API_KEY.")
else:
    logging.error("CRITICAL ERROR: The API key environment variable 'GEMINI_API_KEY' was not found.")


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({})


    # --- UI LAYOUT ---
    gr.Markdown("# 🚀 LinkedIn Organization Dashboard")
    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)
    status_box = gr.Textbox(label="Status", interactive=False, value="Initializing...")

    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.
        No longer generates dashboard HTML as the Home tab is now static.
        """
        status_msg, new_state = load_data_from_bubble(url_token, org_urn_val, current_state)
        return status_msg, 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("1️⃣ 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("3️⃣ Agentic Analysis Report", 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("4️⃣ Agentic OKRs & Tasks", 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
    ]

    initial_load_event = org_urn_display.change(
        fn=initial_data_load_sequence,
        inputs=[url_user_token_display, org_urn_display, token_state],
        outputs=[status_box, token_state],
        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
        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=[orchestration_raw_results_st], # Take the raw results as input
        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)