ui_generators.py

# ui_generators.py
"""
Generates HTML content and Matplotlib plots for the Gradio UI tabs,
and UI components for the Analytics tab.
"""
import pandas as pd
import logging
import matplotlib.pyplot as plt
import matplotlib # To ensure backend is switched before any plt import from other modules if app structure changes
import gradio as gr # Added for UI components

# Switch backend for Matplotlib to Agg for Gradio compatibility
matplotlib.use('Agg')


# Assuming config.py contains all necessary constants
from config import (
    BUBBLE_POST_DATE_COLUMN_NAME, BUBBLE_MENTIONS_DATE_COLUMN_NAME, BUBBLE_MENTIONS_ID_COLUMN_NAME,
    FOLLOWER_STATS_TYPE_COLUMN, FOLLOWER_STATS_CATEGORY_COLUMN, FOLLOWER_STATS_ORGANIC_COLUMN,
    FOLLOWER_STATS_PAID_COLUMN, FOLLOWER_STATS_CATEGORY_COLUMN_DT, UI_DATE_FORMAT, UI_MONTH_FORMAT
)

# Configure logging for this module if not already configured at app level
# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(module)s - %(message)s')


def display_main_dashboard(token_state):
    """Generates HTML for the main dashboard display using data from token_state."""
    if not token_state or not token_state.get("token"):
        logging.warning("Dashboard display: Access denied. No token available.")
        return "❌ Access denied. No token available for dashboard."

    html_parts = ["<div style='padding:10px;'><h3>Dashboard Overview</h3>"]

    # Display Recent Posts
    posts_df = token_state.get("bubble_posts_df", pd.DataFrame())
    html_parts.append(f"<h4>Recent Posts ({len(posts_df)} in Bubble):</h4>")
    if not posts_df.empty:
        cols_to_show_posts = [col for col in [BUBBLE_POST_DATE_COLUMN_NAME, 'text', 'sentiment', 'summary_text', 'li_eb_label'] if col in posts_df.columns]
        if not cols_to_show_posts:
            html_parts.append("<p>No relevant post columns found to display.</p>")
        else:
            display_df_posts = posts_df.copy()
            if BUBBLE_POST_DATE_COLUMN_NAME in display_df_posts.columns:
                try:
                    # Ensure the date column is datetime before formatting
                    display_df_posts[BUBBLE_POST_DATE_COLUMN_NAME] = pd.to_datetime(display_df_posts[BUBBLE_POST_DATE_COLUMN_NAME], errors='coerce')
                    display_df_posts = display_df_posts.sort_values(by=BUBBLE_POST_DATE_COLUMN_NAME, ascending=False)
                    # Format for display after sorting
                    display_df_posts[BUBBLE_POST_DATE_COLUMN_NAME] = display_df_posts[BUBBLE_POST_DATE_COLUMN_NAME].dt.strftime(UI_DATE_FORMAT)
                except Exception as e:
                    logging.error(f"Error formatting post dates for display: {e}")
                    html_parts.append("<p>Error formatting post dates.</p>")
            html_parts.append(display_df_posts[cols_to_show_posts].head().to_html(escape=False, index=False, classes="table table-striped table-sm"))
    else:
        html_parts.append("<p>No posts loaded from Bubble.</p>")
    html_parts.append("<hr/>")

    # Display Recent Mentions
    mentions_df = token_state.get("bubble_mentions_df", pd.DataFrame())
    html_parts.append(f"<h4>Recent Mentions ({len(mentions_df)} in Bubble):</h4>")
    if not mentions_df.empty:
        cols_to_show_mentions = [col for col in [BUBBLE_MENTIONS_DATE_COLUMN_NAME, "mention_text", "sentiment_label"] if col in mentions_df.columns]
        if not cols_to_show_mentions:
            html_parts.append("<p>No relevant mention columns found to display.</p>")
        else:
            display_df_mentions = mentions_df.copy()
            if BUBBLE_MENTIONS_DATE_COLUMN_NAME in display_df_mentions.columns:
                try:
                    display_df_mentions[BUBBLE_MENTIONS_DATE_COLUMN_NAME] = pd.to_datetime(display_df_mentions[BUBBLE_MENTIONS_DATE_COLUMN_NAME], errors='coerce')
                    display_df_mentions = display_df_mentions.sort_values(by=BUBBLE_MENTIONS_DATE_COLUMN_NAME, ascending=False)
                    display_df_mentions[BUBBLE_MENTIONS_DATE_COLUMN_NAME] = display_df_mentions[BUBBLE_MENTIONS_DATE_COLUMN_NAME].dt.strftime(UI_DATE_FORMAT)
                except Exception as e:
                    logging.error(f"Error formatting mention dates for display: {e}")
                    html_parts.append("<p>Error formatting mention dates.</p>")
            html_parts.append(display_df_mentions[cols_to_show_mentions].head().to_html(escape=False, index=False, classes="table table-striped table-sm"))
    else:
        html_parts.append("<p>No mentions loaded from Bubble.</p>")
    html_parts.append("<hr/>")

    # Display Follower Statistics Summary
    follower_stats_df = token_state.get("bubble_follower_stats_df", pd.DataFrame())
    html_parts.append(f"<h4>Follower Statistics ({len(follower_stats_df)} entries in Bubble):</h4>")
    if not follower_stats_df.empty:
        monthly_gains = follower_stats_df[follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_gains_monthly'].copy()
        if not monthly_gains.empty and FOLLOWER_STATS_CATEGORY_COLUMN in monthly_gains.columns and \
           FOLLOWER_STATS_ORGANIC_COLUMN in monthly_gains.columns and FOLLOWER_STATS_PAID_COLUMN in monthly_gains.columns:
            try:
                monthly_gains.loc[:, FOLLOWER_STATS_CATEGORY_COLUMN_DT] = pd.to_datetime(monthly_gains[FOLLOWER_STATS_CATEGORY_COLUMN], errors='coerce')
                monthly_gains_display = monthly_gains.sort_values(by=FOLLOWER_STATS_CATEGORY_COLUMN_DT, ascending=False)
                latest_gain = monthly_gains_display.head(1).copy() 
                if not latest_gain.empty:
                    latest_gain.loc[:, FOLLOWER_STATS_CATEGORY_COLUMN] = latest_gain[FOLLOWER_STATS_CATEGORY_COLUMN_DT].dt.strftime(UI_DATE_FORMAT) 
                    html_parts.append("<h5>Latest Monthly Follower Gain:</h5>")
                    html_parts.append(latest_gain[[FOLLOWER_STATS_CATEGORY_COLUMN, FOLLOWER_STATS_ORGANIC_COLUMN, FOLLOWER_STATS_PAID_COLUMN]].to_html(escape=True, index=False, classes="table table-sm"))
                else:
                    html_parts.append("<p>No valid monthly follower gain data to display after processing.</p>")
            except Exception as e:
                logging.error(f"Error formatting follower gain dates for display: {e}", exc_info=True)
                html_parts.append("<p>Error displaying monthly follower gain data.</p>")
        else:
            html_parts.append("<p>No monthly follower gain data or required columns are missing.</p>")

        demographics_count = len(follower_stats_df[follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] != 'follower_gains_monthly'])
        html_parts.append(f"<p>Total demographic entries (seniority, industry, etc.): {demographics_count}</p>")
    else:
        html_parts.append("<p>No follower statistics loaded from Bubble.</p>")

    html_parts.append("</div>")
    return "".join(html_parts)


