app.py

# app.py
import os
import json
import re
import logging
import threading
import html # For escaping HTML in Gradio Markdown
from datetime import datetime
from dotenv import load_dotenv

import gradio as gr

# --- Load Environment Variables ---
load_dotenv()

# --- Local Logic Modules ---
from model_logic import (
    get_available_providers,
    get_model_display_names_for_provider,
    get_default_model_display_name_for_provider,
    call_model_stream,
    MODELS_BY_PROVIDER # Used for model selection logic
)
from memory_logic import (
    load_rules_from_file, save_rule_to_file, delete_rule_from_file, clear_all_rules, # For rules
    load_memories_from_file, save_memory_to_file, clear_all_memories                 # For memories
)
# Assuming websearch_logic.py contains scrape_url, search_and_scrape_duckduckgo, etc.
from websearch_logic import scrape_url, search_and_scrape_duckduckgo, search_and_scrape_google

# --- Logging Setup ---
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(threadName)s - %(message)s')
logger = logging.getLogger(__name__)
for lib_name in ["urllib3", "requests", "huggingface_hub", "PIL.PngImagePlugin", "matplotlib", "gradio_client.client", "multipart.multipart", "httpx"]:
    logging.getLogger(lib_name).setLevel(logging.WARNING)


# --- Application Configuration & Globals ---
WEB_SEARCH_ENABLED = os.getenv("WEB_SEARCH_ENABLED", "true").lower() == "true"
TOOL_DECISION_PROVIDER_ENV = os.getenv("TOOL_DECISION_PROVIDER", "groq")
TOOL_DECISION_MODEL_ID_ENV = os.getenv("TOOL_DECISION_MODEL", "llama3-8b-8192") # This is the model ID

MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", 7))
current_chat_session_history = [] # Stores conversation in OpenAI format: [{"role": ..., "content": ...}, ...]

DEFAULT_SYSTEM_PROMPT = os.getenv(
    "DEFAULT_SYSTEM_PROMPT",
    "You are a helpful AI research assistant. Your primary goal is to answer questions and perform research tasks accurately and thoroughly. You can use tools like web search and page browsing. When providing information from the web, cite your sources if possible. If asked to perform a task beyond your capabilities, explain politely. Be concise unless asked for detail."
)

logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={TOOL_DECISION_PROVIDER_ENV}, ToolDecisionModelID={TOOL_DECISION_MODEL_ID_ENV}")


# --- Helper Functions (from ai-learn, adapted) ---
def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str, list[dict]]:
    """
    Formats a list of insight strings (rules) into a structured prompt string
    and a list of parsed insight objects.
    """
    if not retrieved_insights_list:
        return "No specific guiding principles or learned insights retrieved.", []
    
    parsed = []
    for text in retrieved_insights_list:
        match = re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\](.*)", text.strip(), re.DOTALL | re.IGNORECASE)
        if match:
            parsed.append({
                "type": match.group(1).upper().replace(" ", "_"),
                "score": match.group(2),
                "text": match.group(3).strip(),
                "original": text.strip()
            })
        else: # Default if format is slightly off or just plain text
            parsed.append({
                "type": "GENERAL_LEARNING",
                "score": "0.5", # Default score
                "text": text.strip(),
                "original": text.strip()
            })
    
    try: # Sort by score, descending
        parsed.sort(key=lambda x: float(x["score"]) if x["score"].replace('.', '', 1).isdigit() else -1.0, reverse=True)
    except ValueError:
        logger.warning("FORMAT_INSIGHTS: Sort error due to invalid score format in an insight.")

    grouped = {"CORE_RULE": [], "RESPONSE_PRINCIPLE": [], "BEHAVIORAL_ADJUSTMENT": [], "GENERAL_LEARNING": []}
    for p_item in parsed:
        grouped.get(p_item["type"], grouped["GENERAL_LEARNING"]).append(f"- (Score: {p_item['score']}) {p_item['text']}")
    
    sections = [f"{k.replace('_', ' ').title()}:\n" + "\n".join(v) for k, v in grouped.items() if v]
    formatted_prompt_str = "\n\n".join(sections) if sections else "No guiding principles retrieved."
    return formatted_prompt_str, parsed


def retrieve_memories_simple_keywords(query: str, k=3) -> list[dict]:
    """Simple keyword-based memory retrieval from loaded memories."""
    all_memories = load_memories_from_file() # Assumes memory_logic.py handles loading
    if not query or not all_memories: return []
    
    query_terms = set(query.lower().split())
    if not query_terms: return []

    scored_memories = []
    for mem in all_memories:
        score = 0
        text_to_search = f"{mem.get('user_input','')} {mem.get('bot_response','')} {mem.get('metrics',{}).get('takeaway','')}".lower()
        for term in query_terms:
            if term in text_to_search:
                score += 1
        if score > 0:
            scored_memories.append({"score": score, "memory": mem})
    
    scored_memories.sort(key=lambda x: x["score"], reverse=True)
    logger.debug(f"Retrieved {len(scored_memories[:k])} memories with simple keyword search for query: '{query[:50]}...'")
    return [item["memory"] for item in scored_memories[:k]]

def retrieve_insights_simple_keywords(query: str, k_insights=5) -> list[str]:
    """Simple keyword-based insight/rule retrieval from loaded rules."""
    all_rules = load_rules_from_file() # Assumes memory_logic.py handles loading
    if not query or not all_rules: return []

    query_terms = set(query.lower().split())
    if not query_terms: return []

    scored_rules = []
    for rule_text in all_rules:
        score = 0
        match = re.match(r"\[.*?\](.*)", rule_text.strip()) # Extract text part of rule
        text_content_to_search = match.group(1).strip().lower() if match else rule_text.lower()
        for term in query_terms:
            if term in text_content_to_search:
                score += 1
        if score > 0:
            scored_rules.append({"score": score, "rule": rule_text})

    scored_rules.sort(key=lambda x: x["score"], reverse=True)
    logger.debug(f"Retrieved {len(scored_rules[:k_insights])} insights with simple keyword search for query: '{query[:50]}...'")
    return [item["rule"] for item in scored_rules[:k_insights]]


