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

# --- MODIFIED: Import the reconstruction function ---
# This function is now central to fetching the detailed OKR data.
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. Report unavailable."
    def extract_key_results_for_selection(okrs_dict): return []
    def format_single_okr_for_display(okr_data, **kwargs): return "Agentic modules not loaded. OKR display unavailable."

logger = logging.getLogger(__name__)

def load_and_display_agentic_results(token_state: dict):
    """
    Loads pre-computed agentic results from the state, populates the report library dropdown,
    and displays the LATEST report and its fully reconstructed OKRs by default.

    This function is called on the initial application load.

    Args:
        token_state: The main state dictionary containing the bubble_agentic_analysis_data DataFrame.

    Returns:
        A tuple of Gradio updates matching the `agentic_display_outputs` list in `app.py`.
    """
    # Default empty/initial return values that match the output components list. The order is critical.
    initial_updates = (
        "No agentic analysis data found in Bubble.",          # 1. agentic_report_display_md
        gr.update(choices=[], value=None, interactive=False), # 2. report_selector_dd
        gr.update(choices=[], value=[], interactive=False),   # 3. key_results_cbg
        "No OKRs to display.",                                # 4. okr_detail_display_md
        None,                                                 # 5. orchestration_raw_results_st
        [],                                                   # 6. selected_key_result_ids_st
        [],                                                   # 7. key_results_for_selection_st
        "Status: No agentic analysis data found."             # 8. agentic_pipeline_status_md
    )

    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 in the state.")
        return initial_updates

    try:
        # --- 1. Prepare Report Library ---
        if 'Created Date' not in agentic_df.columns or '_id' not in agentic_df.columns:
            raise KeyError("Required columns ('Created Date', '_id') not found in agentic data.")
        
        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

        # --- 2. Process the Latest Report by Default ---
        latest_report_series = agentic_df.iloc[0]
        latest_report_id = latest_report_series['_id']
        
        # Format the main report text for display
        report_display_md = format_report_for_display(latest_report_series)
        
        # Update the report library dropdown and select the latest one
        report_selector_update = gr.update(choices=report_choices, value=latest_report_id, interactive=True)

        # --- 3. MODIFIED: Fetch, Reconstruct, and Prepare OKRs from the Latest Report ---
        raw_results_state = None
        okr_details_md = "No OKRs found in the latest report."
        key_results_cbg_update = gr.update(choices=[], value=[], interactive=False)
        all_krs_state = []
        
        # Call the key function to get all linked data (OKRs, KRs, Tasks) from Bubble
        reconstructed_data = fetch_and_reconstruct_data_from_bubble(latest_report_series)
        
        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:
                    # Populate the Key Results checkbox group
                    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)
                    
                    # Format all OKRs for the initial detailed display
                    okrs_list = actionable_okrs_dict.get("okrs", [])
                    output_md_parts = [
                        format_single_okr_for_display(okr_data, okr_main_index=okr_idx)
                        for okr_idx, okr_data 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 full 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 the latest from {latest_report_series['Created Date'].strftime('%Y-%m-%d')}."
        
        # --- 4. Return all UI and State Updates ---
        return (
            report_display_md,
            report_selector_update,
            key_results_cbg_update,
            okr_details_md,
            raw_results_state, # This is now the fully reconstructed dictionary
            [], # Reset selected KRs
            all_krs_state,
            status_update
        )

    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)