#ui_main_page_enhancements.py
import gradio as gr
import pandas as pd # Needed for state variables that hold DataFrames

# Custom CSS for professional, modern UI
CUSTOM_CSS = """
/* === ROOT VARIABLES === */
:root {
    --primary-color: #0A66C2;
    --primary-hover: #004182;
    --secondary-color: #5E9ED6;
    --accent-color: #F3F2EF;
    --success-color: #57C4A3;
    --warning-color: #F5B800;
    --error-color: #CC1016;
    --text-primary: #000000DE;
    --text-secondary: #00000099;
    --text-tertiary: #00000061;
    --background-primary: #FFFFFF;
    --background-secondary: #F8F9FA;
    --background-tertiary: #F3F2EF;
    --border-color: #E5E5E5;
    --border-radius: 12px;
    --shadow-light: 0 2px 8px rgba(0, 0, 0, 0.08);
    --shadow-medium: 0 4px 16px rgba(0, 0, 0, 0.12);
    --shadow-heavy: 0 8px 32px rgba(0, 0, 0, 0.16);
    --transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
}

/* === GLOBAL STYLES === */
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%) !important;
    min-height: 100vh;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

.gradio-container .main {
    background: transparent;
    padding: 0;
    max-width: none !important;
}

/* === MAIN HEADER === */
.main-header {
    background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
    color: white;
    padding: 2rem 0;
    margin: -1rem -1rem 2rem -1rem;
    text-align: center;
    box-shadow: var(--shadow-medium);
    position: relative;
    overflow: hidden;
}

.main-header::before {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: url('data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 100 100"><defs><pattern id="grid" width="10" height="10" patternUnits="userSpaceOnUse"><path d="M 10 0 L 0 0 0 10" fill="none" stroke="rgba(255,255,255,0.1)" stroke-width="0.5"/></pattern></defs><rect width="100" height="100" fill="url(%23grid)"/></svg>');
    opacity: 0.3;
}

.main-header h1 {
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    margin: 0 !important;
    text-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
    position: relative;
    z-index: 1;
}

.main-header .subtitle {
    font-size: 1.1rem;
    opacity: 0.9;
    margin-top: 0.5rem;
    position: relative;
    z-index: 1;
}

/* === STATUS BAR === */
.status-container {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 1rem;
    margin-bottom: 1.5rem;
    box-shadow: var(--shadow-light);
    border-left: 4px solid var(--primary-color);
}

/* === TAB CONTAINER === */
.tab-nav {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 0.5rem;
    box-shadow: var(--shadow-light);
    margin-bottom: 1.5rem;
    border: none;
}

.tab-nav button {
    background: transparent !important;
    border: none !important;
    padding: 1rem 1.5rem !important;
    margin: 0 0.25rem !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
    transition: var(--transition) !important;
    color: var(--text-secondary) !important;
    position: relative;
}

.tab-nav button:hover {
    background: var(--background-secondary) !important;
    color: var(--text-primary) !important;
    transform: translateY(-2px);
}

.tab-nav button.selected {
    background: var(--primary-color) !important;
    color: white !important;
    box-shadow: var(--shadow-medium);
}

/* === TAB CONTENT === */
.tabitem {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 2rem;
    box-shadow: var(--shadow-light);
    border: 1px solid var(--border-color);
    margin-top: 0;
}

/* === FORM ELEMENTS === */
.gr-form {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    box-shadow: var(--shadow-light);
    border: 1px solid var(--border-color);
}

input, textarea, select {
    border: 2px solid var(--border-color) !important;
    border-radius: 8px !important;
    padding: 0.75rem 1rem !important;
    font-size: 0.95rem !important;
    transition: var(--transition) !important;
    background: var(--background-primary) !important;
}

input:focus, textarea:focus, select:focus {
    border-color: var(--primary-color) !important;
    box-shadow: 0 0 0 3px rgba(10, 102, 194, 0.1) !important;
    outline: none !important;
}

/* === BUTTONS === */
.gr-button {
    background: var(--primary-color) !important;
    color: white !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 0.75rem 1.5rem !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
    transition: var(--transition) !important;
    cursor: pointer !important;
    box-shadow: var(--shadow-light);
}

.gr-button:hover {
    background: var(--primary-hover) !important;
    transform: translateY(-2px);
    box-shadow: var(--shadow-medium);
}

.gr-button.secondary {
    background: var(--background-secondary) !important;
    color: var(--text-primary) !important;
    border: 2px solid var(--border-color) !important;
}

.gr-button.secondary:hover {
    background: var(--background-tertiary) !important;
    border-color: var(--primary-color) !important;
}

/* === CARDS === */
.card {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    box-shadow: var(--shadow-light);
    border: 1px solid var(--border-color);
    transition: var(--transition);
    margin-bottom: 1rem;
}

.card:hover {
    box-shadow: var(--shadow-medium);
    transform: translateY(-2px);
}

.card-header {
    font-size: 1.25rem;
    font-weight: 700;
    color: var(--text-primary);
    margin-bottom: 0.5rem;
    display: flex;
    align-items: center;
    gap: 0.5rem;
}

.card-content {
    color: var(--text-secondary);
    line-height: 1.6;
}

/* === CHARTS AND PLOTS === */
.plot-container {
    background: var(--background-primary);
    border-radius: var(--border-radius);
    padding: 1rem;
    box-shadow: var(--shadow-light);
    border: 1px solid var(--border-color);
    margin: 1rem 0;
}

/* === MARKDOWN CONTENT === */
.markdown-content h1,
.markdown-content h2,
.markdown-content h3 {
    color: var(--text-primary);
    font-weight: 700;
    margin-top: 1.5rem;
    margin-bottom: 0.75rem;
}

.markdown-content h1 {
    font-size: 2rem;
    border-bottom: 3px solid var(--primary-color);
    padding-bottom: 0.5rem;
}

.markdown-content h2 {
    font-size: 1.5rem;
    color: var(--primary-color);
}

.markdown-content h3 {
    font-size: 1.25rem;
    color: var(--secondary-color);
}

.markdown-content p {
    line-height: 1.7;
    color: var(--text-secondary);
    margin-bottom: 1rem;
}

.markdown-content ul, .markdown-content ol {
    padding-left: 1.5rem;
    margin-bottom: 1rem;
}

.markdown-content li {
    margin-bottom: 0.5rem;
    line-height: 1.6;
    color: var(--text-secondary);
}

.markdown-content code {
    background: var(--background-tertiary);
    padding: 0.2rem 0.4rem;
    border-radius: 4px;
    font-family: 'Fira Code', monospace;
    font-size: 0.9rem;
}

.markdown-content pre {
    background: var(--background-tertiary);
    padding: 1rem;
    border-radius: 8px;
    overflow-x: auto;
    margin: 1rem 0;
}

/* === LOADING STATES === */
.loading-spinner {
    display: inline-block;
    width: 20px;
    height: 20px;
    border: 3px solid rgba(10, 102, 194, 0.3);
    border-radius: 50%;
    border-top-color: var(--primary-color);
    animation: spin 1s ease-in-out infinite;
}

@keyframes spin {
    to { transform: rotate(360deg); }
}

/* === RESPONSIVE DESIGN === */
@media (max-width: 768px) {
    .main-header h1 {
        font-size: 2rem !important;
    }
    
    .tabitem {
        padding: 1rem;
    }
    
    .tab-nav button {
        padding: 0.75rem 1rem !important;
        font-size: 0.9rem !important;
    }
}

/* === ANIMATIONS === */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(20px); }
    to { opacity: 1; transform: translateY(0); }
}

.fade-in {
    animation: fadeIn 0.6s ease-out;
}

/* === ACCESSIBILITY === */
.sr-only {
    position: absolute;
    width: 1px;
    height: 1px;
    padding: 0;
    margin: -1px;
    overflow: hidden;
    clip: rect(0, 0, 0, 0);
    white-space: nowrap;
    border: 0;
}

/* === GRADIO SPECIFIC OVERRIDES === */
.gradio-container .wrap {
    border-radius: var(--border-radius) !important;
    border: 1px solid var(--border-color) !important;
}

.gradio-container .panel {
    background: var(--background-primary) !important;
    border-radius: var(--border-radius) !important;
}

.gradio-container .form {
    background: transparent !important;
}

/* === STATUS INDICATORS === */
.status-success {
    color: var(--success-color) !important;
    background: rgba(87, 196, 163, 0.1) !important;
    padding: 0.5rem 1rem;
    border-radius: 6px;
    border-left: 4px solid var(--success-color);
}

.status-warning {
    color: var(--warning-color) !important;
    background: rgba(245, 184, 0, 0.1) !important;
    padding: 0.5rem 1rem;
    border-radius: 6px;
    border-left: 4px solid var(--warning-color);
}

.status-error {
    color: var(--error-color) !important;
    background: rgba(204, 16, 22, 0.1) !important;
    padding: 0.5rem 1rem;
    border-radius: 6px;
    border-left: 4px solid var(--error-color);
}

/* === HIDE GRADIO FOOTER === */
footer {
    display: none !important;
}

.footer {
    display: none !important;
}

/* === IMPROVE SPACING === */
.block {
    margin-bottom: 1.5rem !important;
}

/* === CUSTOM COMPONENTS === */
.metric-card {
    background: linear-gradient(135deg, var(--primary-color), var(--secondary-color));
    color: white;
    padding: 1.5rem;
    border-radius: var(--border-radius);
    text-align: center;
    box-shadow: var(--shadow-medium);
    margin: 0.5rem;
}

.metric-value {
    font-size: 2rem;
    font-weight: 700;
    margin-bottom: 0.5rem;
}

.metric-label {
    font-size: 0.9rem;
    opacity: 0.9;
}
"""

