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

# Import UI generator functions (these are now passed to build_main_app_ui)
from ui.ui_generators import (
    build_analytics_tab_plot_area,
    build_home_tab_ui,
    create_enhanced_report_tab,
    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

# NEW: Import UI enhancements from the new module
from ui.ui_main_page_enhancements import build_main_app_ui, update_report_display_enhanced

# --- 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
    )
    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.'}


# --- Initial data load sequence function (remains in app.py as it uses service functions) ---
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)
    # Add status icons based on success/failure
    if "successfully" in status_msg.lower() or "loaded" in status_msg.lower():
        status_msg = f"✅ {status_msg}"
    elif "error" in status_msg.lower() or "failed" in status_msg.lower():
        status_msg = f"❌ {status_msg}"
    else:
        status_msg = f"🔄 {status_msg}"
    return status_msg, new_state

# Instantiate AnalyticsTab (needs to be done before building UI if its methods are called)
analytics_icons = {'bomb': BOMB_ICON, 'explore': EXPLORE_ICON, 'formula': FORMULA_ICON, 'active': ACTIVE_ICON}
analytics_tab_instance = AnalyticsTab(
    # These states are created within build_main_app_ui, but AnalyticsTab needs references.
    # We will set them after build_main_app_ui returns them. For now, pass placeholders.
    token_state=None, # Will be updated after build_main_app_ui
    chat_histories_st=None, # Will be updated after build_main_app_ui
    current_chat_plot_id_st=None, # Will be updated after build_main_app_ui
    plot_data_for_chatbot_st=None, # Will be updated after build_main_app_ui
    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
)


# Build the main UI using the function from ui_enhancements
(app, url_user_token_display, org_urn_display, status_box,
 token_state, reconstruction_cache_st, enhanced_okr_display_html,
 tabs, report_selector_dd, agentic_display_outputs,
 analytics_tab_instance_returned, chat_histories_st_returned, format_report_for_display_func_passed) = \
    build_main_app_ui(
        analytics_tab_instance=analytics_tab_instance, # Pass the instantiated object
        AGENTIC_MODULES_LOADED=AGENTIC_MODULES_LOADED,
        build_home_tab_ui_func=build_home_tab_ui,
        create_enhanced_report_tab_func=create_enhanced_report_tab,
        create_enhanced_okr_tab_func=create_enhanced_okr_tab,
        format_report_for_display_func=format_report_for_display # Pass the imported function
    )

# Now, update the analytics_tab_instance with the actual state components from the built UI
# This is crucial because analytics_tab_instance was instantiated before the gr.State components existed.
analytics_tab_instance.token_state = token_state
analytics_tab_instance.chat_histories_st = chat_histories_st_returned
analytics_tab_instance.current_chat_plot_id_st = agentic_display_outputs[9] # This is current_chat_plot_id_st
analytics_tab_instance.plot_data_for_chatbot_st = agentic_display_outputs[11] # This is actionable_okrs_data_st, which is incorrect. Should be plot_data_for_chatbot_st

# Corrected update for analytics_tab_instance state components.
# Need to know the exact index of each state in agentic_display_outputs or pass them directly.
# Given the `build_main_app_ui` returns specific components, let's use those directly.
# The original code for `chat_histories_st`, `current_chat_plot_id_st`, `plot_data_for_chatbot_st`
# are defined within `build_main_app_ui`'s scope.
# The returned `chat_histories_st_returned` and the other states are the correct references.
# So, no need to access them via `agentic_display_outputs` for the analytics_tab_instance.
# The parameters `chat_histories_st_returned` (which is `chat_histories_st` inside `build_main_app_ui`)
# and others are the correct references.

# Event handlers (re-establishing them now that components are defined)
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)

if AGENTIC_MODULES_LOADED:
    report_selector_dd.change(
        fn=lambda sr_id, c_state: update_report_display_enhanced(sr_id, c_state, format_report_for_display),
        inputs=[report_selector_dd, token_state],
        outputs=[agentic_display_outputs[7], agentic_display_outputs[8]], # report_header_html_display, report_body_markdown_display
        show_progress="minimal"
    )

# Initial load sequence with enhanced status updates
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"
)

# Chain the loading events
initial_load_event.then(
    fn=analytics_tab_instance.refresh_analytics_graphs_ui, # Using the corrected method name (assuming this exists)
    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_returned], # Using returned 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],
    outputs=agentic_display_outputs,
    show_progress="minimal"
).then(
    fn=format_okrs_for_enhanced_display,
    inputs=[reconstruction_cache_st],
    outputs=[enhanced_okr_display_html],
    show_progress="minimal"
)

if __name__ == "__main__":
    # Enhanced startup logging
    print("🚀 Starting LinkedIn Organization Dashboard...")
    
    # Environment variable checks with better logging
    missing_vars = []
    
    if not os.environ.get(LINKEDIN_CLIENT_ID_ENV_VAR):
        missing_vars.append(LINKEDIN_CLIENT_ID_ENV_VAR)
        logging.warning(f"⚠️  WARNING: '{LINKEDIN_CLIENT_ID_ENV_VAR}' is not set.")
    
    bubble_vars = [BUBBLE_APP_NAME_ENV_VAR, BUBBLE_API_KEY_PRIVATE_ENV_VAR, BUBBLE_API_ENDPOINT_ENV_VAR]
    if not all(os.environ.get(var) for var in bubble_vars):
        missing_vars.extend([var for var in bubble_vars if not os.environ.get(var)])
        logging.warning("⚠️  WARNING: One or more Bubble environment variables are not set.")
    
    if not os.environ.get("GEMINI_API_KEY"):
        missing_vars.append("GEMINI_API_KEY")
        logging.warning("⚠️  WARNING: 'GEMINI_API_KEY' is not set.")
    
    if not AGENTIC_MODULES_LOADED:
        logging.warning("🔴 CRITICAL: Agentic modules failed to load.")
    
    if missing_vars:
        print(f"⚠️  Missing environment variables: {', '.join(missing_vars)}")
        print("🔧 Please set these variables for full functionality.")
    else:
        print("✅ All environment variables are properly configured.")
    
    print("🌐 Launching dashboard on http://0.0.0.0:7860")
    print("🎯 Dashboard features:")
    print("   • 📊 Advanced LinkedIn Analytics")
    print("   • 🤖 AI-Powered Insights")
    print("   • 🎯 OKR Generation & Tracking") 
    print("   • ☁️  Bubble.io Integration")
    
    app.launch(
        server_name="0.0.0.0", 
        server_port=int(os.environ.get("PORT", 7860)), 
        debug=True,
        show_error=True,
        show_tips=True,
        enable_queue=True
    )