import gradio as gr
import pandas as pd
import logging
from typing import Dict, Any, List, Optional

# Import the reconstruction function that now expects a cache dictionary
from services.report_data_handler import fetch_and_reconstruct_data_from_bubble

# UI formatting functions
try:
    from ui.insights_ui_generator import (
        format_report_for_display,
        # REMOVED: extract_key_results_for_selection, - No longer directly used here for display
        # REMOVED: format_single_okr_for_display - No longer directly used here for display
    )
    # NEW: Import the enhanced OKR display functions
    from ui.okr_ui_generator import format_okrs_for_enhanced_display, get_initial_okr_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 for when modules are not loaded, ensuring return signatures match
    def format_report_for_display(report_data):
        return {'header_html': '<h1>Agentic modules not loaded.</h1>', 'body_markdown': 'Report display unavailable.'}
    # These functions are now used by app.py directly for placeholder returns if modules are not loaded,
    # but their definitions might still be expected in some contexts, so providing minimal ones.
    def extract_key_results_for_selection(okrs_dict): return []
    def format_single_okr_for_display(okr_data, **kwargs): return "Agentic modules not loaded."
    # Placeholder for the new OKR UI display
    def format_okrs_for_enhanced_display(raw_results): return get_initial_okr_display()
    def get_initial_okr_display(): return """
        <div class="okr-container">
            <div class="okr-header">
                <div class="okr-title">🎯 AI-Generated OKRs & Strategic Tasks</div>
                <div class="okr-subtitle">Intelligent objectives and key results based on your LinkedIn analytics</div>
            </div>
            <div class="okr-stats-bar">
                <div class="stat-card"><div class="stat-number">0</div><div class="stat-label">Objectives</div></div>
                <div class="stat-card"><div class="stat-number">0</div><div class="stat-label">Key Results</div></div>
                <div class="stat-card"><div class="stat-number">0</div><div class="stat-label">Tasks</div></div>
                <div class="stat-card"><div class="stat-number">0</div><div class="stat-label">High Priority</div></div>
            </div>
            <div class="okr-content">
                <div class="empty-state">
                    <div class="empty-state-icon">⚠️</div>
                    <div class="empty-state-title">Modules Not Loaded</div>
                    <div class="empty-state-description">
                        Agentic analysis modules could not be loaded. OKR display unavailable.
                    </div>
                </div>
            </div>
        </div>
    """


logger = logging.getLogger(__name__)