# Custom theme
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="sky",
    neutral_hue="slate",
    spacing_size="md",
    radius_size="md"
).set(
    body_background_fill="linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%)",
    block_background_fill="white",
    block_border_width="1px",
    block_border_color="#e5e5e5",
    block_radius="12px",
    block_shadow="0 2px 8px rgba(0, 0, 0, 0.08)",
    button_primary_background_fill="#0A66C2",
    button_primary_background_fill_hover="#004182",
    button_primary_text_color="white",
    button_secondary_background_fill="#f8f9fa",
    button_secondary_background_fill_hover="#f3f2ef",
    input_border_color="#e5e5e5",
    input_border_color_focus="#0A66C2",
    input_border_width="2px"
)

def update_report_display_enhanced(selected_report_id: str, current_token_state: dict, format_report_for_display_func):
    """
    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'.

    Args:
        selected_report_id (str): The ID of the selected report.
        current_token_state (dict): The current state dictionary containing data.
        format_report_for_display_func (callable): The function to format report data.

    Returns:
        tuple: A tuple of gr.update objects for the header and body displays.
    """
    # Define empty states for header and body with improved styling
    empty_header_html = """
    <div class="card">
        <div class="card-header">
            📊 Comprehensive Analysis Report
        </div>
        <div class="card-content">
            <p>AI-Generated Insights from Your LinkedIn Data</p>
            <div class="status-badge">Generated from Bubble.io</div>
        </div>
    </div>
    """
    empty_body_markdown_no_selection = """
    <div class="card" style="text-align: center; padding: 3rem;">
        <div style="font-size: 3rem; margin-bottom: 1rem;">📋</div>
        <h3 style="color: var(--text-primary); margin-bottom: 0.5rem;">Select a Report</h3>
        <p style="color: var(--text-secondary);">
            Choose a report from the dropdown above to view its detailed analysis and insights.
        </p>
    </div>
    """
    empty_body_markdown_no_data = """
    <div class="card" style="text-align: center; padding: 3rem;">
        <div style="font-size: 3rem; margin-bottom: 1rem;">⚠️</div>
        <h3 style="color: var(--warning-color); margin-bottom: 0.5rem;">Data Not Available</h3>
        <p style="color: var(--text-secondary);">
            Analysis data is not loaded or is empty. Please try refreshing the page.
        </p>
    </div>
    """
    empty_body_markdown_not_found = lambda _id: f"""
    <div class="card" style="text-align: center; padding: 3rem;">
        <div style="font-size: 3rem; margin-bottom: 1rem;">❌</div>
        <h3 style="color: var(--error-color); margin-bottom: 0.5rem;">Report Not Found</h3>
        <p style="color: var(--text-secondary);">
            Report with ID '{_id}' was not found in the database.
        </p>
    </div>
    """

    if not selected_report_id:
        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:
        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:
        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_func(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'])
    )