# --- Metrics and Learning ---
def generate_interaction_metrics(user_input: str, bot_response: str, provider: str, model_display_name: str, api_key_override: str = None) -> dict:
    metric_start_time = time.time()
    logger.info(f"Generating metrics with: {provider}/{model_display_name}")
    
    metric_prompt_content = f"User: \"{user_input}\"\nAI: \"{bot_response}\"\nMetrics: \"takeaway\" (3-7 words), \"response_success_score\" (0.0-1.0), \"future_confidence_score\" (0.0-1.0). Output JSON ONLY, ensure it's a single, valid JSON object."
    metric_messages = [{"role": "system", "content": "You are a precise JSON output agent. Output a single JSON object containing interaction metrics as requested by the user. Do not include any explanatory text before or after the JSON object."}, {"role": "user", "content": metric_prompt_content}]
    
    try:
        # Determine model for metrics (can be overridden by METRICS_MODEL env var)
        metrics_provider_final = provider
        metrics_model_display_final = model_display_name

        metrics_model_env = os.getenv("METRICS_MODEL") # Format: "provider_name/model_id"
        if metrics_model_env and "/" in metrics_model_env:
            m_prov, m_id = metrics_model_env.split('/', 1)
            m_disp_name = None
            # Find display name for the model_id under the provider
            for disp_name_candidate, model_id_candidate in MODELS_BY_PROVIDER.get(m_prov.lower(), {}).get("models", {}).items():
                if model_id_candidate == m_id:
                    m_disp_name = disp_name_candidate
                    break
            if m_disp_name:
                metrics_provider_final = m_prov
                metrics_model_display_final = m_disp_name
                logger.info(f"Overriding metrics model to: {metrics_provider_final}/{metrics_model_display_final} from METRICS_MODEL env var.")
            else:
                logger.warning(f"METRICS_MODEL '{metrics_model_env}' specified, but model ID '{m_id}' not found for provider '{m_prov}'. Using interaction model.")

        response_chunks = list(call_model_stream( # call_model_stream is a generator
            provider=metrics_provider_final,
            model_display_name=metrics_model_display_final,
            messages=metric_messages,
            api_key_override=api_key_override, # Pass along if main call used one
            temperature=0.05, # Low temp for precise JSON
            max_tokens=200  # Generous for JSON metrics
        ))
        resp_str = "".join(response_chunks).strip()

        # Extract JSON (robustly, as LLMs can add ```json ... ```)
        json_match_markdown = re.search(r"```json\s*(\{.*?\})\s*```", resp_str, re.DOTALL | re.IGNORECASE)
        json_match_direct = re.search(r"(\{.*?\})", resp_str, re.DOTALL)
        
        final_json_str = None
        if json_match_markdown: final_json_str = json_match_markdown.group(1)
        elif json_match_direct: final_json_str = json_match_direct.group(1)

        if final_json_str:
            metrics_data = json.loads(final_json_str)
        else:
            logger.warning(f"METRICS_GEN: Non-JSON or malformed JSON response from {metrics_provider_final}/{metrics_model_display_final}: '{resp_str}'")
            return {"takeaway": "N/A", "response_success_score": 0.5, "future_confidence_score": 0.5, "error": "metrics format error"}

        parsed_metrics = {
            "takeaway": metrics_data.get("takeaway", "N/A"),
            "response_success_score": float(metrics_data.get("response_success_score", 0.5)),
            "future_confidence_score": float(metrics_data.get("future_confidence_score", 0.5)),
            "error": metrics_data.get("error", None) # Allow None if no error
        }
        logger.info(f"METRICS_GEN: Metrics generated in {time.time() - metric_start_time:.2f}s. Data: {parsed_metrics}")
        return parsed_metrics

    except Exception as e:
        logger.error(f"METRICS_GEN Error: {e}", exc_info=False)
        return {"takeaway": "N/A", "response_success_score": 0.5, "future_confidence_score": 0.5, "error": str(e)}


