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 (
    # display_main_dashboard, # Removed: Dashboard content is now in app.py directly
    build_analytics_tab_plot_area,
    BOMB_ICON, EXPLORE_ICON, FORMULA_ICON, ACTIVE_ICON
)
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,
        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):
        return "Modules not loaded.", gr.update(), "Modules not loaded.", "Modules not loaded.", None, [], [], "Error", {}
    def fetch_and_reconstruct_data_from_bubble(*args, **kwargs):
        return None, {}
    def format_report_for_display(report_data):
        return "Agentic modules not loaded. Report display unavailable."
    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)
    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)
        # dashboard_content = display_main_dashboard(new_state) # Removed this line
        return status_msg, new_state # Removed dashboard_content from outputs

    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
    )
    
    # --- FIXED: New handler only updates the report display ---
    def update_report_display(selected_report_id: str, current_token_state: dict):
        """
        Updates only the report display markdown when a new report is selected.
        The OKR visualization remains unchanged as it's loaded initially.
        """
        if not selected_report_id:
            return gr.update(value="*Please select a report to view its details.*")
        
        agentic_df = current_token_state.get("bubble_agentic_analysis_data")
        if agentic_df is None or agentic_df.empty:
            return gr.update(value="*Analysis data not loaded or is empty.*")

        selected_report_series_df = agentic_df[agentic_df['_id'] == selected_report_id]
        if selected_report_series_df.empty:
            return gr.update(value=f"*Error: Report with ID '{selected_report_id}' not found.*")
        
        # Extract the report data and format it for display
        selected_report_series = selected_report_series_df.iloc[0]
        report_markdown = format_report_for_display(selected_report_series)
        
        return report_markdown

    with gr.Tabs() as tabs:
        # --- NEW HOME TAB ---
        with gr.TabItem("1️⃣ Home", id="tab_home"):
            gr.Markdown("""
                <div style="text-align: center; padding: 20px; background-color: #f0f8ff; border-radius: 10px; margin-bottom: 20px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
                    <h2 style="color: #2c3e50; margin-bottom: 15px;">Welcome to your LinkedIn Employer Brand Analytics Dashboard!</h2>
                    <p style="font-size: 1.1em; color: #34495e; line-height: 1.6;">
                        This powerful tool is designed to help you **measure and enhance your employer brand** on LinkedIn.
                        By leveraging comprehensive analytics, you can dive into your data to understand trends, track performance,
                        and gain actionable insights to improve your presence and attractiveness as an employer.
                    </p>
                    <p style="font-size: 1.0em; color: #555; margin-top: 15px;">
                        Explore the sections below to get a comprehensive overview of your LinkedIn presence and
                        unlock the full potential of your employer branding efforts.
                    </p>
                </div>
            """)

            with gr.Row():
                with gr.Column():
                    gr.Markdown("""
                        <div style="background-color: #e8f5e9; padding: 20px; border-radius: 8px; min-height: 180px; display: flex; flex-direction: column; justify-content: space-between; box-shadow: 0 2px 4px rgba(0,0,0,0.08);">
                            <h3 style="color: #2e7d32; margin-bottom: 10px;">📈 Graphs</h3>
                            <p style="color: #4caf50; flex-grow: 1;">
                                Dive into detailed visualizations of your LinkedIn data. This section provides dynamic charts and
                                interactive plots that help you understand trends and variations in posts, mentions, and
                                follower statistics over time. Identify patterns and make data-driven decisions.
                            </p>
                            <div style="text-align: center; margin-top: 15px;">
                                <img src="https://placehold.co/100x60/4CAF50/ffffff?text=Charts" alt="Graphs Icon" style="margin-bottom: 10px; border-radius: 5px;">
                            </div>
                        </div>
                    """)
                    btn_graphs = gr.Button("Go to Graphs", variant="primary", size="lg")
                with gr.Column():
                    gr.Markdown("""
                        <div style="background-color: #e3f2fd; padding: 20px; border-radius: 8px; min-height: 180px; display: flex; flex-direction: column; justify-content: space-between; box-shadow: 0 2px 4px rgba(0,0,0,0.08);">
                            <h3 style="color: #1976d2; margin-bottom: 10px;">📊 Reports</h3>
                            <p style="color: #2196f3; flex-grow: 1;">
                                Access comprehensive quarterly and weekly reports of your employer brand performance.
                                These pre-generated reports offer in-depth summaries and insights, providing a clear
                                snapshot of your progress over specific periods.
                            </p>
                            <div style="text-align: center; margin-top: 15px;">
                                <img src="https://placehold.co/100x60/2196F3/ffffff?text=Reports" alt="Reports Icon" style="margin-bottom: 10px; border-radius: 5px;">
                            </div>
                        </div>
                    """)
                    btn_reports = gr.Button("Go to Reports", variant="primary", size="lg")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("""
                        <div style="background-color: #fff3e0; padding: 20px; border-radius: 8px; min-height: 180px; display: flex; flex-direction: column; justify-content: space-between; box-shadow: 0 2px 4px rgba(0,0,0,0.08);">
                            <h3 style="color: #ef6c00; margin-bottom: 10px;">🎯 OKR Table</h3>
                            <p style="color: #ff9800; flex-grow: 1;">
                                Discover Objectives and Key Results (OKRs) generated by AI, along with actionable tasks.
                                This section provides concrete recommendations tailored to improve your employer brand,
                                helping you translate insights into measurable actions.
                            </p>
                            <div style="text-align: center; margin-top: 15px;">
                                <img src="https://placehold.co/100x60/FF9800/ffffff?text=OKRs" alt="OKR Icon" style="margin-bottom: 10px; border-radius: 5px;">
                            </div>
                        </div>
                    """)
                    btn_okr = gr.Button("Go to OKR Table", variant="primary", size="lg")

            # Link buttons to tab selection
            btn_graphs.click(fn=lambda: gr.update(selected="tab_analytics"), 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)


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

        with gr.TabItem("3️⃣ Agentic Analysis Report", id="tab_agentic_report", visible=AGENTIC_MODULES_LOADED):
            gr.Markdown("## 🤖 Comprehensive Analysis Report (from Bubble.io)")
            agentic_pipeline_status_md = gr.Markdown("Status: Loading report data...", visible=True)
            gr.Markdown("Questo report è stato pre-generato. Seleziona un report dalla libreria per visualizzarlo.")
            with gr.Row():
                report_selector_dd = gr.Dropdown(label="Report Library", choices=[], interactive=True, info="Select a report.")
            agentic_report_display_md = gr.Markdown("Please select a report from the library to view it.")
            if not AGENTIC_MODULES_LOADED:
                gr.Markdown("🔴 **Error:** Agentic modules could not be 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. Seleziona i Key Results per dettagli.")
            if not AGENTIC_MODULES_LOADED:
                gr.Markdown("🔴 **Error:** Agentic modules could not be loaded.")
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Suggested Key Results")
                    key_results_cbg = gr.CheckboxGroup(label="Select Key Results", choices=[], value=[], interactive=True)
                with gr.Column(scale=3):
                    gr.Markdown("### Detailed OKRs and Tasks")
                    okr_detail_display_md = gr.Markdown("I dettagli OKR appariranno qui.")

            def update_okr_display_on_selection(selected_kr_ids: list, raw_results: dict, all_krs: list):
                if not raw_results or not AGENTIC_MODULES_LOADED:
                    return gr.update(value="Nessun dato di analisi caricato.")
                actionable_okrs = raw_results.get("actionable_okrs")
                if not actionable_okrs or not isinstance(actionable_okrs.get("okrs"), list):
                    return gr.update(value="Nessun OKR trovato.")
                okrs_list, kr_id_map = actionable_okrs["okrs"], {kr['unique_kr_id']: (kr['okr_index'], kr['kr_index']) for kr in all_krs}
                selected_krs_by_okr_idx = defaultdict(list)
                if selected_kr_ids:
                    for kr_id in selected_kr_ids:
                        if kr_id in kr_id_map:
                            okr_idx, kr_idx = kr_id_map[kr_id]
                            selected_krs_by_okr_idx[okr_idx].append(kr_idx)
                output_parts = []
                for okr_idx, okr in enumerate(okrs_list):
                    if not selected_kr_ids:
                        output_parts.append(format_single_okr_for_display(okr, okr_main_index=okr_idx))
                    elif okr_idx in selected_krs_by_okr_idx:
                        accepted_indices = selected_krs_by_okr_idx.get(okr_idx)
                        output_parts.append(format_single_okr_for_display(okr, accepted_kr_indices=accepted_indices, okr_main_index=okr_idx))
                final_md = "\n\n---\n\n".join(output_parts) if output_parts else "Nessun OKR corrisponde alla selezione."
                return gr.update(value=final_md)

            if AGENTIC_MODULES_LOADED:
                key_results_cbg.change(
                    fn=update_okr_display_on_selection,
                    inputs=[key_results_cbg, orchestration_raw_results_st, key_results_for_selection_st],
                    outputs=[okr_detail_display_md]
                )
    
    if AGENTIC_MODULES_LOADED:
        # FIXED: The change event for the report selector now only updates the report display markdown.
        # The OKR visualization is not affected and remains static after the initial load.
        report_selector_dd.change(
            fn=update_report_display,
            inputs=[report_selector_dd, token_state],
            outputs=[agentic_report_display_md],
            show_progress="minimal"
        )

    agentic_display_outputs = [
        agentic_report_display_md, report_selector_dd, key_results_cbg,
        okr_detail_display_md, orchestration_raw_results_st, selected_key_result_ids_st,
        key_results_for_selection_st, agentic_pipeline_status_md, reconstruction_cache_st
    ]

    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], # dashboard_display_html removed
        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],
        outputs=agentic_display_outputs,
        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)