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,
    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):
        status_msg, new_state = load_data_from_bubble(url_token, org_urn_val, current_state)
        dashboard_content = display_main_dashboard(new_state)
        return status_msg, new_state, dashboard_content

    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
    )
    
    # --- MODIFIED: Handler now uses the session cache ---
    def update_report_and_okr_display(selected_report_id: str, current_token_state: dict, session_cache: dict):
        error_return_tuple = (
            gr.update(value="*Please select a report to view its details.*"),
            gr.update(choices=[], value=[], interactive=False),
            gr.update(value="*Please select a report to see OKRs.*"),
            None, [], [], session_cache # Pass cache back unchanged
        )
        if not selected_report_id:
            return error_return_tuple
        
        agentic_df = current_token_state.get("bubble_agentic_analysis_data")
        if agentic_df is None or agentic_df.empty:
            return error_return_tuple

        selected_report_series_df = agentic_df[agentic_df['_id'] == selected_report_id]
        if selected_report_series_df.empty:
            error_return_tuple[0] = gr.update(value=f"*Error: Report with ID {selected_report_id} not found.*")
            return error_return_tuple
            
        selected_report_series = selected_report_series_df.iloc[0]
        report_markdown = format_report_for_display(selected_report_series)
        
        # Use the session cache
        reconstructed_data, updated_cache = fetch_and_reconstruct_data_from_bubble(selected_report_series, session_cache)

        if reconstructed_data:
            raw_results_state = reconstructed_data
            actionable_okrs_dict = reconstructed_data.get("actionable_okrs", {})
            all_krs_state = extract_key_results_for_selection(actionable_okrs_dict)
            if all_krs_state:
                kr_choices = [(kr['kr_description'], kr['unique_kr_id']) for kr in all_krs_state]
                key_results_cbg_update = gr.update(choices=kr_choices, value=[], interactive=True)
                okrs_list = actionable_okrs_dict.get("okrs", [])
                output_md_parts = [format_single_okr_for_display(okr, okr_main_index=i) for i, okr in enumerate(okrs_list)]
                okr_details_md = "\n\n---\n\n".join(output_md_parts) if output_md_parts else "No OKRs defined."
            else:
                key_results_cbg_update = gr.update(choices=[], value=[], interactive=False)
                okr_details_md = "No Key Results found for this report."
                all_krs_state = []
        else:
            key_results_cbg_update = gr.update(choices=[], value=[], interactive=False)
            okr_details_md = "Error: Could not fetch or reconstruct OKR data."
            raw_results_state, all_krs_state = None, []
        
        return (
            report_markdown, key_results_cbg_update, okr_details_md,
            raw_results_state, [], all_krs_state, updated_cache
        )

    with gr.Tabs() as tabs:
        with gr.TabItem("1️⃣ Dashboard", id="tab_dashboard"):
            gr.Markdown("I dati visualizzati in questo pannello sono caricati direttamente da Bubble.io.")
            dashboard_display_html = gr.HTML("<p style='text-align:center;'>Caricamento dashboard...</p>")

        analytics_tab_instance.create_tab_ui()

        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:
        report_selection_outputs = [
            agentic_report_display_md, key_results_cbg, okr_detail_display_md,
            orchestration_raw_results_st, selected_key_result_ids_st, 
            key_results_for_selection_st, reconstruction_cache_st # Pass cache state back
        ]
        report_selector_dd.change(
            fn=update_report_and_okr_display,
            inputs=[report_selector_dd, token_state, reconstruction_cache_st], # Pass cache state in
            outputs=report_selection_outputs,
            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 # Pass cache state back
    ]

    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],
        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], # Pass cache state in
        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)