def load_and_display_agentic_results(token_state: dict, session_cache: dict):
    """
    Loads agentic results from state, populates the report library, and displays
    the LATEST report and its fully reconstructed OKRs by default, using a session-specific cache.

    Args:
        token_state: The main state dictionary with Bubble data.
        session_cache: The session-specific cache for reconstructed data.

    Returns:
        A tuple of Gradio updates, including the updated cache, the new enhanced OKR HTML,
        and the extracted actionable_okrs_dict.
        This function now returns 12 values to match the expected outputs in app.py.
    """
    # Define placeholder content for empty or error states to match 12 outputs
    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">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>
        """

    # Default initial updates for 12 outputs
    initial_updates = (
        gr.update(value="Status: No agentic analysis data found."), # 0: agentic_pipeline_status_md
        gr.update(choices=[], value=None, interactive=False),       # 1: report_selector_dd
        gr.update(choices=[], value=[], interactive=False),         # 2: key_results_cbg (hidden in app.py)
        gr.update(value="No OKRs to display."),                     # 3: okr_detail_display_md (hidden in app.py)
        gr.update(value=None),                                      # 4: orchestration_raw_results_st
        gr.update(value=[]),                                        # 5: selected_key_result_ids_st (hidden in app.py)
        gr.update(value=[]),                                        # 6: key_results_for_selection_st (hidden in app.py)
        gr.update(value=empty_header_html),                         # 7: report_header_html_display
        gr.update(value=empty_body_markdown_no_selection),          # 8: report_body_markdown_display
        gr.update(value=session_cache),                             # 9: reconstruction_cache_st
        gr.update(value=get_initial_okr_display()),                 # 10: enhanced_okr_display_html
        gr.update(value={})                                         # 11: actionable_okrs_data_st (NEW)
    )

    if not AGENTIC_MODULES_LOADED:
        logger.error("Agentic modules not loaded, returning placeholder updates.")
        # Ensure error updates match the 12-item signature
        error_header_html = '<div class="report-title">Error Loading Report</div><div class="report-subtitle">Agentic modules not loaded.</div><div class="status-badge">Error</div>'
        error_body_markdown = '<div class="empty-state"><div class="empty-state-icon">❌</div><div class="empty-state-title">Module Error</div><div class="empty-state-description">Agentic analysis modules could not be loaded. Report display unavailable.</div></div>'
        return (
            gr.update(value="Status: Critical module import error."), # 0
            gr.update(choices=[], value=None, interactive=False), # 1
            gr.update(choices=[], value=[], interactive=False), # 2
            gr.update(value="Agentic modules not loaded. No OKRs to display."), # 3
            gr.update(value=None), # 4
            gr.update(value=[]), # 5
            gr.update(value=[]), # 6
            gr.update(value=error_header_html), # 7
            gr.update(value=error_body_markdown), # 8
            gr.update(value=session_cache), # 9
            gr.update(value=get_initial_okr_display()), # 10: Placeholder for enhanced OKR display
            gr.update(value={}) # 11: Placeholder for actionable_okrs_data_st (NEW)
        )

    agentic_df = token_state.get("bubble_agentic_analysis_data")

    if agentic_df is None or agentic_df.empty:
        logger.warning("Agentic analysis DataFrame is missing or empty.")
        return initial_updates

    try:
        if 'Created Date' not in agentic_df.columns or '_id' not in agentic_df.columns:
            raise KeyError("Required columns ('Created Date', '_id') not found.")

        agentic_df['Created Date'] = pd.to_datetime(agentic_df['Created Date'])
        agentic_df = agentic_df.sort_values(by='Created Date', ascending=False).reset_index(drop=True)

        report_choices = [(f"{row.get('report_type', 'Report')} - {row['Created Date'].strftime('%Y-%m-%d %H:%M')}", row['_id'])
                          for _, row in agentic_df.iterrows()]

        if not report_choices:
            return initial_updates

        quarterly_reports_df = agentic_df[agentic_df['report_type'] == 'Quarter'].copy()
        if quarterly_reports_df.empty:
            logger.warning("No quarterly reports found, falling back to latest available report.")
            latest_report_series = agentic_df.iloc[0] # Use the absolute latest if no quarterly
        else:
            latest_report_series = quarterly_reports_df.iloc[0]

        latest_report_id = latest_report_series['_id']

        # Split the formatted report content into header and body
        formatted_report_parts = format_report_for_display(latest_report_series)
        report_header_content = formatted_report_parts['header_html']
        report_body_content = formatted_report_parts['body_markdown']

        report_selector_update = gr.update(choices=report_choices, value=latest_report_id, interactive=True)

        # --- MODIFIED: Use the session cache for data reconstruction ---
        reconstructed_data, updated_cache = fetch_and_reconstruct_data_from_bubble(latest_report_series, session_cache)


        raw_results_state = None
        okr_details_md = "No OKRs found in the latest report." # This is for the old, hidden component
        enhanced_okr_html_content = get_initial_okr_display() # Default to loading state or empty
        actionable_okrs_dict_for_state = {} # NEW: Initialize the new state variable
        key_results_cbg_update = gr.update(choices=[], value=[], interactive=False)
        all_krs_state = []

        if reconstructed_data:
            raw_results_state = reconstructed_data
            actionable_okrs_dict = raw_results_state.get("actionable_okrs", {})

            # Refined check for actionable_okrs_dict content
            if actionable_okrs_dict and isinstance(actionable_okrs_dict, dict) and actionable_okrs_dict.get("okrs"):
                actionable_okrs_dict_for_state = actionable_okrs_dict # Store for the new state
                # Format for the new enhanced HTML display
                enhanced_okr_html_content = format_okrs_for_enhanced_display(raw_results_state)
            else:
                logger.info(f"No 'actionable_okrs' key or 'okrs' list found in reconstructed data for report ID {latest_report_id}. "
                            f"Content of 'actionable_okrs_dict': {actionable_okrs_dict}")
                enhanced_okr_html_content = get_initial_okr_display() # Show empty state if no OKRs
        else:
            logger.error(f"Failed to reconstruct data for latest report ID {latest_report_id}")
            okr_details_md = "Error: Could not reconstruct OKR data for this report."
            enhanced_okr_html_content = get_initial_okr_display() # Show empty state on reconstruction error

        status_update = f"Status: Loaded {len(agentic_df)} reports. Displaying latest from {latest_report_series['Created Date'].strftime('%Y-%m-%d')}."

        return (
            gr.update(value=status_update),              # 0: agentic_pipeline_status_md
            report_selector_update,                      # 1: report_selector_dd
            key_results_cbg_update,                      # 2: key_results_cbg (kept for compatibility)
            gr.update(value=okr_details_md),             # 3: okr_detail_display_md (kept for compatibility)
            gr.update(value=raw_results_state),          # 4: orchestration_raw_results_st
            gr.update(value=[]),                         # 5: selected_key_result_ids_st (kept for compatibility)
            gr.update(value=all_krs_state),              # 6: key_results_for_selection_st (kept for compatibility)
            gr.update(value=report_header_content),      # 7: report_header_html_display
            gr.update(value=report_body_content),        # 8: report_body_markdown_display
            gr.update(value=updated_cache),              # 9: reconstruction_cache_st
            gr.update(value=enhanced_okr_html_content),   # 10: NEW: The enhanced HTML display for OKRs
            gr.update(value=actionable_okrs_dict_for_state) # 11: NEW: Explicit actionable_okrs_data_st
        )

    except Exception as e:
        logger.error(f"Failed to process and display agentic results: {e}", exc_info=True)
        # Ensure error returns match the 12-item signature
        error_header_html = """
            <div class="report-title">⚠️ Error Loading Report</div>
            <div class="report-subtitle">An error occurred during data processing.</div>
            <div class="status-badge" style="background: #e74c3c;">Error</div>
            """
        error_body_markdown = f"""
            <div class="empty-state">
                <div class="empty-state-icon">🚨</div>
                <div class="empty-state-title">Report Loading Failed</div>
                <div class="empty-state-description">
                    An error occurred while loading or processing the report data: {e}.
                    Please try again or contact support if the issue persists.
                </div>
            </div>
            """
        return (
            gr.update(value=f"Status: An error occurred: {e}"), # 0
            gr.update(choices=[], value=None, interactive=False), # 1
            gr.update(choices=[], value=[], interactive=False), # 2
            gr.update(value="Error: Could not display OKRs due to an error."), # 3
            gr.update(value=None), # 4
            gr.update(value=[]), # 5
            gr.update(value=[]), # 6
            gr.update(value=error_header_html), # 7
            gr.update(value=error_body_markdown), # 8
            gr.update(value=session_cache), # 9
            gr.update(value=get_initial_okr_display()), # 10: Placeholder for enhanced OKR display on error
            gr.update(value={}) # 11: Placeholder for actionable_okrs_data_st on error
        )