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 OKR UI generator functions
from ui.okr_ui_generator import create_enhanced_okr_tab, format_okrs_for_enhanced_display, get_empty_okr_state

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
    )
    # The functions below are no longer needed due to the new OKR display approach and are removed from insights_ui_generator.py
    # from ui.insights_ui_generator import extract_key_results_for_selection, format_single_okr_for_display
    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 length (now 8 values)
        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>
        """
        # Placeholder should return sensible defaults for all outputs
        return (
            gr.update(value="Modules not loaded."), # 0: agentic_pipeline_status_md
            gr.update(choices=[], value=None),       # 1: report_selector_dd
            None, # 2: orchestration_raw_results_st - Placeholder should return None or an empty dict if it's the source for OKR display
            [],                                      # 3: selected_key_result_ids_st (internal state, not UI)
            [],                                      # 4: key_results_for_selection_st (internal state, not UI)
            gr.update(value=empty_header_html),      # 5: report_header_html_display
            gr.update(value=empty_body_markdown),    # 6: report_body_markdown_display
            {}                                       # 7: reconstruction_cache_st
        )
    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.'}


# --- 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)
    # The following states are no longer directly used for UI display, but might be kept for
    # internal logic if needed by other parts of the agentic pipeline that are not UI-related.
    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'])
        )

    # NEW: Function to update the enhanced OKR display
    def update_okr_tab_display(raw_results: dict):
        """
        Updates the enhanced OKR HTML display based on the raw agentic results.
        This replaces the old update_okr_display_on_selection which used a checkbox group.
        """
        return gr.update(value=format_okrs_for_enhanced_display(raw_results))


    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):
            # NEW: Call the new function to create the enhanced OKR tab UI
            okr_display_html = create_enhanced_okr_tab()

            # The old UI components key_results_cbg and okr_detail_display_md are removed
            # No need for individual KR selection anymore, the whole structure is displayed.


    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
    # REMOVED: okr_display_html from this list, as it's updated separately
    agentic_display_outputs = [
        agentic_pipeline_status_md,       # 0: Status Markdown (hidden)
        report_selector_dd,               # 1: Dropdown for selecting reports
        orchestration_raw_results_st,     # 2: Raw results state (source for OKR display)
        selected_key_result_ids_st,       # 3: Selected KR IDs state (internal state)
        key_results_for_selection_st,     # 4: All KRs for selection state (internal state)
        report_header_html_display,       # 5: New HTML output for report header
        report_body_markdown_display,     # 6: New Markdown output for report body
        reconstruction_cache_st,          # 7: Reconstruction cache state
    ]

    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: Outputs list updated to remove okr_display_html
        outputs=agentic_display_outputs,
        show_progress="minimal"
    ).then(
        # NEW: Call the update function for the enhanced OKR tab display, now as the only source for okr_display_html
        fn=update_okr_tab_display,
        inputs=[orchestration_raw_results_st], # Use the raw results state
        outputs=[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)