def run_mentions_tab_display(token_state):
    """Generates HTML and a plot for the Mentions tab."""
    logging.info("Updating Mentions Tab display.")
    if not token_state or not token_state.get("token"):
        logging.warning("Mentions tab: Access denied. No token.")
        return "❌ Access denied. No token available for mentions.", None

    mentions_df = token_state.get("bubble_mentions_df", pd.DataFrame())
    if mentions_df.empty:
        logging.info("Mentions tab: No mentions data in Bubble.")
        return "<p style='text-align:center;'>No mentions data in Bubble. Try syncing.</p>", None

    html_parts = ["<h3 style='text-align:center;'>Recent Mentions</h3>"]
    display_columns = [col for col in [BUBBLE_MENTIONS_DATE_COLUMN_NAME, "mention_text", "sentiment_label", BUBBLE_MENTIONS_ID_COLUMN_NAME] if col in mentions_df.columns]

    mentions_df_display = mentions_df.copy()
    if BUBBLE_MENTIONS_DATE_COLUMN_NAME in mentions_df_display.columns:
        try:
            mentions_df_display[BUBBLE_MENTIONS_DATE_COLUMN_NAME] = pd.to_datetime(mentions_df_display[BUBBLE_MENTIONS_DATE_COLUMN_NAME], errors='coerce')
            mentions_df_display = mentions_df_display.sort_values(by=BUBBLE_MENTIONS_DATE_COLUMN_NAME, ascending=False)
            mentions_df_display[BUBBLE_MENTIONS_DATE_COLUMN_NAME] = mentions_df_display[BUBBLE_MENTIONS_DATE_COLUMN_NAME].dt.strftime(UI_DATE_FORMAT)
        except Exception as e:
            logging.error(f"Error formatting mention dates for tab display: {e}")
            html_parts.append("<p>Error formatting mention dates.</p>")

    if not display_columns or mentions_df_display[display_columns].empty:
        html_parts.append("<p>Required columns for mentions display are missing or no data after processing.</p>")
    else:
        html_parts.append(mentions_df_display[display_columns].head(20).to_html(escape=False, index=False, classes="table table-sm"))

    mentions_html_output = "\n".join(html_parts)
    fig = None # Initialize fig to None
    if not mentions_df.empty and "sentiment_label" in mentions_df.columns:
        try:
            fig_plot, ax = plt.subplots(figsize=(6,4))
            sentiment_counts = mentions_df["sentiment_label"].value_counts()
            sentiment_counts.plot(kind='bar', ax=ax, color=['#4CAF50', '#FFC107', '#F44336', '#9E9E9E', '#2196F3'])
            ax.set_title("Mention Sentiment Distribution")
            ax.set_ylabel("Count")
            plt.xticks(rotation=45, ha='right')
            plt.tight_layout()
            fig = fig_plot # Assign the figure object
            logging.info("Mentions tab: Sentiment distribution plot generated.")
        except Exception as e:
            logging.error(f"Error generating mentions plot: {e}", exc_info=True)
            # fig remains None
        finally:
            if fig is not None and fig is not plt: # Ensure we don't close the global plt if fig is just an alias
                 plt.close(fig) # Close the specific figure to free memory
            elif fig is None and plt.get_fignums(): # If fig is None but some global figure might exist from error
                 plt.close('all')


    return mentions_html_output, fig