def build_main_app_ui(
    analytics_tab_instance,
    AGENTIC_MODULES_LOADED,
    build_home_tab_ui_func,
    create_enhanced_report_tab_func,
    create_enhanced_okr_tab_func,
    format_report_for_display_func # New argument to pass the formatting function
):
    """
    Builds the main Gradio application UI with enhanced styling and structure.

    Args:
        analytics_tab_instance (AnalyticsTab): An instance of the AnalyticsTab class.
        AGENTIC_MODULES_LOADED (bool): Flag indicating if agentic modules are loaded.
        build_home_tab_ui_func (callable): Function to build the Home tab UI.
        create_enhanced_report_tab_func (callable): Function to create the enhanced Report tab.
        create_enhanced_okr_tab_func (callable): Function to create the enhanced OKR tab.
        format_report_for_display_func (callable): Function to format report data for display.

    Returns:
        tuple: A tuple containing:
            - app (gr.Blocks): The main Gradio application object.
            - url_user_token_display (gr.Textbox): Hidden component for user token.
            - org_urn_display (gr.Textbox): Hidden component for organization URN.
            - status_box (gr.Textbox): Component for system status.
            - token_state (gr.State): Gradio state for token and data.
            - reconstruction_cache_st (gr.State): State for reconstructed OKR data cache.
            - enhanced_okr_display_html (gr.HTML): HTML component for enhanced OKR display.
            - tabs (gr.Tabs): The main tabs component.
            - report_selector_dd (gr.Dropdown): Dropdown for selecting reports.
            - agentic_display_outputs (list): List of Gradio components for agentic results.
            - analytics_tab_instance (AnalyticsTab): The analytics tab instance itself.
            - chat_histories_st (gr.State): State for chat histories.
    """
    app = gr.Blocks(
        theme=custom_theme,
        css=CUSTOM_CSS,
        title="LinkedIn Organization Dashboard",
        head="""
        <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet">
        <link href="https://fonts.googleapis.com/css2?family=Fira+Code:wght@400;500&display=swap" rel="stylesheet">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        """
    )

    with 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({})
        # NEW: State to hold the actionable_okrs dictionary explicitly
        actionable_okrs_data_st = gr.State({})

        # --- UI LAYOUT ---
        # Main Header
        with gr.Row():
            with gr.Column():
                gr.HTML("""
                <div class="main-header">
                    <h1>🚀 LinkedIn Organization Dashboard</h1>
                    <div class="subtitle">Advanced Analytics & AI-Powered Insights for LinkedIn Organizations</div>
                </div>
                """)
        
        # Hidden components for token management
        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 section with better styling
        with gr.Row():
            with gr.Column():
                gr.HTML('<div class="status-container">')
                status_box = gr.Textbox(
                    label="🔄 System Status", 
                    interactive=False, 
                    value="🔄 Initializing dashboard...",
                    elem_classes=["status-display"]
                )
                gr.HTML('</div>')

        with gr.Tabs(elem_classes=["tab-nav"]) as tabs:
            # --- HOME TAB ---
            with gr.TabItem("🏠 Home", id="tab_home", elem_classes=["tabitem"]):
                # 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_func()

                # 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
            analytics_tab_instance.create_tab_ui() # This is the "Graphs" tab, assuming its ID is "tab_analytics_module"

            # --- AGENTIC ANALYSIS REPORT TAB ---
            with gr.TabItem("📊 Analysis Reports", id="tab_agentic_report", visible=AGENTIC_MODULES_LOADED, elem_classes=["tabitem"]):
                # The create_enhanced_report_tab function handles the CSS and HTML structure
                agentic_pipeline_status_md, report_selector_dd, report_header_html_display, report_body_markdown_display = \
                    create_enhanced_report_tab_func(AGENTIC_MODULES_LOADED)

            # --- AGENTIC OKRS TAB ---
            with gr.TabItem("🎯 OKRs & Tasks", id="tab_agentic_okrs", visible=AGENTIC_MODULES_LOADED, elem_classes=["tabitem"]):
                gr.HTML("""
                <div class="card">
                    <div class="card-header">
                        🎯 AI Generated OKRs and Actionable Tasks
                    </div>
                    <div class="card-content">
                        <p>Based on AI analysis, the agent has proposed the following OKRs and actionable tasks from Bubble.io data.</p>
                    </div>
                </div>
                """)

                if not AGENTIC_MODULES_LOADED:
                    gr.HTML("""
                    <div class="card" style="text-align: center; padding: 2rem;">
                        <div style="font-size: 2rem; margin-bottom: 1rem;">🔴</div>
                        <h3 style="color: var(--error-color);">Module Loading Error</h3>
                        <p>Agentic modules could not be loaded. This tab is currently unavailable.</p>
                    </div>
                    """)

                # 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_func()

        # Add footer with enhanced styling
        gr.HTML("""
        <div style="text-align: center; padding: 2rem; margin-top: 3rem; border-top: 1px solid var(--border-color); color: var(--text-tertiary);">
            <p>© 2024 LinkedIn Organization Dashboard - Powered by AI & Advanced Analytics</p>
            <p style="font-size: 0.9rem; margin-top: 0.5rem;">
                🚀 Built with Gradio • 🔗 LinkedIn API • 🤖 AI Analytics • ☁️ Bubble.io Integration
            </p>
        </div>
        """)
    
    # 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
        actionable_okrs_data_st                # 11: NEW: The actionable_okrs dictionary state
    ]

    return (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, chat_histories_st,
            format_report_for_display_func) # Return format_report_for_display_func as well