# --- Core Interaction Processing ---
def process_user_interaction_gradio(user_input: str, provider_name: str, model_display_name: str, 
                                    chat_history_for_prompt: list[dict], custom_system_prompt: str = None, 
                                    ui_api_key_override: str = None):
    process_start_time = time.time()
    request_id = os.urandom(4).hex()
    logger.info(f"PUI_GRADIO [{request_id}] Start. User: '{user_input[:50]}...' Provider: {provider_name}/{model_display_name} Hist_len:{len(chat_history_for_prompt)}")

    # Prepare history string for prompts (limited turns)
    history_str_parts = []
    for t_msg in chat_history_for_prompt[-(MAX_HISTORY_TURNS * 2):]: # Use last N turns for prompt context
        role_display = "User" if t_msg['role'] == 'user' else "AI"
        history_str_parts.append(f"{role_display}: {t_msg['content']}")
    history_str_for_prompt = "\n".join(history_str_parts)

    yield "status", "<i>[Checking guidelines (learned rules)...]</i>"
    # Use simple keyword retrieval for insights/rules
    insights_query = f"{user_input}\n{history_str_for_prompt}" # Combine current query and history for context
    initial_insights = retrieve_insights_simple_keywords(insights_query, k_insights=5)
    initial_insights_ctx_str, parsed_initial_insights_list = format_insights_for_prompt(initial_insights)
    logger.info(f"PUI_GRADIO [{request_id}]: Initial RAG (insights) found {len(initial_insights)} items. Context preview: {initial_insights_ctx_str[:150]}...")

    action_type, action_input_dict = "quick_respond", {}
    user_input_lower = user_input.lower()

    # --- Tool/Action Decision Logic ---
    time_before_tool_decision = time.time()

    # 1. Heuristic: Direct URL for scraping
    if WEB_SEARCH_ENABLED and ("http://" in user_input or "https://" in user_input):
        url_match = re.search(r'(https?://[^\s]+)', user_input)
        if url_match:
            action_type = "scrape_url_and_report"
            action_input_dict = {"url": url_match.group(1)}
            logger.info(f"PUI_GRADIO [{request_id}]: Heuristic: URL detected. Action: {action_type}.")

    # 2. Heuristic: Simple keywords (if no URL matched)
    simple_keywords = ["hello", "hi", "hey", "thanks", "thank you", "ok", "okay", "yes", "no", "bye"]
    if action_type == "quick_respond" and len(user_input.split()) <= 3 and any(kw in user_input_lower for kw in simple_keywords) and not "?" in user_input:
        action_type = "quick_respond" # Already default, but explicit
        logger.info(f"PUI_GRADIO [{request_id}]: Heuristic: Simple keyword. Action: {action_type}.")
    
    # 3. LLM-based Tool Decision (if web search enabled and query seems to need it)
    needs_llm_decision = (WEB_SEARCH_ENABLED and (
        len(user_input.split()) > 3 or 
        "?" in user_input or 
        any(w in user_input_lower for w in ["what is", "how to", "explain", "search for", "find info", "who is", "why"])
    ))

    if action_type == "quick_respond" and needs_llm_decision:
        yield "status", "<i>[LLM choosing best approach...]</i>"
        
        # Prepare prompt for tool decision LLM
        history_snippet_for_tool = "\n".join([f"{msg['role']}: {msg['content'][:100]}" for msg in chat_history_for_prompt[-2:]]) # Short snippet
        guideline_snippet_for_tool = initial_insights_ctx_str[:200].replace('\n', ' ')

        tool_decision_sys_prompt = "You are a precise routing agent. Your task is to choose the single most appropriate action from the list to address the user's query. Output ONLY a single JSON object with 'action' and 'action_input' keys. Example: {\"action\": \"search_duckduckgo_and_report\", \"action_input\": {\"search_engine_query\": \"efficient LLM fine-tuning\"}}"
        tool_decision_user_prompt = f"""User Query: "{user_input}"
Recent Conversation Snippet:
{history_snippet_for_tool}
Key Guidelines (summary): {guideline_snippet_for_tool}...

Available Actions & Required Inputs:
1. `quick_respond`: For simple chat, greetings, or if no external info/memory is needed. (Input: N/A)
2. `answer_using_conversation_memory`: If the query refers to past specific details of THIS conversation. (Input: N/A)
3. `search_duckduckgo_and_report`: For general knowledge, facts, current events, or if user asks to search. (Input: `search_engine_query`: string)
4. `scrape_url_and_report`: If user explicitly provides a URL to summarize or analyze. (Input: `url`: string)

Select one action and its input. Output JSON only."""
        
        tool_decision_messages = [
            {"role":"system", "content": tool_decision_sys_prompt},
            {"role":"user", "content": tool_decision_user_prompt}
        ]
        
        # Determine tool decision LLM provider and model display name
        tool_provider = TOOL_DECISION_PROVIDER_ENV
        tool_model_id = TOOL_DECISION_MODEL_ID_ENV
        tool_model_display = None
        for disp_name, mod_id_val in MODELS_BY_PROVIDER.get(tool_provider.lower(), {}).get("models", {}).items():
            if mod_id_val == tool_model_id:
                tool_model_display = disp_name; break
        if not tool_model_display:
            tool_model_display = get_default_model_display_name_for_provider(tool_provider)
            logger.warning(f"Tool decision model ID '{tool_model_id}' not mapped for provider '{tool_provider}'. Using default: {tool_model_display}")

        if not tool_model_display:
            logger.error(f"Could not find any model for tool decision provider {tool_provider}. Defaulting to quick_respond.")
        else:
            try:
                logger.info(f"PUI_GRADIO [{request_id}]: Calling tool decision LLM: {tool_provider}/{tool_model_display}")
                tool_resp_chunks = list(call_model_stream(
                    provider=tool_provider, model_display_name=tool_model_display,
                    messages=tool_decision_messages, temperature=0.0, max_tokens=150
                ))
                tool_resp_raw = "".join(tool_resp_chunks).strip()
                
                json_match_tool = re.search(r"\{.*\}", tool_resp_raw, re.DOTALL)
                if json_match_tool:
                    action_data = json.loads(json_match_tool.group(0))
                    action_type = action_data.get("action", "quick_respond")
                    action_input_dict = action_data.get("action_input", {})
                    if not isinstance(action_input_dict, dict): action_input_dict = {} # Ensure it's a dict
                    logger.info(f"PUI_GRADIO [{request_id}]: LLM Tool Decision: Action='{action_type}', Input='{action_input_dict}'")
                else:
                    logger.warning(f"PUI_GRADIO [{request_id}]: Tool decision LLM non-JSON. Raw: {tool_resp_raw}")
            except Exception as e_tool_llm:
                logger.error(f"PUI_GRADIO [{request_id}]: Tool decision LLM error: {e_tool_llm}", exc_info=False)
    
    # 4. Fallback if web search disabled, consider memory for longer queries
    if action_type == "quick_respond" and not WEB_SEARCH_ENABLED:
         if len(user_input.split()) > 4 or "?" in user_input or any(w in user_input_lower for w in ["remember","recall"]): 
            action_type="answer_using_conversation_memory"
            logger.info(f"PUI_GRADIO [{request_id}]: Web search disabled, heuristic for memory. Action: {action_type}")

    logger.info(f"PUI_GRADIO [{request_id}]: Tool decision logic took {time.time() - time_before_tool_decision:.3f}s. Final Action: {action_type}, Input: {action_input_dict}")
    yield "status", f"<i>[Path: {action_type}. Preparing response...]</i>"

    # --- Action Execution & Prompt Construction for Main LLM ---
    final_system_prompt_str = custom_system_prompt or DEFAULT_SYSTEM_PROMPT
    final_user_prompt_content_str = ""
    scraped_content_for_prompt = "" # For web search results
    memory_context_for_prompt = ""  # For conversation memory

    if action_type == "quick_respond":
        final_system_prompt_str += " Respond directly to the user's query using the provided guidelines and conversation history for context. Be concise and helpful."
        final_user_prompt_content_str = f"Current Conversation History (User/AI turns):\n{history_str_for_prompt}\n\nGuiding Principles (Learned Rules):\n{initial_insights_ctx_str}\n\nUser's Current Query: \"{user_input}\"\n\nYour concise and helpful response:"

    elif action_type == "answer_using_conversation_memory":
        yield "status", "<i>[Searching conversation memory (simple keywords)...]</i>"
        mem_query_context = history_str_for_prompt[-1000:] # Last 1000 chars for context
        memory_query = f"User query: {user_input}\nRelevant conversation context:\n{mem_query_context}"
        retrieved_mems = retrieve_memories_simple_keywords(memory_query, k=2)
        
        if retrieved_mems:
            memory_context_for_prompt = "Relevant Past Interactions (for your reference only, prioritize current query):\n" + "\n".join(
                [f"- User: {m.get('user_input','')} -> AI: {m.get('bot_response','')} (Takeaway: {m.get('metrics',{}).get('takeaway','N/A')}, Timestamp: {m.get('timestamp','N/A')})" for m in retrieved_mems]
            )
        else:
            memory_context_for_prompt = "No highly relevant past interactions found in memory for this specific query."
        logger.info(f"PUI_GRADIO [{request_id}]: Memory retrieval found {len(retrieved_mems)} items.")
        
        final_system_prompt_str += " Respond by incorporating relevant information from 'Memory Context' and your general guidelines. Focus on the user's current query."
        final_user_prompt_content_str = f"Current Conversation History:\n{history_str_for_prompt}\n\nGuiding Principles:\n{initial_insights_ctx_str}\n\nMemory Context (from previous related interactions):\n{memory_context_for_prompt}\n\nUser's Current Query: \"{user_input}\"\n\nYour helpful response (draw from memory context if applicable, otherwise answer generally):"

    elif WEB_SEARCH_ENABLED and action_type in ["search_duckduckgo_and_report", "scrape_url_and_report"]: # "search_google_and_report" would be similar
        query_or_url = action_input_dict.get("search_engine_query") if "search" in action_type else action_input_dict.get("url")
        
        if not query_or_url:
            logger.warning(f"PUI_GRADIO [{request_id}]: Missing input for {action_type}. Falling back.")
            action_type = "quick_respond" # Fallback to quick_respond logic above
            final_system_prompt_str += " Respond directly. (Note: A web action was attempted but failed due to missing input)."
            final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
        else:
            yield "status", f"<i>[Web: '{query_or_url[:60]}'...]</i>"
            web_results_data = []
            max_web_results = 1 if action_type == "scrape_url_and_report" else 2

            try:
                if action_type == "search_duckduckgo_and_report":
                    web_results_data = search_and_scrape_duckduckgo(query_or_url, num_results=max_web_results)
                elif action_type == "scrape_url_and_report":
                    scrape_res = scrape_url(query_or_url)
                    if scrape_res and (scrape_res.get("content") or scrape_res.get("error")): web_results_data = [scrape_res]
                # Add elif for search_and_scrape_google if used
            except Exception as e_web_tool:
                logger.error(f"PUI_GRADIO [{request_id}]: Error during web tool '{action_type}': {e_web_tool}", exc_info=True)
                web_results_data = [{"url": query_or_url, "title": "Tool Execution Error", "content": None, "error": str(e_web_tool)}]

            if web_results_data:
                scraped_parts = []
                for i, r_item in enumerate(web_results_data):
                    content_item = r_item.get('content') or r_item.get('error') or 'N/A'
                    max_len_per_source = 3500 # Limit length per source
                    scraped_parts.append(f"Source {i+1}:\nURL: {r_item.get('url','N/A')}\nTitle: {r_item.get('title','N/A')}\nContent Snippet:\n{content_item[:max_len_per_source]}{'...' if len(content_item) > max_len_per_source else ''}\n---")
                scraped_content_for_prompt = "\n".join(scraped_parts)
            else:
                scraped_content_for_prompt = f"No results or content found from {action_type} for '{query_or_url}'."
            
            yield "status", "<i>[Synthesizing web report...]</i>"
            final_system_prompt_str += " Generate a report/answer from web content, history, & guidelines. Cite URLs as [Source X]."
            final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\n\nWeb Content Found:\n{scraped_content_for_prompt}\n\nUser's Query: \"{user_input}\"\n\nYour report/response (cite sources like [Source 1], [Source 2]):"
    
    else: # Fallback if action_type is somehow unknown
        logger.warning(f"PUI_GRADIO [{request_id}]: Unknown action_type '{action_type}'. Defaulting.")
        final_system_prompt_str += " Respond directly. (Internal state error)."
        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"


    # --- Final LLM Call for Response Generation ---
    final_llm_messages_for_api = []
    if final_system_prompt_str: final_llm_messages_for_api.append({"role": "system", "content": final_system_prompt_str})
    final_llm_messages_for_api.append({"role": "user", "content": final_user_prompt_content_str}) # The detailed user prompt

    # Log prompt being sent (truncated)
    logger.debug(f"PUI_GRADIO [{request_id}]: Final LLM System Prompt: {final_system_prompt_str[:200]}...")
    logger.debug(f"PUI_GRADIO [{request_id}]: Final LLM User Prompt (content part): {final_user_prompt_content_str[:200]}... {final_user_prompt_content_str[-200:] if len(final_user_prompt_content_str)>400 else ''}")

    streamed_response_accumulator = ""
    time_before_main_llm = time.time()
    try:
        response_iterator = call_model_stream(
            provider=provider_name,
            model_display_name=model_display_name,
            messages=final_llm_messages_for_api,
            api_key_override=ui_api_key_override,
            temperature=0.6,  # Can be adjusted
            max_tokens=2500   # Max tokens for the response
        )
        for chunk in response_iterator:
            if isinstance(chunk, str) and chunk.startswith("Error:"): # Check for errors from call_model_stream
                error_message = f"\n\nLLM API Error: {chunk}\n"
                streamed_response_accumulator += error_message
                yield "response_chunk", error_message
                logger.error(f"PUI_GRADIO [{request_id}]: Error from model stream: {chunk}")
                break 
            streamed_response_accumulator += chunk
            yield "response_chunk", chunk
            
    except Exception as e_final_llm:
        logger.error(f"PUI_GRADIO [{request_id}]: Final LLM call raised exception: {e_final_llm}", exc_info=True)
        error_chunk = f"\n\n(Error during final response generation: {str(e_final_llm)[:150]})"
        streamed_response_accumulator += error_chunk
        yield "response_chunk", error_chunk
    
    logger.info(f"PUI_GRADIO [{request_id}]: Main LLM stream took {time.time() - time_before_main_llm:.3f}s.")
    final_bot_response_text = streamed_response_accumulator.strip() or "(No response generated or error occurred.)"
    
    logger.info(f"PUI_GRADIO [{request_id}]: Processing finished. Total wall time: {time.time() - process_start_time:.2f}s. Response length: {len(final_bot_response_text)}")
    yield "final_response_and_insights", {"response": final_bot_response_text, "insights_used": parsed_initial_insights_list}