def run_follower_stats_tab_display(token_state):
    """Generates HTML and plots for the Follower Stats tab."""
    logging.info("Updating Follower Stats Tab display.")
    if not token_state or not token_state.get("token"):
        logging.warning("Follower stats tab: Access denied. No token.")
        return "❌ Access denied. No token available for follower stats.", None, None, None

    follower_stats_df_orig = token_state.get("bubble_follower_stats_df", pd.DataFrame())
    if follower_stats_df_orig.empty:
        logging.info("Follower stats tab: No follower stats data in Bubble.")
        return "<p style='text-align:center;'>No follower stats data in Bubble. Try syncing.</p>", None, None, None

    follower_stats_df = follower_stats_df_orig.copy()
    html_parts = ["<div style='padding:10px;'><h3 style='text-align:center;'>Follower Statistics Overview</h3>"]

    plot_monthly_gains = None
    plot_seniority_dist = None
    plot_industry_dist = None

    # --- Monthly Gains Table & Plot ---
    monthly_gains_df = follower_stats_df[
        (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_gains_monthly') &
        (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
        (follower_stats_df[FOLLOWER_STATS_ORGANIC_COLUMN].notna()) &
        (follower_stats_df[FOLLOWER_STATS_PAID_COLUMN].notna())
    ].copy()

    if not monthly_gains_df.empty:
        try:
            monthly_gains_df.loc[:, FOLLOWER_STATS_CATEGORY_COLUMN_DT] = pd.to_datetime(monthly_gains_df[FOLLOWER_STATS_CATEGORY_COLUMN], errors='coerce')
            monthly_gains_df_sorted_table = monthly_gains_df.sort_values(by=FOLLOWER_STATS_CATEGORY_COLUMN_DT, ascending=False)

            html_parts.append("<h4>Monthly Follower Gains (Last 13 Months):</h4>")
            table_display_df = monthly_gains_df_sorted_table.copy()
            table_display_df.loc[:,FOLLOWER_STATS_CATEGORY_COLUMN] = table_display_df[FOLLOWER_STATS_CATEGORY_COLUMN_DT].dt.strftime(UI_MONTH_FORMAT)
            html_parts.append(table_display_df[[FOLLOWER_STATS_CATEGORY_COLUMN, FOLLOWER_STATS_ORGANIC_COLUMN, FOLLOWER_STATS_PAID_COLUMN]].head(13).to_html(escape=True, index=False, classes="table table-sm"))

            monthly_gains_df_sorted_plot = monthly_gains_df.sort_values(by=FOLLOWER_STATS_CATEGORY_COLUMN_DT, ascending=True).copy()
            monthly_gains_df_sorted_plot.loc[:, '_plot_month'] = monthly_gains_df_sorted_plot[FOLLOWER_STATS_CATEGORY_COLUMN_DT].dt.strftime(UI_MONTH_FORMAT)
            plot_data = monthly_gains_df_sorted_plot.groupby('_plot_month').agg(
                organic=(FOLLOWER_STATS_ORGANIC_COLUMN, 'sum'),
                paid=(FOLLOWER_STATS_PAID_COLUMN, 'sum')
            ).reset_index().sort_values(by='_plot_month') # Ensure month order for plot

            fig_gains, ax_gains = plt.subplots(figsize=(10,5))
            ax_gains.plot(plot_data['_plot_month'], plot_data['organic'], marker='o', linestyle='-', label='Organic Gain')
            ax_gains.plot(plot_data['_plot_month'], plot_data['paid'], marker='x', linestyle='--', label='Paid Gain')
            ax_gains.set_title("Monthly Follower Gains Over Time")
            ax_gains.set_ylabel("Follower Count")
            ax_gains.set_xlabel("Month (YYYY-MM)")
            plt.xticks(rotation=45, ha='right')
            ax_gains.legend()
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.tight_layout()
            plot_monthly_gains = fig_gains
            logging.info("Follower stats tab: Monthly gains plot generated.")
        except Exception as e:
            logging.error(f"Error processing or plotting monthly gains: {e}", exc_info=True)
            html_parts.append("<p>Error displaying monthly follower gain data.</p>")
        finally:
            if plot_monthly_gains is not None and plot_monthly_gains is not plt:
                 plt.close(plot_monthly_gains)
            elif plot_monthly_gains is None and plt.get_fignums():
                 plt.close('all')
    else:
        html_parts.append("<p>No monthly follower gain data available or required columns missing.</p>")
    html_parts.append("<hr/>")

    # --- Seniority Table & Plot ---
    seniority_df = follower_stats_df[
        (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_seniority') &
        (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
        (follower_stats_df[FOLLOWER_STATS_ORGANIC_COLUMN].notna())
    ].copy()
    if not seniority_df.empty:
        try:
            seniority_df_sorted = seniority_df.sort_values(by=FOLLOWER_STATS_ORGANIC_COLUMN, ascending=False)
            html_parts.append("<h4>Followers by Seniority (Top 10 Organic):</h4>")
            html_parts.append(seniority_df_sorted[[FOLLOWER_STATS_CATEGORY_COLUMN, FOLLOWER_STATS_ORGANIC_COLUMN, FOLLOWER_STATS_PAID_COLUMN]].head(10).to_html(escape=True, index=False, classes="table table-sm"))

            fig_seniority, ax_seniority = plt.subplots(figsize=(8,5))
            top_n_seniority = seniority_df_sorted.nlargest(10, FOLLOWER_STATS_ORGANIC_COLUMN)
            ax_seniority.bar(top_n_seniority[FOLLOWER_STATS_CATEGORY_COLUMN], top_n_seniority[FOLLOWER_STATS_ORGANIC_COLUMN], color='skyblue')
            ax_seniority.set_title("Follower Distribution by Seniority (Top 10 Organic)")
            ax_seniority.set_ylabel("Organic Follower Count")
            plt.xticks(rotation=45, ha='right')
            plt.grid(axis='y', linestyle='--', alpha=0.7)
            plt.tight_layout()
            plot_seniority_dist = fig_seniority
            logging.info("Follower stats tab: Seniority distribution plot generated.")
        except Exception as e:
            logging.error(f"Error processing or plotting seniority data: {e}", exc_info=True)
            html_parts.append("<p>Error displaying follower seniority data.</p>")
        finally:
            if plot_seniority_dist is not None and plot_seniority_dist is not plt:
                 plt.close(plot_seniority_dist)
            elif plot_seniority_dist is None and plt.get_fignums():
                 plt.close('all')
    else:
        html_parts.append("<p>No follower seniority data available or required columns missing.</p>")
    html_parts.append("<hr/>")

    # --- Industry Table & Plot ---
    industry_df = follower_stats_df[
        (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_industry') &
        (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
        (follower_stats_df[FOLLOWER_STATS_ORGANIC_COLUMN].notna())
    ].copy()
    if not industry_df.empty:
        try:
            industry_df_sorted = industry_df.sort_values(by=FOLLOWER_STATS_ORGANIC_COLUMN, ascending=False)
            html_parts.append("<h4>Followers by Industry (Top 10 Organic):</h4>")
            html_parts.append(industry_df_sorted[[FOLLOWER_STATS_CATEGORY_COLUMN, FOLLOWER_STATS_ORGANIC_COLUMN, FOLLOWER_STATS_PAID_COLUMN]].head(10).to_html(escape=True, index=False, classes="table table-sm"))

            fig_industry, ax_industry = plt.subplots(figsize=(8,5))
            top_n_industry = industry_df_sorted.nlargest(10, FOLLOWER_STATS_ORGANIC_COLUMN)
            ax_industry.bar(top_n_industry[FOLLOWER_STATS_CATEGORY_COLUMN], top_n_industry[FOLLOWER_STATS_ORGANIC_COLUMN], color='lightcoral')
            ax_industry.set_title("Follower Distribution by Industry (Top 10 Organic)")
            ax_industry.set_ylabel("Organic Follower Count")
            plt.xticks(rotation=45, ha='right')
            plt.grid(axis='y', linestyle='--', alpha=0.7)
            plt.tight_layout()
            plot_industry_dist = fig_industry
            logging.info("Follower stats tab: Industry distribution plot generated.")
        except Exception as e:
            logging.error(f"Error processing or plotting industry data: {e}", exc_info=True)
            html_parts.append("<p>Error displaying follower industry data.</p>")
        finally:
            if plot_industry_dist is not None and plot_industry_dist is not plt:
                 plt.close(plot_industry_dist)
            elif plot_industry_dist is None and plt.get_fignums():
                 plt.close('all')

    else:
        html_parts.append("<p>No follower industry data available or required columns missing.</p>")

    html_parts.append("</div>")
    follower_html_output = "\n".join(html_parts)
    return follower_html_output, plot_monthly_gains, plot_seniority_dist, plot_industry_dist


# --- NEW UI GENERATION LOGIC FOR ANALYTICS TAB ---
def create_analytics_plot_panel(label, plot_id_str):
    """
    Creates a Gradio Column representing a single plot panel.
    This panel contains the plot and its associated "bomb" button to the right.
    Returns the panel (gr.Column), the plot component, and the bomb button.
    """
    with gr.Column() as panel_col: # This is the component that will go into the two-column rows
        with gr.Row(equal_height=False, variant="panel"): # A sub-row for plot and button
            with gr.Column(scale=8): # Adjust scale as needed for plot vs button
                plot_component = gr.Plot(label=label)
            with gr.Column(scale=2, min_width=70, vertical_alignment="center"): # min_width for button, try to center vertically
                bomb_button = gr.Button("💣", variant="secondary", size="sm", elem_id=f"bomb_{plot_id_str}")
    return panel_col, plot_component, bomb_button

def build_analytics_tab_plot_area(plot_configs):
    """
    Builds the main plot area for the Analytics tab, arranging plot panels into rows of two.
    This function should be called within the "Analytics" TabItem context in app.py,
    specifically within the main column designated for plots.
    Returns a dictionary of plot UI objects.
    """
    logging.info(f"Building plot area for {len(plot_configs)} analytics plots.")
    plot_ui_objects = {} # Stores {"plot_id": {"plot_component": gr.Plot, "bomb_button": gr.Button, "label": str}}
    
    current_section_title = ""
    
    # Iterate through plot_configs to create sections and then rows of plots
    for i in range(0, len(plot_configs), 2): # Process in pairs for two plots per row
        # Handle section header
        # Check if the section of the first plot in the potential pair is new
        if plot_configs[i]["section"] != current_section_title:
            current_section_title = plot_configs[i]["section"]
            gr.Markdown(f"### {current_section_title}")

        with gr.Row(): # Create a new row for up to two plot panels
            # First plot in the pair
            config1 = plot_configs[i]
            panel_col1, plot_component1, bomb_button1 = \
                create_analytics_plot_panel(config1["label"], config1["id"])
            plot_ui_objects[config1["id"]] = {
                "plot_component": plot_component1,
                "bomb_button": bomb_button1,
                "label": config1["label"]
            }
            logging.debug(f"Created UI panel for plot_id: {config1['id']}")

            # Second plot in the pair, if it exists
            if i + 1 < len(plot_configs):
                config2 = plot_configs[i+1]
                # Check if the second plot starts a new section (edge case if sections aren't aligned with pairs)
                if config2["section"] != current_section_title:
                    # This row is finished, start a new section header for the next plot
                    # This might leave the current row with only one plot if sections change mid-pair
                    # For simplicity, we assume sections align well or the header appears before the row with the new section's first plot
                    # The current loop structure will create a new header before the next gr.Row if section changes
                    pass # The next iteration's section check will handle this.

                panel_col2, plot_component2, bomb_button2 = \
                    create_analytics_plot_panel(config2["label"], config2["id"])
                plot_ui_objects[config2["id"]] = {
                    "plot_component": plot_component2,
                    "bomb_button": bomb_button2,
                    "label": config2["label"]
                }
                logging.debug(f"Created UI panel for plot_id: {config2['id']}")
            else:
                # If there's no second plot, this row will just have one panel
                pass
        
    logging.info(f"Finished building plot area. Total plot objects: {len(plot_ui_objects)}")
    return plot_ui_objects