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,
        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
    def format_report_for_display(report_data): return "Agentic modules not loaded."
    def extract_key_results_for_selection(okrs_dict): return []
    def format_single_okr_for_display(okr_data, **kwargs): return "Agentic modules not loaded."

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.
    """
    initial_updates = (
        "No agentic analysis data found in Bubble.",
        gr.update(choices=[], value=None, interactive=False),
        gr.update(choices=[], value=[], interactive=False),
        "No OKRs to display.",
        None, [], [], "Status: No agentic analysis data found.",
        session_cache # Return the cache unchanged
    )

    if not AGENTIC_MODULES_LOADED:
        error_updates = list(initial_updates)
        error_updates[7] = "Status: Critical module import error."
        return tuple(error_updates)

    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()
        latest_report_series = quarterly_reports_df.iloc[0]
        latest_report_id = latest_report_series['_id']
        
        report_display_md = format_report_for_display(latest_report_series)
        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, okr_details_md = None, "No OKRs found in the latest report."
        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", {})
            if actionable_okrs_dict:
                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 okr_details_md
        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."

        status_update = f"Status: Loaded {len(agentic_df)} reports. Displaying latest from {latest_report_series['Created Date'].strftime('%Y-%m-%d')}."
        
        return (
            report_display_md, report_selector_update, key_results_cbg_update,
            okr_details_md, raw_results_state, [], all_krs_state, status_update,
            updated_cache # Return the potentially updated cache
        )

    except Exception as e:
        logger.error(f"Failed to process and display agentic results: {e}", exc_info=True)
        error_updates = list(initial_updates)
        error_updates[0] = f"An error occurred while loading reports: {e}"
        error_updates[7] = f"Status: Error - {e}"
        return tuple(error_updates)