# --- Deferred Learning Task ---
def deferred_learning_and_memory_task(user_input: str, bot_response: str, provider_name: str, model_display_name: str, 
                                      parsed_insights_for_reflection: list[dict], ui_api_key_override: str = None):
    deferred_start_time = time.time()
    task_id = os.urandom(4).hex()
    logger.info(f"DEFERRED_LEARNING [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
    
    try:
        # 1. Generate Interaction Metrics
        metrics = generate_interaction_metrics(user_input, bot_response, provider_name, model_display_name, ui_api_key_override)
        logger.info(f"DEFERRED_LEARNING [{task_id}]: Metrics generated: {metrics}")

        # 2. Save memory entry (using memory_logic.py)
        save_memory_to_file(user_input, bot_response, metrics)

        # 3. Insight/Rule Generation (Reflection)
        summary_for_reflection = f"User:\"{user_input}\"\nAI:\"{bot_response}\"\nMetrics(takeaway):{metrics.get('takeaway','N/A')}, Success Score:{metrics.get('response_success_score','N/A')}"
        
        # Get context from existing rules (simple keyword retrieval)
        reflection_context_query = f"{summary_for_reflection}\n{user_input}" # Query for existing rules
        relevant_existing_rules = retrieve_insights_simple_keywords(reflection_context_query, k_insights=10) # Max 10 rules for context
        existing_rules_ctx_str = "\n".join([f"- \"{rule}\"" for rule in relevant_existing_rules]) if relevant_existing_rules else "No specific existing rules found as highly relevant for direct comparison."
        
        # Use the long system prompt for insight generation from ai-learn
        insight_sys_prompt = """You are an expert AI knowledge base curator... (Your full long prompt here from ai-learn's deferred_learning_and_memory)... Output ONLY the JSON list."""
        insight_user_prompt = f"""Interaction Summary:\n{summary_for_reflection}\n
Potentially Relevant Existing Rules (Review these carefully for consolidation or refinement):\n{existing_rules_ctx_str}\n
Guiding principles that were considered during THIS interaction (these might offer clues for new rules or refinements):\n{json.dumps([p['original'] for p in parsed_insights_for_reflection if 'original' in p]) if parsed_insights_for_reflection else "None"}\n
Task: Based on your reflection process... (Your full long task description here from ai-learn's deferred_learning_and_memory)... Output JSON only."""
        
        insight_gen_messages = [{"role": "system", "content": insight_sys_prompt}, {"role": "user", "content": insight_user_prompt}]

        # Determine model for insight generation (can be overridden by INSIGHT_MODEL_OVERRIDE env var)
        insight_provider_final = provider_name
        insight_model_display_final = model_display_name
        insight_model_env = os.getenv("INSIGHT_MODEL_OVERRIDE") # Format: "provider_name/model_id"
        if insight_model_env and "/" in insight_model_env:
            i_prov, i_id = insight_model_env.split('/', 1)
            i_disp_name = None
            for disp_name_candidate, model_id_candidate in MODELS_BY_PROVIDER.get(i_prov.lower(), {}).get("models", {}).items():
                if model_id_candidate == i_id: i_disp_name = disp_name_candidate; break
            if i_disp_name:
                insight_provider_final = i_prov
                insight_model_display_final = i_disp_name
                logger.info(f"Overriding insight generation model to: {insight_provider_final}/{insight_model_display_final} from INSIGHT_MODEL_OVERRIDE.")
        
        logger.info(f"DEFERRED_LEARNING [{task_id}]: Generating insights with {insight_provider_final}/{insight_model_display_final}")
        raw_insight_ops_chunks = list(call_model_stream(
            provider=insight_provider_final, model_display_name=insight_model_display_final,
            messages=insight_gen_messages, api_key_override=ui_api_key_override,
            temperature=0.05, max_tokens=2000 # Generous for multiple JSON operations
        ))
        raw_insight_ops_json_str = "".join(raw_insight_ops_chunks).strip()
        
        # Parse operations JSON
        operations = []
        json_ops_match_md = re.search(r"```json\s*(\[.*?\])\s*```", raw_insight_ops_json_str, re.DOTALL | re.IGNORECASE)
        json_ops_match_direct = re.search(r"(\[.*?\])", raw_insight_ops_json_str, re.DOTALL)
        
        final_ops_json_str = None
        if json_ops_match_md: final_ops_json_str = json_ops_match_md.group(1)
        elif json_ops_match_direct: final_ops_json_str = json_ops_match_direct.group(1)

        if final_ops_json_str:
            try: operations = json.loads(final_ops_json_str)
            except json.JSONDecodeError as e_json_ops:
                logger.error(f"DEFERRED_LEARNING [{task_id}]: JSONDecodeError for insight ops '{final_ops_json_str[:200]}...': {e_json_ops}")
        else: logger.warning(f"DEFERRED_LEARNING [{task_id}]: Insight LLM output not a JSON list: {raw_insight_ops_json_str[:200]}...")

        if not isinstance(operations, list):
            logger.warning(f"DEFERRED_LEARNING [{task_id}]: Parsed insight ops not a list. Type: {type(operations)}. No-op."); operations = []
        
        insights_processed_count = 0
        if operations:
            logger.info(f"DEFERRED_LEARNING [{task_id}]: LLM provided {len(operations)} insight operation(s).")
            for op_idx, op in enumerate(operations):
                if not isinstance(op, dict): continue
                action = op.get("action", "").strip().lower()
                insight_text = op.get("insight", "").strip()
                # Basic validation of insight_text format (e.g., starts with [TYPE|SCORE])
                if not insight_text or not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I):
                    logger.warning(f"DEFERRED_LEARNING [{task_id}]: Op {op_idx} has invalid insight format: '{insight_text[:70]}...'. Skip.")
                    continue
                
                if action == "add":
                    if save_rule_to_file(insight_text): insights_processed_count += 1
                elif action == "update":
                    old_insight_text_to_replace = op.get("old_insight_to_replace", "").strip()
                    if not old_insight_text_to_replace:
                        logger.warning(f"DEFERRED_LEARNING [{task_id}]: 'update' op {op_idx} missing 'old_insight_to_replace'. Attempting as 'add'.")
                        if save_rule_to_file(insight_text): insights_processed_count += 1
                    else:
                        if old_insight_text_to_replace == insight_text:
                            logger.info(f"DEFERRED_LEARNING [{task_id}]: Update op {op_idx} has identical old/new. Skip.")
                            continue
                        delete_rule_from_file(old_insight_text_to_replace) # Best effort delete
                        if save_rule_to_file(insight_text): insights_processed_count += 1
                time.sleep(0.01) # Small delay between file operations
            logger.info(f"DEFERRED_LEARNING [{task_id}]: Finished processing. Insights added/updated: {insights_processed_count}")
        else:
            logger.info(f"DEFERRED_LEARNING [{task_id}]: No insight operations proposed by LLM or parsing failed.")

    except Exception as e_deferred:
        logger.error(f"DEFERRED_LEARNING [{task_id}]: CRITICAL ERROR in deferred task: {e_deferred}", exc_info=True)
    logger.info(f"DEFERRED_LEARNING [{task_id}]: END. Total time: {time.time() - deferred_start_time:.2f}s")


