"""AI research setup component for the API key manager."""

import streamlit as st
from loguru import logger
from typing import Dict, Any
from ..manager import APIKeyManager
import os
from dotenv import load_dotenv # Keep if api_key_manager uses it
import sys

# Configure logger (assuming configured elsewhere or keep minimal here)
logger.add(sys.stderr, level="INFO") # Keep simple example if needed

# Helper function to validate a specific research provider's key
def _validate_research_key(provider_name: str, key_value: str) -> bool:
    """Validate the API key for a given research provider."""
    if not key_value:
        logger.debug(f"Validation: Key for {provider_name} is empty.")
        return False
    try:
        logger.debug(f"Validating key for {provider_name}...")
        # Ensure the function exists in validation.py before calling
        if provider_name == "serpapi":
            if callable(getattr(sys.modules[__name__], 'test_serpapi_key', None)):
                is_valid = test_serpapi_key(key_value)
            else:
                 logger.error("test_serpapi_key not found in validation module")
                 is_valid = False
        elif provider_name == "tavily":
            if callable(getattr(sys.modules[__name__], 'test_tavily_key', None)):
                is_valid = test_tavily_key(key_value)
            else:
                 logger.error("test_tavily_key not found in validation module")
                 is_valid = False
        elif provider_name == "metaphor":
            if callable(getattr(sys.modules[__name__], 'test_metaphor_key', None)):
                is_valid = test_metaphor_key(key_value)
            else:
                 logger.error("test_metaphor_key not found in validation module")
                 is_valid = False
        elif provider_name == "firecrawl":
             if callable(getattr(sys.modules[__name__], 'test_firecrawl_key', None)):
                is_valid = test_firecrawl_key(key_value)
             else:
                 logger.error("test_firecrawl_key not found in validation module")
                 is_valid = False
        else:
            logger.warning(f"Validation not implemented for research provider: {provider_name}")
            return False # Default to False for unknown providers
        
        logger.info(f"Validation result for {provider_name}: {'Valid' if is_valid else 'Invalid'}")
        return is_valid
    except Exception as e:
        logger.error(f"Error validating key for {provider_name}: {e}", exc_info=True)
        return False

# Callback function for handling API key input changes
def _handle_research_key_change(provider_name: str, api_key_manager):
    """Save and validate research API key when input changes."""
    key_input_widget_key = f"{provider_name}_key_input"
    status_widget_key = f"{provider_name}_status"
    
    if key_input_widget_key not in st.session_state:
        logger.warning(f"Input widget key '{key_input_widget_key}' not found in session state.")
        return
        
    key_value = st.session_state[key_input_widget_key]
    current_status = st.session_state.get(status_widget_key)
    
    logger.debug(f"Handling research key change for {provider_name}. Key: {'***' if key_value else 'Empty'}. Current status: {current_status}")

    if not key_value:
        api_key_manager.save_api_key(provider_name, "")
        st.session_state[status_widget_key] = "unsaved"
        logger.info(f"Cleared API key for {provider_name}.")
        return

    st.session_state[status_widget_key] = "saving"
    st.rerun()
    
    try:
        logger.debug(f"Saving key for {provider_name}...")
        api_key_manager.save_api_key(provider_name, key_value)
        logger.info(f"Saved API key for {provider_name}.")

        is_valid = _validate_research_key(provider_name, key_value)

        if is_valid:
            st.session_state[status_widget_key] = "valid"
        else:
            st.session_state[status_widget_key] = "invalid"
            
    except Exception as e:
        logger.error(f"Error during saving/validation for {provider_name}: {e}", exc_info=True)
        st.session_state[status_widget_key] = "error"