# --- Gradio Chat Handler ---
def handle_gradio_chat_submit(user_message_text: str,
                              gradio_chat_history_list: list[tuple[str | None, str | None]],
                              selected_provider_name: str,
                              selected_model_display_name: str,
                              ui_api_key_text: str | None,
                              custom_system_prompt_text: str):
    
    # Initialize UI update variables
    cleared_input_text = "" # To clear the user input box
    updated_gradio_history = list(gradio_chat_history_list) # Copy current display history
    current_status_text = "Initializing..."
    # Default values for output components to yield immediately
    # These should match the types of the output components in demo.launch
    # Use dummy values that match the component types if needed
    default_detected_outputs_md = gr.Markdown(value="*Processing...*")
    default_formatted_output_text = gr.Textbox(value="*Waiting for AI response...*")
    default_download_button = gr.DownloadButton(interactive=False, value=None, visible=False)

    if not user_message_text.strip():
        current_status_text = "Error: Cannot send an empty message."
        # Add error to Gradio display history
        updated_gradio_history.append((user_message_text or "(Empty)", current_status_text))
        yield (cleared_input_text, updated_gradio_history, current_status_text, 
               default_detected_outputs_md, default_formatted_output_text, default_download_button)
        return

    # Add user message to Gradio display history with a thinking placeholder
    updated_gradio_history.append((user_message_text, "<i>Thinking...</i>"))
    yield (cleared_input_text, updated_gradio_history, current_status_text,
           default_detected_outputs_md, default_formatted_output_text, default_download_button)

    # Prepare history for internal processing (OpenAI format)
    # current_chat_session_history is the global list storing conversation in OpenAI format
    internal_processing_history = list(current_chat_session_history)
    internal_processing_history.append({"role": "user", "content": user_message_text})
    
    # Truncate internal_processing_history if too long (maintain MAX_HISTORY_TURNS)
    if len(internal_processing_history) > (MAX_HISTORY_TURNS * 2 + 1): # +1 for potential system prompt
        # Simple truncation from the beginning, preserving last N turns
        # More sophisticated: keep system prompt if present, then truncate older user/assistant pairs
        internal_processing_history = internal_processing_history[-(MAX_HISTORY_TURNS * 2):]


    final_bot_response_text_accumulated = ""
    parsed_insights_used_in_response = [] # List of insight dicts

    try:
        # Call the core processing generator
        interaction_processor_gen = process_user_interaction_gradio(
            user_input=user_message_text,
            provider_name=selected_provider_name,
            model_display_name=selected_model_display_name,
            chat_history_for_prompt=internal_processing_history, # Pass the internal history
            custom_system_prompt=custom_system_prompt_text.strip() if custom_system_prompt_text else None,
            ui_api_key_override=ui_api_key_text.strip() if ui_api_key_text else None
        )

        # Stream updates to Gradio UI
        current_bot_display_message = ""
        for update_type, update_data in interaction_processor_gen:
            if update_type == "status":
                current_status_text = update_data
                # Update last bot message in Gradio history with status
                if updated_gradio_history and updated_gradio_history[-1][0] == user_message_text:
                    updated_gradio_history[-1] = (user_message_text, f"{current_bot_display_message} <i>{current_status_text}</i>" if current_bot_display_message else f"<i>{current_status_text}</i>")
            
            elif update_type == "response_chunk":
                current_bot_display_message += update_data
                if updated_gradio_history and updated_gradio_history[-1][0] == user_message_text:
                    updated_gradio_history[-1] = (user_message_text, current_bot_display_message)
            
            elif update_type == "final_response_and_insights":
                final_bot_response_text_accumulated = update_data["response"]
                parsed_insights_used_in_response = update_data["insights_used"]
                
                current_status_text = "Response complete."
                if not current_bot_display_message and final_bot_response_text_accumulated : # If no chunks streamed (e.g. error or very short non-streamed response)
                    current_bot_display_message = final_bot_response_text_accumulated

                if updated_gradio_history and updated_gradio_history[-1][0] == user_message_text:
                     updated_gradio_history[-1] = (user_message_text, current_bot_display_message or "(No textual response)")
                
                # Update the dedicated report display area
                default_formatted_output_text = gr.Textbox(value=current_bot_display_message)
                
                # Update download button
                if current_bot_display_message and not current_bot_display_message.startswith("Error:"):
                    report_filename = f"ai_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md" # Use markdown for rich text
                    default_download_button = gr.DownloadButton(label="Download Report (.md)", value=current_bot_display_message, filename=report_filename, visible=True, interactive=True)
                else:
                    default_download_button = gr.DownloadButton(interactive=False, value=None, visible=False)
                
                # Update detected outputs preview (e.g., insights used)
                insights_preview_md = "### Insights Considered During Response:\n"
                if parsed_insights_used_in_response:
                    for insight_obj in parsed_insights_used_in_response[:3]: # Show top 3
                        insights_preview_md += f"- **[{insight_obj.get('type','N/A')}|{insight_obj.get('score','N/A')}]** {insight_obj.get('text','N/A')[:100]}...\n"
                else: insights_preview_md += "*No specific learned insights were retrieved as highly relevant for this query.*"
                default_detected_outputs_md = gr.Markdown(value=insights_preview_md)

            # Yield all UI components that need updating
            yield (cleared_input_text, updated_gradio_history, current_status_text,
                   default_detected_outputs_md, default_formatted_output_text, default_download_button)
            
            if update_type == "final_response_and_insights":
                break # Generator finished for this interaction

    except Exception as e_handler:
        logger.error(f"Error in Gradio chat handler: {e_handler}", exc_info=True)
        current_status_text = f"Error: {str(e_handler)[:100]}"
        if updated_gradio_history and updated_gradio_history[-1][0] == user_message_text:
            updated_gradio_history[-1] = (user_message_text, current_status_text)
        else: # Should not happen if placeholder was added
             updated_gradio_history.append((user_message_text, current_status_text))
        yield (cleared_input_text, updated_gradio_history, current_status_text,
               default_detected_outputs_md, default_formatted_output_text, default_download_button)
        return # Exit on error

    # After successful response generation & streaming:
    if final_bot_response_text_accumulated and not final_bot_response_text_accumulated.startswith("Error:"):
        # Update the global internal chat history (OpenAI format)
        current_chat_session_history.append({"role": "user", "content": user_message_text})
        current_chat_session_history.append({"role": "assistant", "content": final_bot_response_text_accumulated})
        
        # Trim global history if it exceeds max turns
        if len(current_chat_session_history) > (MAX_HISTORY_TURNS * 2):
            current_chat_session_history = current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]

        # Start deferred learning task in a background thread
        logger.info(f"Starting deferred learning task for user: '{user_message_text[:40]}...'")
        deferred_thread = threading.Thread(
            target=deferred_learning_and_memory_task,
            args=(user_message_text, final_bot_response_text_accumulated,
                  selected_provider_name, selected_model_display_name,
                  parsed_insights_used_in_response,
                  ui_api_key_text.strip() if ui_api_key_text else None),
            daemon=True # Exits when main program exits
        )
        deferred_thread.start()
        current_status_text = "Response complete. Background learning initiated."
    else:
        current_status_text = "Processing finished, but no final response or an error occurred."

    # Final yield to ensure UI reflects the very last status
    yield (cleared_input_text, updated_gradio_history, current_status_text,
           default_detected_outputs_md, default_formatted_output_text, default_download_button)


# --- Gradio UI Helper Functions for Memory/Rules (File-based) ---
def ui_view_rules_action():
    rules_list = load_rules_from_file()
    if not rules_list: return "No rules/insights learned or stored yet."
    # Format for display in TextArea, one rule per line or separated by '---'
    return "\n\n---\n\n".join(rules_list)

def ui_upload_rules_action(uploaded_file_obj, progress=gr.Progress()):
    if not uploaded_file_obj: return "No file provided for rules upload."
    
    try:
        content = uploaded_file_obj.decode('utf-8') # Gradio File component gives bytes
    except AttributeError: # If it's already a string (e.g. from temp file path)
        try:
            with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f: # .name if it's a temp file object
                content = f.read()
        except Exception as e_read:
            logger.error(f"Error reading uploaded rules file: {e_read}")
            return f"Error reading file: {e_read}"
    except Exception as e_decode:
        logger.error(f"Error decoding uploaded rules file: {e_decode}")
        return f"Error decoding file content: {e_decode}"

    if not content.strip(): return "Uploaded rules file is empty."

    added_count, skipped_count, error_count = 0, 0, 0
    # Try splitting by '---' first, then by newline if that yields only one item
    potential_rules = content.split("\n\n---\n\n")
    if len(potential_rules) == 1 and "\n" in content: # Fallback to newline delimiter
        potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
    
    total_to_process = len(potential_rules)
    progress(0, desc="Starting rules upload...")

    for idx, rule_text in enumerate(potential_rules):
        rule_text = rule_text.strip()
        if not rule_text: continue
        
        # memory_logic.save_rule_to_file handles validation and duplicate checks
        # We need a more nuanced return from save_rule_to_file to distinguish reasons for not saving
        # For now, let's assume True means added, False means not added (any reason)
        # A better save_rule_to_file could return: "added", "duplicate", "invalid_format", "error"
        
        # Re-check for existing before trying to save for more accurate "skipped_count"
        existing_rules_snapshot = load_rules_from_file() # Could be slow if called repeatedly
        if rule_text in existing_rules_snapshot:
            skipped_count +=1
        elif save_rule_to_file(rule_text): # save_rule_to_file will log its own errors/skips
            added_count += 1
        else: # Failed to save for other reasons (format error logged by save_rule_to_file, or file write error)
            error_count += 1
        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} rules...")

    msg = f"Rules Upload: Processed {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors/Not Added: {error_count}."
    logger.info(msg)
    return msg