def render_ai_research_setup(api_key_manager: APIKeyManager) -> Dict[str, Any]:
    """Render the AI research setup step with immediate feedback."""
    logger.info("[render_ai_research_setup] Rendering AI research setup component")
    
    research_providers = ["serpapi", "tavily", "metaphor", "firecrawl"]
    
    # Initialize statuses
    for provider in research_providers:
        status_key = f"{provider}_status"
        if status_key not in st.session_state:
            existing_key = api_key_manager.get_api_key(provider)
            if existing_key:
                if _validate_research_key(provider, existing_key):
                    st.session_state[status_key] = "valid"
                else:
                     st.session_state[status_key] = "invalid" 
            else:
                 st.session_state[status_key] = "unsaved"

    st.markdown("""
        <div class='setup-header'>
            <h2>Step 3: Configure AI Research Tools (Optional)</h2>
            <p>Set up API keys for enhanced web research, crawling, and analysis. These are optional but recommended.</p>
        </div>
    """, unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    
    # --- SerpAPI --- 
    with col1:
        st.subheader("SerpAPI")
        st.markdown("Access real-time search engine results. Get key: [SerpAPI](https://serpapi.com)")
        serpapi_key_val = api_key_manager.get_api_key("serpapi")
        st.text_input(
            "SerpAPI Key",
            value=serpapi_key_val if serpapi_key_val else "",
            type="password",
            key="serpapi_key_input",
            on_change=_handle_research_key_change,
            args=("serpapi", api_key_manager)
        )
        # Feedback Area
        serpapi_status = st.session_state.get("serpapi_status", "unsaved")
        feedback_placeholder_serpapi = st.empty()
        if serpapi_status == "saving":
            feedback_placeholder_serpapi.info("Validating SerpAPI key...", icon="⏳")
        elif serpapi_status == "valid":
            feedback_placeholder_serpapi.success("SerpAPI key saved and valid!", icon="✅")
        elif serpapi_status == "invalid":
            feedback_placeholder_serpapi.error("Invalid SerpAPI key.", icon="❌")
        elif serpapi_status == "error":
            feedback_placeholder_serpapi.error("Error saving/validating SerpAPI key.", icon="⚠️")

    # --- Firecrawl --- 
    with col1:
        st.subheader("Firecrawl")
        st.markdown("Web content extraction and crawling. Get key: [Firecrawl](https://www.firecrawl.dev/account)")
        firecrawl_key_val = api_key_manager.get_api_key("firecrawl")
        st.text_input(
            "Firecrawl API Key",
            value=firecrawl_key_val if firecrawl_key_val else "",
            type="password",
            key="firecrawl_key_input",
            on_change=_handle_research_key_change,
            args=("firecrawl", api_key_manager)
        )
        # Feedback Area
        firecrawl_status = st.session_state.get("firecrawl_status", "unsaved")
        feedback_placeholder_firecrawl = st.empty()
        if firecrawl_status == "saving":
            feedback_placeholder_firecrawl.info("Validating Firecrawl key...", icon="⏳")
        elif firecrawl_status == "valid":
            feedback_placeholder_firecrawl.success("Firecrawl key saved and valid!", icon="✅")
        elif firecrawl_status == "invalid":
            feedback_placeholder_firecrawl.error("Invalid Firecrawl key.", icon="❌")
        elif firecrawl_status == "error":
            feedback_placeholder_firecrawl.error("Error saving/validating Firecrawl key.", icon="⚠️")
            
    # --- Tavily --- 
    with col2:
        st.subheader("Tavily AI")
        st.markdown("AI-powered search & summarization. Get key: [Tavily](https://tavily.com)")
        tavily_key_val = api_key_manager.get_api_key("tavily")
        st.text_input(
            "Tavily API Key",
            value=tavily_key_val if tavily_key_val else "",
            type="password",
            key="tavily_key_input",
            on_change=_handle_research_key_change,
            args=("tavily", api_key_manager)
        )
        # Feedback Area
        tavily_status = st.session_state.get("tavily_status", "unsaved")
        feedback_placeholder_tavily = st.empty()
        if tavily_status == "saving":
            feedback_placeholder_tavily.info("Validating Tavily key...", icon="⏳")
        elif tavily_status == "valid":
            feedback_placeholder_tavily.success("Tavily key saved and valid!", icon="✅")
        elif tavily_status == "invalid":
            feedback_placeholder_tavily.error("Invalid Tavily key.", icon="❌")
        elif tavily_status == "error":
            feedback_placeholder_tavily.error("Error saving/validating Tavily key.", icon="⚠️")

    # --- Metaphor/Exa --- 
    with col2:
        st.subheader("Metaphor/Exa")
        st.markdown("Neural search for deep research. Get key: [Metaphor/Exa](https://metaphor.systems)")
        metaphor_key_val = api_key_manager.get_api_key("metaphor")
        st.text_input(
            "Metaphor/Exa API Key",
            value=metaphor_key_val if metaphor_key_val else "",
            type="password",
            key="metaphor_key_input",
            on_change=_handle_research_key_change,
            args=("metaphor", api_key_manager)
        )
        # Feedback Area
        metaphor_status = st.session_state.get("metaphor_status", "unsaved")
        feedback_placeholder_metaphor = st.empty()
        if metaphor_status == "saving":
            feedback_placeholder_metaphor.info("Validating Metaphor/Exa key...", icon="⏳")
        elif metaphor_status == "valid":
            feedback_placeholder_metaphor.success("Metaphor/Exa key saved and valid!", icon="✅")
        elif metaphor_status == "invalid":
            feedback_placeholder_metaphor.error("Invalid Metaphor/Exa key.", icon="❌")
        elif metaphor_status == "error":
            feedback_placeholder_metaphor.error("Error saving/validating Metaphor/Exa key.", icon="⚠️")

    # --- Coming Soon --- 
    with st.expander("🔜 Coming Soon - More Search Options", expanded=False):
        st.info("Integrations for Bing Search and Google Search APIs are planned.")

    st.info("You can skip this step if you don't need these research tools. Click Continue to proceed.")
        
    return {}