def ui_view_memories_action():
    memories_list_of_dicts = load_memories_from_file()
    if not memories_list_of_dicts: return [] # gr.JSON expects a list or dict
    return memories_list_of_dicts

def ui_upload_memories_action(uploaded_file_obj, progress=gr.Progress()):
    if not uploaded_file_obj: return "No file provided for memories upload."

    try:
        content = uploaded_file_obj.decode('utf-8')
    except AttributeError:
        try:
            with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f:
                content = f.read()
        except Exception as e_read:
            logger.error(f"Error reading uploaded memories file: {e_read}")
            return f"Error reading file: {e_read}"
    except Exception as e_decode:
        logger.error(f"Error decoding uploaded memories file: {e_decode}")
        return f"Error decoding file content: {e_decode}"

    if not content.strip(): return "Uploaded memories file is empty."

    added_count, format_error_count, save_error_count = 0, 0, 0
    memory_objects_to_process = []

    try: # Attempt to parse as a single JSON list first
        parsed_json = json.loads(content)
        if isinstance(parsed_json, list):
            memory_objects_to_process = parsed_json
        else: # If it's a single object, wrap it in a list
            memory_objects_to_process = [parsed_json]
    except json.JSONDecodeError: # If not a single JSON list, try JSON Lines
        logger.info("Failed to parse memories as single JSON list, trying JSON Lines format.")
        for line in content.splitlines():
            line = line.strip()
            if line:
                try:
                    mem_obj = json.loads(line)
                    memory_objects_to_process.append(mem_obj)
                except json.JSONDecodeError:
                    logger.warning(f"Skipping malformed JSON line in memories upload: {line[:100]}")
                    format_error_count += 1
    
    if not memory_objects_to_process and format_error_count == 0: # No objects found and no parsing errors
        return "No valid memory objects found in the uploaded file."

    total_to_process = len(memory_objects_to_process)
    progress(0, desc="Starting memories upload...")

    for idx, mem_data in enumerate(memory_objects_to_process):
        if not isinstance(mem_data, dict) or not all(k in mem_data for k in ["user_input", "bot_response", "metrics"]): # Timestamp optional for upload
            format_error_count += 1
            continue
        
        # For file-based, duplicate check on save might be too slow.
        # memory_logic.save_memory_to_file just appends.
        if save_memory_to_file(mem_data["user_input"], mem_data["bot_response"], mem_data["metrics"]):
            added_count += 1
        else:
            save_error_count += 1 # Error during file write
        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} memories...")

    msg = f"Memories Upload: Processed {total_to_process}. Added: {added_count}, Format Errors: {format_error_count}, Save Errors: {save_error_count}."
    logger.info(msg)
    return msg


# --- Gradio UI Definition ---
custom_theme = gr.themes.Base(primary_hue="teal", secondary_hue="purple", neutral_hue="zinc", text_size="sm", spacing_size="sm", radius_size="sm", font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"])
custom_css = """
body { font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 96% !important; margin: auto !important; padding-top: 1rem !important; }
footer { display: none !important; }
.gr-button { white-space: nowrap; }
.gr-input, .gr-textarea textarea, .gr-dropdown input { border-radius: 8px !important; }
.gr-chatbot .message { border-radius: 10px !important; box-shadow: 0 2px 5px rgba(0,0,0,0.08) !important; }
.prose { h1 { font-size: 1.8rem; margin-bottom: 0.6em; margin-top: 0.8em; } h2 { font-size: 1.4rem; margin-bottom: 0.5em; margin-top: 0.7em; } h3 { font-size: 1.15rem; margin-bottom: 0.4em; margin-top: 0.6em; } p { margin-bottom: 0.8em; line-height: 1.65; } ul, ol { margin-left: 1.5em; margin-bottom: 0.8em; } code { background-color: #f1f5f9; padding: 0.2em 0.45em; border-radius: 4px; font-size: 0.9em; } pre > code { display: block; padding: 0.8em; overflow-x: auto; background-color: #f8fafc; border: 1px solid #e2e8f0; border-radius: 6px;}}
.compact-group .gr-input-label, .compact-group .gr-dropdown-label { font-size: 0.8rem !important; padding-bottom: 2px !important;} /* Example to make labels smaller in a group */
"""

with gr.Blocks(theme=custom_theme, css=custom_css, title="AI Research Mega Agent v3") as demo:
    gr.Markdown("# 🚀 AI Research Mega Agent (Dynamic Models & File Memory)", elem_classes="prose")
    
    # --- Provider and Model Selection ---
    available_providers_list = get_available_providers()
    default_provider = available_providers_list[0] if available_providers_list else None
    default_models_for_provider = get_model_display_names_for_provider(default_provider) if default_provider else []
    default_model_for_provider = get_default_model_display_name_for_provider(default_provider) if default_provider else None

    with gr.Row():
        with gr.Column(scale=1, min_width=320): # Sidebar
            gr.Markdown("## ⚙️ Configuration", elem_classes="prose")
            
            with gr.Accordion("API & Model Settings", open=True):
                with gr.Group(elem_classes="compact-group"):
                    gr.Markdown("### LLM Provider & Model", elem_classes="prose")
                    provider_select_dd = gr.Dropdown(
                        label="Select LLM Provider", choices=available_providers_list, 
                        value=default_provider, interactive=True
                    )
                    model_select_dd = gr.Dropdown(
                        label="Select Model", choices=default_models_for_provider,
                        value=default_model_for_provider, interactive=True
                    )
                    api_key_override_tb = gr.Textbox(
                        label="API Key Override (Optional)", type="password",
                        placeholder="Enter API key for selected provider",
                        info="Overrides .env if provided here for the session."
                    )
                with gr.Group(elem_classes="compact-group"):
                    gr.Markdown("### System Prompt (Optional)", elem_classes="prose")
                    system_prompt_tb = gr.Textbox(
                        label="Custom System Prompt Base", lines=6, value=DEFAULT_SYSTEM_PROMPT,
                        interactive=True, info="Base prompt for the AI. Internal logic may add more context."
                    )

            with gr.Accordion("Knowledge Management (File-based)", open=False):
                gr.Markdown("### Rules (Learned Insights)", elem_classes="prose")
                view_rules_btn = gr.Button("View All Rules")
                upload_rules_file_obj = gr.File(label="Upload Rules File (.txt or .jsonl)", file_types=[".txt", ".jsonl"], scale=2)
                rules_status_tb = gr.Textbox(label="Rules Action Status", interactive=False, lines=2)
                clear_rules_btn = gr.Button("⚠️ Clear All Rules", variant="stop")
                
                gr.Markdown("### Memories (Past Interactions)", elem_classes="prose")
                view_memories_btn = gr.Button("View All Memories")
                upload_memories_file_obj = gr.File(label="Upload Memories File (.jsonl)", file_types=[".jsonl"], scale=2)
                memories_status_tb = gr.Textbox(label="Memories Action Status", interactive=False, lines=2)
                clear_memories_btn = gr.Button("⚠️ Clear All Memories", variant="stop")

        with gr.Column(scale=3): # Main Chat Area
            gr.Markdown("## 💬 AI Research Assistant Chat", elem_classes="prose")
            main_chatbot_display = gr.Chatbot(
                label="AI Research Chat", height=650, bubble_full_width=False,
                avatar_images=(None, "https://raw.githubusercontent.com/huggingface/brand-assets/main/hf-logo-with-title.png"), # Example bot avatar
                show_copy_button=True, render_markdown=True, sanitize_html=True, elem_id="main_chatbot"
            )
            with gr.Row():
                user_message_tb = gr.Textbox(
                    show_label=False, placeholder="Ask your research question or give an instruction...",
                    scale=7, lines=1, max_lines=5, autofocus=True, elem_id="user_message_input"
                )
                send_chat_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
            
            agent_status_tb = gr.Textbox(label="Agent Status", interactive=False, lines=1, value="Ready. Initializing systems...")
            
            with gr.Tabs():
                with gr.TabItem("📝 Generated Report/Output"):
                    gr.Markdown("The AI's full response or generated report will appear here.", elem_classes="prose")
                    formatted_report_tb = gr.Textbox(label="Current Research Output", lines=20, interactive=True, show_copy_button=True, value="*AI responses will appear here...*")
                    download_report_btn = gr.DownloadButton(label="Download Report", interactive=False, visible=False)
                
                with gr.TabItem("🔍 Intermediate Details / Data Viewer"):
                    gr.Markdown("View intermediate details, loaded data, or debug information.", elem_classes="prose")
                    detected_outputs_md_display = gr.Markdown(value="*Insights used or other intermediate details will show here...*")
                    gr.HTML("<hr style='margin: 1em 0;'>") # Separator
                    gr.Markdown("### Current Rules Viewer", elem_classes="prose")
                    rules_display_ta = gr.TextArea(label="Loaded Rules/Insights (Snapshot)", lines=10, interactive=False, max_lines=20)
                    gr.HTML("<hr style='margin: 1em 0;'>")
                    gr.Markdown("### Current Memories Viewer", elem_classes="prose")
                    memories_display_json_viewer = gr.JSON(label="Loaded Memories (Snapshot)")

    # --- Event Handlers ---
    
    # Update model dropdown when provider changes
    def dynamic_update_model_dropdown(selected_provider_name_dyn: str):
        models_for_provider_dyn = get_model_display_names_for_provider(selected_provider_name_dyn)
        default_model_dyn = get_default_model_display_name_for_provider(selected_provider_name_dyn)
        return gr.Dropdown(choices=models_for_provider_dyn, value=default_model_dyn, interactive=True)
    
    provider_select_dd.change(fn=dynamic_update_model_dropdown, inputs=provider_select_dd, outputs=model_select_dd)

    # Chat submission
    chat_inputs_list = [
        user_message_tb, main_chatbot_display,
        provider_select_dd, model_select_dd, api_key_override_tb,
        system_prompt_tb
    ]
    chat_outputs_list = [
        user_message_tb, main_chatbot_display, agent_status_tb,
        detected_outputs_md_display, formatted_report_tb, download_report_btn
    ]
    
    send_chat_btn.click(fn=handle_gradio_chat_submit, inputs=chat_inputs_list, outputs=chat_outputs_list)
    user_message_tb.submit(fn=handle_gradio_chat_submit, inputs=chat_inputs_list, outputs=chat_outputs_list)

    # Rules/Insights Management
    view_rules_btn.click(fn=ui_view_rules_action, outputs=rules_display_ta)
    upload_rules_file_obj.upload(fn=ui_upload_rules_action, inputs=[upload_rules_file_obj], outputs=[rules_status_tb], show_progress="full").then(
        fn=ui_view_rules_action, outputs=rules_display_ta # Refresh view after upload
    )
    clear_rules_btn.click(fn=lambda: "All rules files cleared." if clear_all_rules() else "Error clearing rules files.", outputs=rules_status_tb).then(
        fn=ui_view_rules_action, outputs=rules_display_ta # Refresh view
    )

    # Memories Management
    view_memories_btn.click(fn=ui_view_memories_action, outputs=memories_display_json_viewer)
    upload_memories_file_obj.upload(fn=ui_upload_memories_action, inputs=[upload_memories_file_obj], outputs=[memories_status_tb], show_progress="full").then(
        fn=ui_view_memories_action, outputs=memories_display_json_viewer # Refresh view after upload
    )
    clear_memories_btn.click(fn=lambda: "All memory files cleared." if clear_all_memories() else "Error clearing memory files.", outputs=memories_status_tb).then(
        fn=ui_view_memories_action, outputs=memories_display_json_viewer # Refresh view
    )
    
    # Initial status update on app load
    def app_load_init_status():
        # memory_logic.py creates DATA_DIR if not exists.
        # No complex loading like FAISS, so just confirm ready.
        logger.info("App loaded. File-based memory system is active.")
        return "AI Systems Initialized. Using File Memory. Ready."
    
    demo.load(fn=app_load_init_status, inputs=None, outputs=agent_status_tb)


# --- Main Application Execution ---
if __name__ == "__main__":
    logger.info("Starting Gradio AI Research Mega Agent Application (v3)...")
    
    # memory_logic.py handles its own directory creation.
    # No explicit data loading into globals needed here if handlers load on demand.

    app_port = int(os.getenv("GRADIO_PORT", 7860))
    app_server_name = os.getenv("GRADIO_SERVER_NAME", "0.0.0.0")
    app_debug_mode = os.getenv("GRADIO_DEBUG", "False").lower() == "true"
    app_share_mode = os.getenv("GRADIO_SHARE", "False").lower() == "true"

    logger.info(f"Launching Gradio server on http://{app_server_name}:{app_port}. Debug: {app_debug_mode}, Share: {app_share_mode}")
    
    # .queue() is important for streaming and handling multiple users
    demo.queue().launch(
        server_name=app_server_name,
        server_port=app_port,
        debug=app_debug_mode,
        share=app_share_mode,
        # prevent_thread_lock=True # May not be needed, test without first
        # auth=("user", "password") # Example for basic auth
    )
    logger.info("Gradio application has been shut down.")