app.py

import os
# --- CONFIGURATION TOGGLES ---
# Set these values to configure the application's behavior.

# Set to True to disable destructive actions (clearing all data, saving edited rules, and all uploads).
DEMO_MODE = False
# Select the storage backend: "HF_DATASET", "SQLITE", or "RAM".
# This will override the .env file setting for STORAGE_BACKEND.
MEMORY_STORAGE_TYPE = "RAM" 

# If using HF_DATASET, specify the repository names here.
# These will override the .env file settings.
HF_DATASET_MEMORY_REPO = "broadfield-dev/ai-brain"
HF_DATASET_RULES_REPO = "broadfield-dev/ai-rules"

# Set environment variables based on the toggles above BEFORE importing other modules
os.environ['STORAGE_BACKEND'] = MEMORY_STORAGE_TYPE
if MEMORY_STORAGE_TYPE == "HF_DATASET":
    os.environ['HF_MEMORY_DATASET_REPO'] = HF_DATASET_MEMORY_REPO
    os.environ['HF_RULES_DATASET_REPO'] = HF_DATASET_RULES_REPO
# --- END CONFIGURATION ---


import json
import re
import logging
from datetime import datetime
from dotenv import load_dotenv
import gradio as gr
import time
import tempfile
import xml.etree.ElementTree as ET

load_dotenv() # Load .env file, but our settings above will take precedence if set.

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
)

    
from memory_logic import (
    initialize_memory_system,
    add_memory_entry, retrieve_memories_semantic, get_all_memories_cached, clear_all_memory_data_backend,
    add_rule_entry, retrieve_rules_semantic, remove_rule_entry, get_all_rules_cached, clear_all_rules_data_backend,
    save_faiss_indices_to_disk, STORAGE_BACKEND as MEMORY_STORAGE_BACKEND, SQLITE_DB_PATH as MEMORY_SQLITE_PATH,
    HF_MEMORY_DATASET_REPO as MEMORY_HF_MEM_REPO, HF_RULES_DATASET_REPO as MEMORY_HF_RULES_REPO
)
from websearch_logic import scrape_url, search_and_scrape_duckduckgo, search_and_scrape_google

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", "sentence_transformers", "faiss", "datasets"]:
    if logging.getLogger(lib_name): logging.getLogger(lib_name).setLevel(logging.WARNING)

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")
MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", 7))
current_chat_session_history = []

DEFAULT_SYSTEM_PROMPT = os.getenv(
    "DEFAULT_SYSTEM_PROMPT",
    "Your Name is Node. You are a Helpful AI Assistant, and your goal is to improve yourself, so that you can improve everybody elses life. Your job 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."
)

# --- Startup Loading File Paths ---
LOAD_RULES_FILE = os.getenv("LOAD_RULES_FILE")
LOAD_MEMORIES_FILE = os.getenv("LOAD_MEMORIES_FILE")
logger.info(f"App Config: WebSearch={WEB_SEARCH_ENABLED}, ToolDecisionProvider={TOOL_DECISION_PROVIDER_ENV}, ToolDecisionModelID={TOOL_DECISION_MODEL_ID_ENV}, MemoryBackend={MEMORY_STORAGE_BACKEND}")
logger.info(f"Startup loading: Rules from {LOAD_RULES_FILE or 'None'}, Memories from {LOAD_MEMORIES_FILE or 'None'}")


# --- Helper Functions ---
def format_insights_for_prompt(retrieved_insights_list: list[str]) -> tuple[str, list[dict]]:
    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:
            parsed.append({"type": "GENERAL_LEARNING", "score": "0.5", "text": text.strip(), "original": text.strip()})
    try:
        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.")
    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]
    return "\n\n".join(sections) if sections else "No guiding principles retrieved.", parsed

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:
        metrics_provider_final, metrics_model_display_final = provider, model_display_name
        metrics_model_env = os.getenv("METRICS_MODEL")
        if metrics_model_env and "/" in metrics_model_env:
            m_prov, m_id = metrics_model_env.split('/', 1)
            m_disp_name = next((dn for dn, mid in MODELS_BY_PROVIDER.get(m_prov.lower(), {}).get("models", {}).items() if mid == m_id), None)
            if m_disp_name: metrics_provider_final, metrics_model_display_final = m_prov, m_disp_name
            else: logger.warning(f"METRICS_MODEL '{metrics_model_env}' not found, using interaction model.")
        response_chunks = list(call_model_stream(provider=metrics_provider_final, model_display_name=metrics_model_display_final, messages=metric_messages, api_key_override=api_key_override, temperature=0.05, max_tokens=200))
        resp_str = "".join(response_chunks).strip()
        json_match = re.search(r"```json\s*(\{.*?\})\s*```", resp_str, re.DOTALL | re.IGNORECASE) or re.search(r"(\{.*?\})", resp_str, re.DOTALL)
        if json_match: metrics_data = json.loads(json_match.group(1))
        else:
            logger.warning(f"METRICS_GEN: Non-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")}
        logger.info(f"METRICS_GEN: 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)}


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)}")
    history_str_for_prompt = "\n".join([f"{('User' if t_msg['role'] == 'user' else 'AI')}: {t_msg['content']}" for t_msg in chat_history_for_prompt[-(MAX_HISTORY_TURNS * 2):]])
    yield "status", "<i>[Checking guidelines (semantic search)...]</i>"
    initial_insights = retrieve_rules_semantic(f"{user_input}\n{history_str_for_prompt}", k=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)}. Context: {initial_insights_ctx_str[:150]}...")
    action_type, action_input_dict = "quick_respond", {}
    user_input_lower = user_input.lower()
    time_before_tool_decision = time.time()
    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, action_input_dict = "scrape_url_and_report", {"url": url_match.group(1)}
    if action_type == "quick_respond" and len(user_input.split()) <= 3 and any(kw in user_input_lower for kw in ["hello", "hi", "thanks", "ok", "bye"]) and not "?" in user_input: pass
    elif action_type == "quick_respond" and 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"])):
        yield "status", "<i>[LLM choosing best approach...]</i>"
        history_snippet = "\n".join([f"{msg['role']}: {msg['content'][:100]}" for msg in chat_history_for_prompt[-2:]])
        guideline_snippet = initial_insights_ctx_str[:200].replace('\n', ' ')
        tool_sys_prompt = "You are a precise routing agent... Output JSON only. Example: {\"action\": \"search_duckduckgo_and_report\", \"action_input\": {\"search_engine_query\": \"query\"}}"
        tool_user_prompt = f"User Query: \"{user_input}\nRecent History:\n{history_snippet}\nGuidelines: {guideline_snippet}...\nAvailable Actions: quick_respond, answer_using_conversation_memory, search_duckduckgo_and_report, scrape_url_and_report.\nSelect one action and input. Output JSON."
        tool_decision_messages = [{"role":"system", "content": tool_sys_prompt}, {"role":"user", "content": tool_user_prompt}]
        tool_provider, tool_model_id = TOOL_DECISION_PROVIDER_ENV, TOOL_DECISION_MODEL_ID_ENV
        tool_model_display = next((dn for dn, mid in MODELS_BY_PROVIDER.get(tool_provider.lower(), {}).get("models", {}).items() if mid == tool_model_id), None)
        if not tool_model_display: tool_model_display = get_default_model_display_name_for_provider(tool_provider)
        if tool_model_display:
            try:
                logger.info(f"PUI_GRADIO [{request_id}]: 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_input_dict = action_data.get("action", "quick_respond"), action_data.get("action_input", {})
                    if not isinstance(action_input_dict, dict): action_input_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: logger.error(f"PUI_GRADIO [{request_id}]: Tool decision LLM error: {e}", exc_info=False)
        else: logger.error(f"No model for tool decision provider {tool_provider}.")
    elif action_type == "quick_respond" and not WEB_SEARCH_ENABLED and (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}]: Tool decision logic took {time.time() - time_before_tool_decision:.3f}s. Action: {action_type}, Input: {action_input_dict}")
    yield "status", f"<i>[Path: {action_type}. Preparing response...]</i>"
    final_system_prompt_str, final_user_prompt_content_str = custom_system_prompt or DEFAULT_SYSTEM_PROMPT, ""
    if action_type == "quick_respond":
        final_system_prompt_str += " Respond directly using guidelines & history."
        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
    elif action_type == "answer_using_conversation_memory":
        yield "status", "<i>[Searching conversation memory (semantic)...]</i>"
        retrieved_mems = retrieve_memories_semantic(f"User query: {user_input}\nContext:\n{history_str_for_prompt[-1000:]}", k=2)
        memory_context = "Relevant Past Interactions:\n" + "\n".join([f"- User:{m.get('user_input','')}->AI:{m.get('bot_response','')} (Takeaway:{m.get('metrics',{}).get('takeaway','N/A')})" for m in retrieved_mems]) if retrieved_mems else "No relevant past interactions found."
        final_system_prompt_str += " Respond using Memory Context, guidelines, & history."
        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nMemory Context:\n{memory_context}\nQuery: \"{user_input}\"\nResponse (use memory context if relevant):"
    elif WEB_SEARCH_ENABLED and action_type in ["search_duckduckgo_and_report", "scrape_url_and_report"]:
        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:
            final_system_prompt_str += " Respond directly (web action failed: no 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, max_results = [], 1 if action_type == "scrape_url_and_report" else 2
            try:
                if action_type == "search_duckduckgo_and_report": web_results = search_and_scrape_duckduckgo(query_or_url, num_results=max_results)
                elif action_type == "scrape_url_and_report":
                    res = scrape_url(query_or_url)
                    if res and (res.get("content") or res.get("error")): web_results = [res]
            except Exception as e: web_results = [{"url": query_or_url, "title": "Tool Error", "error": str(e)}]
            scraped_content = "\n".join([f"Source {i+1}:\nURL:{r.get('url','N/A')}\nTitle:{r.get('title','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:3500]}\n---" for i,r in enumerate(web_results)]) if web_results else f"No results from {action_type} for '{query_or_url}'."
            yield "status", "<i>[Synthesizing web report...]</i>"
            final_system_prompt_str += " Generate 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}\nWeb Content:\n{scraped_content}\nQuery: \"{user_input}\"\nReport/Response (cite sources [Source X]):"
    else:
        final_system_prompt_str += " Respond directly (unknown action path)."
        final_user_prompt_content_str = f"History:\n{history_str_for_prompt}\nGuidelines:\n{initial_insights_ctx_str}\nQuery: \"{user_input}\"\nResponse:"
    final_llm_messages = [{"role": "system", "content": final_system_prompt_str}, {"role": "user", "content": final_user_prompt_content_str}]
    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 Start: {final_user_prompt_content_str[:200]}...")
    streamed_response, time_before_llm = "", time.time()
    try:
        for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_llm_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=2500):
            if isinstance(chunk, str) and chunk.startswith("Error:"): streamed_response += f"\n{chunk}\n"; yield "response_chunk", f"\n{chunk}\n"; break
            streamed_response += chunk; yield "response_chunk", chunk
    except Exception as e: streamed_response += f"\n\n(Error: {str(e)[:150]})"; yield "response_chunk", f"\n\n(Error: {str(e)[:150]})"
    logger.info(f"PUI_GRADIO [{request_id}]: Main LLM stream took {time.time() - time_before_llm:.3f}s.")
    final_bot_text = streamed_response.strip() or "(No response or error.)"
    logger.info(f"PUI_GRADIO [{request_id}]: Finished. Total: {time.time() - process_start_time:.2f}s. Resp len: {len(final_bot_text)}")
    yield "final_response_and_insights", {"response": final_bot_text, "insights_used": parsed_initial_insights_list}

def perform_post_interaction_learning(user_input: str, bot_response: str, provider: str, model_disp_name: str, insights_reflected: list[dict], api_key_override: str = None):
    task_id = os.urandom(4).hex()
    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: START User='{user_input[:40]}...', Bot='{bot_response[:40]}...'")
    learning_start_time = time.time()
    significant_learnings_summary = [] # To store summaries of new core learnings

    try:
        metrics = generate_interaction_metrics(user_input, bot_response, provider, model_disp_name, api_key_override)
        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Metrics: {metrics}")
        add_memory_entry(user_input, metrics, bot_response)

        summary = f"User:\"{user_input}\"\nAI:\"{bot_response}\"\nMetrics(takeaway):{metrics.get('takeaway','N/A')},Success:{metrics.get('response_success_score','N/A')}"
        existing_rules_ctx = "\n".join([f"- \"{r}\"" for r in retrieve_rules_semantic(f"{summary}\n{user_input}", k=10)]) or "No existing rules context."

        insight_sys_prompt = """You are an expert AI knowledge base curator. Your primary function is to meticulously analyze an interaction and update the AI's guiding principles (insights/rules) to improve its future performance and self-understanding.
**CRITICAL OUTPUT REQUIREMENT: You MUST output a single, valid XML structure representing a list of operation objects.**
The root element should be `<operations_list>`. Each operation should be an `<operation>` element.
If no operations are warranted, output an empty list: `<operations_list></operations_list>`.
ABSOLUTELY NO other text, explanations, or markdown should precede or follow this XML structure.
Each `<operation>` element must contain the following child elements:
1.  `<action>`: A string, either `"add"` (for entirely new rules) or `"update"` (to replace an existing rule with a better one).
2.  `<insight>`: The full, refined insight text including its `[TYPE|SCORE]` prefix (e.g., `[CORE_RULE|1.0] My name is Lumina, an AI assistant.`). Multi-line insight text can be placed directly within this tag; XML handles newlines naturally.
3.  `<old_insight_to_replace>`: (ONLY for `"update"` action) The *exact, full text* of an existing insight that the new `<insight>` should replace. If action is `"add"`, this element should be omitted or empty.
**XML Structure Example:**
<operations_list>
  <operation>
    <action>update</action>
    <insight>[CORE_RULE|1.0] I am Lumina, an AI assistant.
My purpose is to help with research.</insight>
    <old_insight_to_replace>[CORE_RULE|0.9] My name is Assistant.</old_insight_to_replace>
  </operation>
  <operation>
    <action>add</action>
    <insight>[RESPONSE_PRINCIPLE|0.8] User prefers short answers.
Provide details only when asked.</insight>
  </operation>
</operations_list>
**Your Reflection Process (Consider each step and generate operations accordingly):**
- **STEP 1: CORE IDENTITY/PURPOSE:** Review the interaction and existing rules. Identify if the interaction conflicts with, clarifies, or reinforces your core identity (name, fundamental nature, primary purpose). If necessary, propose updates or additions to CORE_RULEs. Aim for a single, consistent set of CORE_RULEs over time by updating older versions.
- **STEP 2: NEW LEARNINGS:** Based *only* on the "Interaction Summary", identify concrete, factual information, user preferences, or skills demonstrated that were not previously known or captured. These should be distinct, actionable learnings. Formulate these as new [GENERAL_LEARNING] or specific [BEHAVIORAL_ADJUSTMENT] rules. Do NOT add rules that are already covered by existing relevant rules.
- **STEP 3: REFINEMENT/ADJUSTMENT:** Review existing non-core rules ([RESPONSE_PRINCIPLE], [BEHAVIORAL_ADJUSTMENT], [GENERAL_LEARNING]) retrieved as "Potentially Relevant Existing Rules". Determine if the interaction indicates any of these rules need refinement, adjustment, or correction. Update existing rules if a better version exists.
**General Guidelines for Insight Content and Actions:**
- Ensure the `<insight>` field always contains the properly formatted insight string: `[TYPE|SCORE] Text`.
- Be precise with `<old_insight_to_replace>` – it must *exactly* match an existing rule string.
- Aim for a comprehensive set of operations.
"""
        insight_user_prompt = f"""Interaction Summary:\n{summary}\n
Potentially Relevant Existing Rules (Review these carefully. Your main goal is to consolidate CORE_RULEs and then identify other changes/additions based on the Interaction Summary and these existing rules):\n{existing_rules_ctx}\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 insights_reflected if 'original' in p]) if insights_reflected else "None"}\n
Task: Based on your three-step reflection process (Core Identity, New Learnings, Refinements):
1.  **Consolidate CORE_RULEs:** Merge similar identity/purpose rules from "Potentially Relevant Existing Rules" into single, definitive statements using "update" operations. Replace multiple old versions with the new canonical one.
2.  **Add New Learnings:** Identify and "add" any distinct new facts, skills, or important user preferences learned from the "Interaction Summary".
3.  **Update Existing Principles:** "Update" any non-core principles from "Potentially Relevant Existing Rules" if the "Interaction Summary" provided a clear refinement.
Combine all findings into a single, valid XML structure as specified in the system prompt (root `<operations_list>`, with child `<operation>` elements). Output XML ONLY.
"""
        insight_msgs = [{"role":"system", "content":insight_sys_prompt}, {"role":"user", "content":insight_user_prompt}]
        insight_prov, insight_model_disp = provider, model_disp_name
        insight_env_model = os.getenv("INSIGHT_MODEL_OVERRIDE")
        if insight_env_model and "/" in insight_env_model:
            i_p, i_id = insight_env_model.split('/', 1)
            i_d_n = next((dn for dn, mid in MODELS_BY_PROVIDER.get(i_p.lower(), {}).get("models", {}).items() if mid == i_id), None)
            if i_d_n: insight_prov, insight_model_disp = i_p, i_d_n
        logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Generating insights with {insight_prov}/{insight_model_disp} (expecting XML)")

        raw_ops_xml_full = "".join(list(call_model_stream(provider=insight_prov, model_display_name=insight_model_disp, messages=insight_msgs, api_key_override=api_key_override, temperature=0.0, max_tokens=3500))).strip()

        ops_data_list, processed_count = [], 0

        xml_match = re.search(r"```xml\s*(<operations_list>.*</operations_list>)\s*```", raw_ops_xml_full, re.DOTALL | re.IGNORECASE) or \
                    re.search(r"(<operations_list>.*</operations_list>)", raw_ops_xml_full, re.DOTALL | re.IGNORECASE)

        if xml_match:
            xml_content_str = xml_match.group(1)
            try:
                root = ET.fromstring(xml_content_str)
                if root.tag == "operations_list":
                    for op_element in root.findall("operation"):
                        action_el = op_element.find("action")
                        insight_el = op_element.find("insight")
                        old_insight_el = op_element.find("old_insight_to_replace")

                        action = action_el.text.strip().lower() if action_el is not None and action_el.text else None
                        insight_text = insight_el.text.strip() if insight_el is not None and insight_el.text else None
                        old_insight_text = old_insight_el.text.strip() if old_insight_el is not None and old_insight_el.text else None

                        if action and insight_text:
                            ops_data_list.append({
                                "action": action,
                                "insight": insight_text,
                                "old_insight_to_replace": old_insight_text
                            })
                        else:
                            logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Skipped XML operation due to missing action or insight text. Action: {action}, Insight: {insight_text}")
                else:
                    logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML root tag is not <operations_list>. Found: {root.tag}. XML content:\n{xml_content_str}")
            except ET.ParseError as e:
                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: XML parsing error: {e}. XML content that failed:\n{xml_content_str}")
            except Exception as e_xml_proc:
                logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: Error processing parsed XML: {e_xml_proc}. XML content:\n{xml_content_str}")
        else:
            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No <operations_list> XML structure found in LLM output. Full raw output:\n{raw_ops_xml_full}")

        if ops_data_list:
            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: LLM provided {len(ops_data_list)} insight ops from XML.")
            for op_idx, op_data in enumerate(ops_data_list):
                action = op_data["action"]
                insight_text = op_data["insight"]
                old_insight = op_data["old_insight_to_replace"]

                if not re.match(r"\[(CORE_RULE|RESPONSE_PRINCIPLE|BEHAVIORAL_ADJUSTMENT|GENERAL_LEARNING)\|([\d\.]+?)\]", insight_text, re.I|re.DOTALL):
                    logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to invalid insight_text format from XML: '{insight_text[:100]}...'")
                    continue

                rule_added_or_updated = False
                if action == "add":
                    success, status_msg = add_rule_entry(insight_text)
                    if success:
                        processed_count +=1
                        rule_added_or_updated = True
                        if insight_text.upper().startswith("[CORE_RULE"):
                            significant_learnings_summary.append(f"New Core Rule Added: {insight_text}")
                    else: logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (add from XML): Failed to add rule '{insight_text[:50]}...'. Status: {status_msg}")
                elif action == "update":
                    removed_old = False
                    if old_insight:
                        if old_insight != insight_text:
                            remove_success = remove_rule_entry(old_insight)
                            if not remove_success:
                                 logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Failed to remove old rule '{old_insight[:50]}...' before adding new.")
                            else:
                                removed_old = True
                        else:
                             logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Old insight is identical to new insight. Skipping removal.")

                    success, status_msg = add_rule_entry(insight_text)
                    if success:
                        processed_count +=1
                        rule_added_or_updated = True
                        if insight_text.upper().startswith("[CORE_RULE"):
                             significant_learnings_summary.append(f"Core Rule Updated (Old: {'Removed' if removed_old else 'Not removed/Same'}, New: {insight_text})")
                    else: logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx} (update from XML): Failed to add/update rule '{insight_text[:50]}...'. Status: {status_msg}")
                else:
                    logger.warning(f"POST_INTERACTION_LEARNING [{task_id}]: Op {op_idx}: Skipped op due to unknown action '{action}' from XML.")

            # After processing all rules, if there were significant learnings, add a special memory
            if significant_learnings_summary:
                learning_digest = "SYSTEM CORE LEARNING DIGEST:\n" + "\n".join(significant_learnings_summary)
                # Create a synthetic metrics object for this system memory
                system_metrics = {
                    "takeaway": "Core knowledge refined.",
                    "response_success_score": 1.0, # Assuming successful internal update
                    "future_confidence_score": 1.0,
                    "type": "SYSTEM_REFLECTION"
                }
                add_memory_entry(
                    user_input="SYSTEM_INTERNAL_REFLECTION_TRIGGER", # Fixed identifier
                    metrics=system_metrics,
                    bot_response=learning_digest
                )
                logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Added CORE_LEARNING_DIGEST to memories: {learning_digest[:100]}...")

            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: Processed {processed_count} insight ops out of {len(ops_data_list)} received from XML.")
        else:
            logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: No valid insight operations derived from LLM's XML output.")

    except Exception as e: logger.error(f"POST_INTERACTION_LEARNING [{task_id}]: CRITICAL ERROR in learning task: {e}", exc_info=True)
    logger.info(f"POST_INTERACTION_LEARNING [{task_id}]: END. Total: {time.time() - learning_start_time:.2f}s")


def handle_gradio_chat_submit(user_msg_txt: str, gr_hist_list: list, sel_prov_name: str, sel_model_disp_name: str, ui_api_key: str|None, cust_sys_prompt: str):
    global current_chat_session_history
    cleared_input, updated_gr_hist, status_txt = "", list(gr_hist_list), "Initializing..."
    # Initialize all potential output components with default/current values
    updated_rules_text = ui_refresh_rules_display_fn() # Get current rules state
    updated_mems_json = ui_refresh_memories_display_fn() # Get current memories state
    def_detect_out_md = gr.Markdown(visible=False)
    def_fmt_out_txt = gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True)
    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)

    if not user_msg_txt.strip():
        status_txt = "Error: Empty message."
        updated_gr_hist.append((user_msg_txt or "(Empty)", status_txt))
        # Ensure all outputs are provided on early exit
        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)
        return

    updated_gr_hist.append((user_msg_txt, "<i>Thinking...</i>"))
    # Initial yield to update chat UI with thinking message and show current knowledge base state
    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)

    internal_hist = list(current_chat_session_history); internal_hist.append({"role": "user", "content": user_msg_txt})
    # History truncation logic (keep MAX_HISTORY_TURNS pairs + optional system prompt)
    hist_len_check = MAX_HISTORY_TURNS * 2
    if internal_hist and internal_hist[0]["role"] == "system": hist_len_check +=1
    if len(internal_hist) > hist_len_check:
         current_chat_session_history = ([internal_hist[0]] if internal_hist[0]["role"] == "system" else []) + internal_hist[-(MAX_HISTORY_TURNS * 2):]
         internal_hist = list(current_chat_session_history) # Use truncated history for current turn processing

    final_bot_resp_acc, insights_used_parsed = "", []
    temp_dl_file_path = None

    try:
        processor_gen = process_user_interaction_gradio(user_input=user_msg_txt, provider_name=sel_prov_name, model_display_name=sel_model_disp_name, chat_history_for_prompt=internal_hist, custom_system_prompt=cust_sys_prompt.strip() or None, ui_api_key_override=ui_api_key.strip() if ui_api_key else None)
        curr_bot_disp_msg = ""
        for upd_type, upd_data in processor_gen:
            if upd_type == "status":
                status_txt = upd_data
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                    # Update the status alongside the streaming message
                    updated_gr_hist[-1] = (user_msg_txt, f"{curr_bot_disp_msg} <i>{status_txt}</i>" if curr_bot_disp_msg else f"<i>{status_txt}</i>")
            elif upd_type == "response_chunk":
                curr_bot_disp_msg += upd_data
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg) # Update chat with streamed chunk
            elif upd_type == "final_response_and_insights":
                final_bot_resp_acc, insights_used_parsed = upd_data["response"], upd_data["insights_used"]
                status_txt = "Response generated. Processing learning..."
                # Ensure the final chat message reflects the full response
                if not curr_bot_disp_msg and final_bot_resp_acc : curr_bot_disp_msg = final_bot_resp_acc
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                    updated_gr_hist[-1] = (user_msg_txt, curr_bot_disp_msg or "(No text)")

                # Update detailed response box and download button
                def_fmt_out_txt = gr.Textbox(value=curr_bot_disp_msg, interactive=True, show_copy_button=True)

                if curr_bot_disp_msg and not curr_bot_disp_msg.startswith("Error:"):
                    try:
                        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
                            tmpfile.write(curr_bot_disp_msg)
                            temp_dl_file_path = tmpfile.name
                        def_dl_btn = gr.DownloadButton(value=temp_dl_file_path, visible=True, interactive=True)
                    except Exception as e:
                        logger.error(f"Error creating temp file for download: {e}", exc_info=False)
                        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False, label="Download Error")
                else:
                    def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)

                # Update insights display
                insights_md_content = "### Insights Considered (Pre-Response):\n" + ("\n".join([f"- **[{i.get('type','N/A')}|{i.get('score','N/A')}]** {i.get('text','N/A')[:100]}..." for i in insights_used_parsed[:3]]) if insights_used_parsed else "*None specific.*")
                def_detect_out_md = gr.Markdown(value=insights_md_content, visible=True if insights_used_parsed else False)

            # Yield intermediate updates for the UI
            # Pass the *current* state of rules and memories display components
            yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)

            # Stop processing generator after final_response_and_insights
            if upd_type == "final_response_and_insights": break

    except Exception as e:
        logger.error(f"Chat handler error during main processing: {e}", exc_info=True); status_txt = f"Error: {str(e)[:100]}"
        error_message_for_chat = f"Sorry, an error occurred during response generation: {str(e)[:100]}"
        if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
            updated_gr_hist[-1] = (user_msg_txt, error_message_for_chat)
        else:
            updated_gr_hist.append((user_msg_txt, error_message_for_chat))
        def_fmt_out_txt = gr.Textbox(value=error_message_for_chat, interactive=True)
        def_dl_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
        def_detect_out_md = gr.Markdown(value="*Error processing request.*", visible=True)

        # Provide the current state of rules/memories on error path yield
        current_rules_text_on_error = ui_refresh_rules_display_fn()
        current_mems_json_on_error = ui_refresh_memories_display_fn()

        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_on_error, current_mems_json_on_error)
        # Clean up temp file if created before error
        if temp_dl_file_path and os.path.exists(temp_dl_file_path):
            try: os.unlink(temp_dl_file_path)
            except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path} after error: {e_unlink}")
        return # Exit the function after error handling

    # --- Post-Interaction Learning ---
    if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
        # Add the successful turn to the internal history
        current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
        # History truncation again after adding
        hist_len_check = MAX_HISTORY_TURNS * 2
        if current_chat_session_history and current_chat_session_history[0]["role"] == "system": hist_len_check +=1
        if len(current_chat_session_history) > hist_len_check:
            current_chat_session_history = ([current_chat_session_history[0]] if current_chat_session_history[0]["role"] == "system" else []) + current_chat_session_history[-(MAX_HISTORY_TURNS * 2):]

        status_txt = "<i>[Performing post-interaction learning...]</i>"
        # Yield status before synchronous learning
        current_rules_text_before_learn = ui_refresh_rules_display_fn()
        current_mems_json_before_learn = ui_refresh_memories_display_fn()
        yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, current_rules_text_before_learn, current_mems_json_before_learn)

        try:
            perform_post_interaction_learning(
                user_input=user_msg_txt,
                bot_response=final_bot_resp_acc,
                provider=sel_prov_name,
                model_disp_name=sel_model_disp_name,
                insights_reflected=insights_used_parsed,
                api_key_override=ui_api_key.strip() if ui_api_key else None
            )
            status_txt = "Response & Learning Complete."
        except Exception as e_learn:
            logger.error(f"Error during post-interaction learning: {e_learn}", exc_info=True)
            status_txt = "Response complete. Error during learning."

    elif final_bot_resp_acc.startswith("Error:"):
        status_txt = final_bot_resp_acc
        # If it was an error response from the generator, it's already in updated_gr_hist[-1]
        # The other output components (fmt_report_tb, dl_btn, detect_out_md) are already set by the generator loop or default state
    else:
        status_txt = "Processing finished; no valid response or error occurred during main phase."


    # Final yield after learning (or error handling)
    # This final yield updates the UI one last time with the true final status
    # AND crucially refreshes the Rules and Memories displays in case they changed during learning.
    updated_rules_text = ui_refresh_rules_display_fn()
    updated_mems_json = ui_refresh_memories_display_fn()

    yield (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)

    # Clean up the temporary download file after the final yield
    if temp_dl_file_path and os.path.exists(temp_dl_file_path):
        try: os.unlink(temp_dl_file_path)
        except Exception as e_unlink: logger.error(f"Error deleting temp download file {temp_dl_file_path}: {e_unlink}")


# --- Startup Loading Functions ---
def load_rules_from_file(filepath: str | None):
    """Loads rules from a local file (.txt or .jsonl) and adds them to the system."""
    if not filepath:
        logger.info("LOAD_RULES_FILE environment variable not set. Skipping rules loading from file.")
        return 0, 0, 0 # added, skipped, errors

    if not os.path.exists(filepath):
        logger.warning(f"LOAD_RULES: Specified rules file not found: {filepath}. Skipping loading.")
        return 0, 0, 0

    added_count, skipped_count, error_count = 0, 0, 0
    potential_rules = []

    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        logger.error(f"LOAD_RULES: Error reading file {filepath}: {e}", exc_info=False)
        return 0, 0, 1 # Indicate read error

    if not content.strip():
        logger.info(f"LOAD_RULES: File {filepath} is empty. Skipping loading.")
        return 0, 0, 0

    file_name_lower = filepath.lower()

    if file_name_lower.endswith(".txt"):
        potential_rules = content.split("\n\n---\n\n")
        # Also handle simple line breaks if '---' separator is not used
        if len(potential_rules) == 1 and "\n" in content:
             potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
    elif file_name_lower.endswith(".jsonl"):
        for line_num, line in enumerate(content.splitlines()):
            line = line.strip()
            if line:
                try:
                    # Expecting each line to be a JSON string containing the rule text
                    rule_text_in_json_string = json.loads(line)
                    if isinstance(rule_text_in_json_string, str):
                        potential_rules.append(rule_text_in_json_string)
                    else:
                        logger.warning(f"LOAD_RULES (JSONL): Line {line_num+1} in {filepath} did not contain a string value. Got: {type(rule_text_in_json_string)}")
                        error_count +=1
                except json.JSONDecodeError:
                    logger.warning(f"LOAD_RULES (JSONL): Line {line_num+1} in {filepath} failed to parse as JSON: {line[:100]}")
                    error_count +=1
    else:
        logger.error(f"LOAD_RULES: Unsupported file type for rules: {filepath}. Must be .txt or .jsonl")
        return 0, 0, 1 # Indicate type error

    valid_potential_rules = [r.strip() for r in potential_rules if r.strip()]
    total_to_process = len(valid_potential_rules)

    if total_to_process == 0 and error_count == 0:
        logger.info(f"LOAD_RULES: No valid rule segments found in {filepath} to process.")
        return 0, 0, 0
    elif total_to_process == 0 and error_count > 0:
         logger.warning(f"LOAD_RULES: No valid rule segments found to process. Encountered {error_count} parsing/format errors in {filepath}.")
         return 0, 0, error_count # Indicate only errors

    logger.info(f"LOAD_RULES: Attempting to add {total_to_process} potential rules from {filepath}...")
    for idx, rule_text in enumerate(valid_potential_rules):
        success, status_msg = add_rule_entry(rule_text)
        if success:
            added_count += 1
        elif status_msg == "duplicate":
            skipped_count += 1
        else:
            logger.warning(f"LOAD_RULES: Failed to add rule from {filepath} (segment {idx+1}): '{rule_text[:50]}...'. Status: {status_msg}")
            error_count += 1

    logger.info(f"LOAD_RULES: Finished processing {filepath}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors: {error_count}.")
    return added_count, skipped_count, error_count

def load_memories_from_file(filepath: str | None):
    """Loads memories from a local file (.json or .jsonl) and adds them to the system."""
    if not filepath:
        logger.info("LOAD_MEMORIES_FILE environment variable not set. Skipping memories loading from file.")
        return 0, 0, 0 # added, format_errors, save_errors

    if not os.path.exists(filepath):
        logger.warning(f"LOAD_MEMORIES: Specified memories file not found: {filepath}. Skipping loading.")
        return 0, 0, 0

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

    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        logger.error(f"LOAD_MEMORIES: Error reading file {filepath}: {e}", exc_info=False)
        return 0, 1, 0 # Indicate read error

    if not content.strip():
        logger.info(f"LOAD_MEMORIES: File {filepath} is empty. Skipping loading.")
        return 0, 0, 0

    file_ext = os.path.splitext(filepath.lower())[1]

    if file_ext == ".json":
        try:
            parsed_json = json.loads(content)
            if isinstance(parsed_json, list):
                memory_objects_to_process = parsed_json
            elif isinstance(parsed_json, dict):
                 # If it's a single object, process it as a list of one
                 memory_objects_to_process = [parsed_json]
            else:
                logger.warning(f"LOAD_MEMORIES (.json): File content is not a JSON list or object in {filepath}. Type: {type(parsed_json)}")
                format_error_count = 1
        except json.JSONDecodeError as e:
            logger.warning(f"LOAD_MEMORIES (.json): Invalid JSON file {filepath}. Error: {e}")
            format_error_count = 1
    elif file_ext == ".jsonl":
        for line_num, line in enumerate(content.splitlines()):
            line = line.strip()
            if line:
                try:
                    memory_objects_to_process.append(json.loads(line))
                except json.JSONDecodeError:
                    logger.warning(f"LOAD_MEMORIES (.jsonl): Line {line_num+1} in {filepath} parse error: {line[:100]}")
                    format_error_count += 1
    else:
        logger.error(f"LOAD_MEMORIES: Unsupported file type for memories: {filepath}. Must be .json or .jsonl")
        return 0, 1, 0 # Indicate type error

    total_to_process = len(memory_objects_to_process)

    if total_to_process == 0 and format_error_count > 0 :
         logger.warning(f"LOAD_MEMORIES: File parsing failed for {filepath}. Found {format_error_count} format errors and no processable objects.")
         return 0, format_error_count, 0
    elif total_to_process == 0:
         logger.info(f"LOAD_MEMORIES: No memory objects found in {filepath} after parsing.")
         return 0, 0, 0


    logger.info(f"LOAD_MEMORIES: Attempting to add {total_to_process} memory objects from {filepath}...")
    for idx, mem_data in enumerate(memory_objects_to_process):
        # Validate minimum structure
        if isinstance(mem_data, dict) and all(k in mem_data for k in ["user_input", "bot_response", "metrics"]):
            # Add entry without generating new embeddings if possible (assuming file contains embeddings)
            # NOTE: The current add_memory_entry function *always* generates embeddings.
            # If performance is an issue with large files, memory_logic might need
            # an optimized bulk import function that reuses existing embeddings or
            # generates them in batches. For now, we use the existing function.
            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"]) # add_memory_entry needs user_input, metrics, bot_response
            if success:
                added_count += 1
            else:
                # add_memory_entry currently doesn't return detailed error status
                logger.warning(f"LOAD_MEMORIES: Failed to save memory object from {filepath} (segment {idx+1}). Data: {str(mem_data)[:100]}")
                save_error_count += 1
        else:
            logger.warning(f"LOAD_MEMORIES: Skipped invalid memory object structure in {filepath} (segment {idx+1}): {str(mem_data)[:100]}")
            format_error_count += 1

    logger.info(f"LOAD_MEMORIES: Finished processing {filepath}. Added: {added_count}, Format/Structure Errors: {format_error_count}, Save Errors: {save_error_count}.")
    return added_count, format_error_count, save_error_count


# --- UI Functions for Rules and Memories (ui_refresh_..., ui_download_..., ui_upload_...) ---
def ui_refresh_rules_display_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."

def ui_download_rules_action_fn():
    rules_content = "\n\n---\n\n".join(get_all_rules_cached())
    if not rules_content.strip():
        gr.Warning("No rules to download.")
        return gr.DownloadButton(value=None, interactive=False, label="No Rules")
    try:
        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt", encoding='utf-8') as tmpfile:
            tmpfile.write(rules_content)
            return tmpfile.name
    except Exception as e:
        logger.error(f"Error creating rules download file: {e}")
        gr.Error(f"Failed to prepare rules for download: {e}")
        return gr.DownloadButton(value=None, interactive=False, label="Error")

def ui_upload_rules_action_fn(uploaded_file_obj, progress=gr.Progress()):
    if not uploaded_file_obj: return "No file provided for rules upload."
    try:
        with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f: content = f.read()
    except Exception as e_read: return f"Error reading file: {e_read}"
    if not content.strip(): return "Uploaded rules file is empty."
    added_count, skipped_count, error_count = 0,0,0

    potential_rules = []
    file_name_lower = uploaded_file_obj.name.lower()

    if file_name_lower.endswith(".txt"):
        potential_rules = content.split("\n\n---\n\n")
        if len(potential_rules) == 1 and "\n" in content:
             potential_rules = [r.strip() for r in content.splitlines() if r.strip()]
    elif file_name_lower.endswith(".jsonl"):
        for line_num, line in enumerate(content.splitlines()):
            line = line.strip()
            if line:
                try:
                    rule_text_in_json_string = json.loads(line)
                    if isinstance(rule_text_in_json_string, str):
                        potential_rules.append(rule_text_in_json_string)
                    else:
                        logger.warning(f"Rule Upload (JSONL): Line {line_num+1} did not contain a string value. Got: {type(rule_text_in_json_string)}")
                        error_count +=1
                except json.JSONDecodeError:
                    logger.warning(f"Rule Upload (JSONL): Line {line_num+1} failed to parse as JSON: {line[:100]}")
                    error_count +=1
    else:
        return "Unsupported file type for rules. Please use .txt or .jsonl."

    valid_potential_rules = [r.strip() for r in potential_rules if r.strip()]
    total_to_process = len(valid_potential_rules)

    if total_to_process == 0 and error_count == 0:
        return "No valid rules found in file to process."
    elif total_to_process == 0 and error_count > 0:
        return f"No valid rules found to process. Encountered {error_count} parsing/format errors."

    progress(0, desc="Starting rules upload...")
    for idx, rule_text in enumerate(valid_potential_rules):
        success, status_msg = add_rule_entry(rule_text)
        if success: added_count += 1
        elif status_msg == "duplicate": skipped_count += 1
        else: error_count += 1
        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} rules...")

    msg = f"Rules Upload: Total valid rule segments processed: {total_to_process}. Added: {added_count}, Skipped (duplicates): {skipped_count}, Errors (parsing/add): {error_count}."
    logger.info(msg); return msg

def ui_refresh_memories_display_fn(): return get_all_memories_cached() or []

def ui_download_memories_action_fn():
    memories = get_all_memories_cached()
    if not memories:
        gr.Warning("No memories to download.")
        return gr.DownloadButton(value=None, interactive=False, label="No Memories")

    jsonl_content = ""
    for mem_dict in memories:
        try: jsonl_content += json.dumps(mem_dict) + "\n"
        except Exception as e: logger.error(f"Error serializing memory for download: {mem_dict}, Error: {e}")

    if not jsonl_content.strip():
        gr.Warning("No valid memories to serialize for download.")
        return gr.DownloadButton(value=None, interactive=False, label="No Data")
    try:
        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".jsonl", encoding='utf-8') as tmpfile:
            tmpfile.write(jsonl_content)
            return tmpfile.name
    except Exception as e:
        logger.error(f"Error creating memories download file: {e}")
        gr.Error(f"Failed to prepare memories for download: {e}")
        return gr.DownloadButton(value=None, interactive=False, label="Error")

def ui_upload_memories_action_fn(uploaded_file_obj, progress=gr.Progress()):
    if not uploaded_file_obj: return "No file provided for memories upload."
    try:
        with open(uploaded_file_obj.name, 'r', encoding='utf-8') as f: content = f.read()
    except Exception as e_read: return f"Error reading file: {e_read}"
    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 = []

    file_ext = os.path.splitext(uploaded_file_obj.name.lower())[1]

    if file_ext == ".json":
        try:
            parsed_json = json.loads(content)
            if isinstance(parsed_json, list):
                memory_objects_to_process = parsed_json
            elif isinstance(parsed_json, dict):
                 memory_objects_to_process = [parsed_json]
            else:
                logger.warning(f"Memories Upload (.json): File content is not a JSON list or object. Type: {type(parsed_json)}")
                format_error_count = 1
        except json.JSONDecodeError as e:
            logger.warning(f"Memories Upload (.json): Invalid JSON file. Error: {e}")
            format_error_count = 1
    elif file_ext == ".jsonl":
        for line_num, line in enumerate(content.splitlines()):
            line = line.strip()
            if line:
                try:
                    memory_objects_to_process.append(json.loads(line))
                except json.JSONDecodeError:
                    logger.warning(f"Memories Upload (.jsonl): Line {line_num+1} parse error: {line[:100]}")
                    format_error_count += 1
    else:
        return "Unsupported file type for memories. Please use .json or .jsonl."

    if not memory_objects_to_process and format_error_count > 0 :
         return f"Memories Upload: File parsing failed. Found {format_error_count} format errors and no processable objects."
    elif not memory_objects_to_process:
         return "No valid memory objects found in the uploaded file."

    total_to_process = len(memory_objects_to_process)
    if total_to_process == 0: return "No memory objects to process (after parsing)."

    progress(0, desc="Starting memories upload...")
    for idx, mem_data in enumerate(memory_objects_to_process):
        if isinstance(mem_data, dict) and all(k in mem_data for k in ["user_input", "bot_response", "metrics"]):
            success, _ = add_memory_entry(mem_data["user_input"], mem_data["metrics"], mem_data["bot_response"])
            if success: added_count += 1
            else: save_error_count += 1
        else:
            logger.warning(f"Memories Upload: Skipped invalid memory object structure: {str(mem_data)[:100]}")
            format_error_count += 1
        progress((idx + 1) / total_to_process, desc=f"Processed {idx+1}/{total_to_process} memories...")

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

def save_edited_rules_action_fn(edited_rules_text: str, progress=gr.Progress()):
    # --- DEMO MODE CHANGE ---
    if DEMO_MODE:
        gr.Warning("Saving edited rules is disabled in Demo Mode.")
        return "Saving edited rules is disabled in Demo Mode."

    if not edited_rules_text.strip():
        return "No rules text to save."

    potential_rules = edited_rules_text.split("\n\n---\n\n")
    if len(potential_rules) == 1 and "\n" in edited_rules_text:
        potential_rules = [r.strip() for r in edited_rules_text.splitlines() if r.strip()]

    if not potential_rules:
        return "No rules found to process from editor."

    added, skipped, errors = 0, 0, 0
    unique_rules_to_process = sorted(list(set(filter(None, [r.strip() for r in potential_rules]))))

    total_unique = len(unique_rules_to_process)
    if total_unique == 0: return "No unique, non-empty rules found in editor text."

    progress(0, desc=f"Saving {total_unique} unique rules from editor...")

    for idx, rule_text in enumerate(unique_rules_to_process):
        success, status_msg = add_rule_entry(rule_text)
        if success: added += 1
        elif status_msg == "duplicate": skipped += 1
        else: errors += 1
        progress((idx + 1) / total_unique, desc=f"Processed {idx+1}/{total_unique} rules...")

    return f"Editor Save: Added: {added}, Skipped (duplicates): {skipped}, Errors/Invalid: {errors} from {total_unique} unique rules in text."

def app_load_fn():
    logger.info("App loading. Initializing systems...")
    initialize_memory_system()
    logger.info("Memory system initialized.")

    # --- Load Rules from File ---
    rules_added, rules_skipped, rules_errors = load_rules_from_file(LOAD_RULES_FILE)
    rules_load_msg = f"Rules: Added {rules_added}, Skipped {rules_skipped}, Errors {rules_errors} from {LOAD_RULES_FILE or 'None'}."
    logger.info(rules_load_msg)

    # --- Load Memories from File ---
    mems_added, mems_format_errors, mems_save_errors = load_memories_from_file(LOAD_MEMORIES_FILE)
    mems_load_msg = f"Memories: Added {mems_added}, Format Errors {mems_format_errors}, Save Errors {mems_save_errors} from {LOAD_MEMORIES_FILE or 'None'}."
    logger.info(mems_load_msg)

    final_status = f"AI Systems Initialized. {rules_load_msg} {mems_load_msg} Ready."

    # Initial population of all relevant UI components AFTER loading
    rules_on_load = ui_refresh_rules_display_fn()
    mems_on_load = ui_refresh_memories_display_fn()

    # Return values for outputs defined in demo.load
    return (
        final_status, # agent_stat_tb
        rules_on_load, # rules_disp_ta
        mems_on_load,  # mems_disp_json
        gr.Markdown(visible=False), # detect_out_md (initial state)
        gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True), # fmt_report_tb (initial state)
        gr.DownloadButton(interactive=False, value=None, visible=False), # dl_report_btn (initial state)
    )


# --- Gradio UI Definition ---
with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
        .gr-button { margin: 5px; }
        .gr-textbox, .gr-text-area, .gr-dropdown, .gr-json { border-radius: 8px; }
        .gr-group { border: 1px solid #e0e0e0; border-radius: 8px; padding: 10px; }
        .gr-row { gap: 10px; }
        .gr-tab { border-radius: 8px; }
        .status-text { font-size: 0.9em; color: #555; }
        .gr-json { max-height: 300px; overflow-y: auto; } /* Added scrolling for JSON */
    """
) as demo:
    # --- DEMO MODE CHANGE ---
    gr.Markdown(
        f"""
        # 🤖 AI Research Agent {'(DEMO MODE)' if DEMO_MODE else ''}
        Your intelligent assistant for research and knowledge management
        ### Special thanks to [Groq](https://groq.com) for their blazing fast inference
        
        """,
        elem_classes=["header"]
    )

    is_sqlite = MEMORY_STORAGE_BACKEND == "SQLITE"
    is_hf_dataset = MEMORY_STORAGE_BACKEND == "HF_DATASET"

    with gr.Row(variant="compact"):
        agent_stat_tb = gr.Textbox(
            label="Agent Status", value="Initializing systems...", interactive=False,
            elem_classes=["status-text"], scale=4
        )
        with gr.Column(scale=1, min_width=150):
            memory_backend_info_tb = gr.Textbox(
                label="Memory Backend", value=MEMORY_STORAGE_BACKEND, interactive=False,
                elem_classes=["status-text"]
            )
            sqlite_path_display = gr.Textbox(
                label="SQLite Path", value=MEMORY_SQLITE_PATH, interactive=False,
                visible=is_sqlite, elem_classes=["status-text"]
            )
            hf_repos_display = gr.Textbox(
                label="HF Repos", value=f"M: {MEMORY_HF_MEM_REPO}, R: {MEMORY_HF_RULES_REPO}",
                interactive=False, visible=is_hf_dataset, elem_classes=["status-text"]
            )

    with gr.Row():
        with gr.Sidebar():
            gr.Markdown("## ⚙️ Configuration")
            with gr.Group():
                gr.Markdown("### AI Model Settings")
                api_key_tb = gr.Textbox(
                    label="AI Provider API Key (Override)", type="password", placeholder="Uses .env if blank"
                )
                available_providers = get_available_providers()
                default_provider = available_providers[0] if available_providers else None
                prov_sel_dd = gr.Dropdown(
                    label="AI Provider", choices=available_providers,
                    value=default_provider, interactive=True
                )
                default_model_display = get_default_model_display_name_for_provider(default_provider) if default_provider else None
                model_sel_dd = gr.Dropdown(
                    label="AI Model",
                    choices=get_model_display_names_for_provider(default_provider) if default_provider else [],
                    value=default_model_display,
                    interactive=True
                )
            with gr.Group():
                gr.Markdown("### System Prompt")
                sys_prompt_tb = gr.Textbox(
                    label="System Prompt Base", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True
                )
            if MEMORY_STORAGE_BACKEND == "RAM":
                save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")

        with gr.Column(scale=3):
            with gr.Tabs():
                with gr.TabItem("💬 Chat & Research"):
                    with gr.Group():
                        gr.Markdown("### AI Chat Interface")
                        main_chat_disp = gr.Chatbot(
                            label=None, height=400, bubble_full_width=False,
                            avatar_images=(None, "https://raw.githubusercontent.com/huggingface/brand-assets/main/hf-logo-with-title.png"),
                            show_copy_button=True, render_markdown=True, sanitize_html=True
                        )
                        with gr.Row(variant="compact"):
                            user_msg_tb = gr.Textbox(
                                show_label=False, placeholder="Ask your research question...",
                                scale=7, lines=1, max_lines=3
                            )
                            send_btn = gr.Button("Send", variant="primary", scale=1, min_width=100)
                        with gr.Accordion("📝 Detailed Response & Insights", open=False):
                            fmt_report_tb = gr.Textbox(
                                label="Full AI Response", lines=8, interactive=True, show_copy_button=True
                            )
                            dl_report_btn = gr.DownloadButton(
                                "Download Report", value=None, interactive=False, visible=False
                            )
                            detect_out_md = gr.Markdown(visible=False)

                with gr.TabItem("🧠 Knowledge Base"):
                    with gr.Row(equal_height=True):
                        with gr.Column():
                            gr.Markdown("### 📜 Rules Management")
                            rules_disp_ta = gr.TextArea(
                                label="Current Rules", lines=10,
                                placeholder="Rules will appear here.",
                                interactive=True
                            )
                            gr.Markdown("To edit rules, modify the text above and click 'Save Edited Text', or upload a new file.")
                            save_edited_rules_btn = gr.Button("💾 Save Edited Text", variant="primary", interactive=not DEMO_MODE)
                            with gr.Row(variant="compact"):
                                dl_rules_btn = gr.DownloadButton("⬇️ Download Rules", value=None)
                                clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop", visible=not DEMO_MODE)
                            # --- DEMO MODE CHANGE ---
                            upload_rules_fobj = gr.File(
                                label="Upload Rules File (.txt with '---' separators, or .jsonl of rule strings)",
                                file_types=[".txt", ".jsonl"],
                                interactive=not DEMO_MODE
                            )
                            rules_stat_tb = gr.Textbox(
                                label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"]
                            )

                        with gr.Column():
                            gr.Markdown("### 📚 Memories Management")
                            mems_disp_json = gr.JSON(
                                label="Current Memories", value=[]
                            )
                            gr.Markdown("To add memories, upload a .jsonl or .json file.")
                            with gr.Row(variant="compact"):
                                dl_mems_btn = gr.DownloadButton("⬇️ Download Memories", value=None)
                                clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop", visible=not DEMO_MODE)
                            # --- DEMO MODE CHANGE ---
                            upload_mems_fobj = gr.File(
                                label="Upload Memories File (.jsonl of memory objects, or .json array of objects)",
                                file_types=[".jsonl", ".json"],
                                interactive=not DEMO_MODE
                            )
                            mems_stat_tb = gr.Textbox(
                                label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"]
                            )

    def dyn_upd_model_dd(sel_prov_dyn: str):
        models_dyn = get_model_display_names_for_provider(sel_prov_dyn)
        def_model_dyn = get_default_model_display_name_for_provider(sel_prov_dyn)
        return gr.Dropdown(choices=models_dyn, value=def_model_dyn, interactive=True)

    prov_sel_dd.change(fn=dyn_upd_model_dd, inputs=prov_sel_dd, outputs=model_sel_dd)

    # Inputs for the main chat submission function
    chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
    # Outputs for the main chat submission function (includes knowledge base displays)
    chat_outs = [user_msg_tb, main_chat_disp, agent_stat_tb, detect_out_md, fmt_report_tb, dl_report_btn, rules_disp_ta, mems_disp_json]

    chat_event_args = {"fn": handle_gradio_chat_submit, "inputs": chat_ins, "outputs": chat_outs}

    send_btn.click(**chat_event_args)
    user_msg_tb.submit(**chat_event_args)

    # Rules Management events
    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False)

    save_edited_rules_btn.click(
        fn=save_edited_rules_action_fn,
        inputs=[rules_disp_ta],
        outputs=[rules_stat_tb],
        show_progress="full"
    ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False)

    upload_rules_fobj.upload(
        fn=ui_upload_rules_action_fn,
        inputs=[upload_rules_fobj],
        outputs=[rules_stat_tb],
        show_progress="full"
    ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False)

    clear_rules_btn.click(
        fn=lambda: ("All rules cleared." if clear_all_rules_data_backend() else "Error clearing rules."),
        outputs=rules_stat_tb,
        show_progress=False
    ).then(fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta, show_progress=False)

    # Memories Management events
    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn, show_progress=False)

    upload_mems_fobj.upload(
        fn=ui_upload_memories_action_fn,
        inputs=[upload_mems_fobj],
        outputs=[mems_stat_tb],
        show_progress="full"
    ).then(fn=ui_refresh_memories_display_fn, outputs=mems_disp_json, show_progress=False)

    clear_mems_btn.click(
        fn=lambda: ("All memories cleared." if clear_all_memory_data_backend() else "Error clearing memories."),
        outputs=mems_stat_tb,
        show_progress=False
    ).then(fn=ui_refresh_memories_display_fn, outputs=mems_disp_json, show_progress=False)

    # FAISS save button visibility and action (RAM backend only)
    if MEMORY_STORAGE_BACKEND == "RAM" and 'save_faiss_sidebar_btn' in locals():
        def save_faiss_action_with_feedback_sidebar_fn():
            try:
                save_faiss_indices_to_disk()
                gr.Info("Attempted to save FAISS indices to disk.")
            except Exception as e:
                logger.error(f"Error saving FAISS indices: {e}", exc_info=True)
                gr.Error(f"Error saving FAISS indices: {e}")

        save_faiss_sidebar_btn.click(fn=save_faiss_action_with_feedback_sidebar_fn, inputs=None, outputs=None, show_progress=False)


    # --- Initial Load Event ---
    app_load_outputs = [
        agent_stat_tb,
        rules_disp_ta,
        mems_disp_json,
        detect_out_md,
        fmt_report_tb,
        dl_report_btn
    ]
    demo.load(fn=app_load_fn, inputs=None, outputs=app_load_outputs, show_progress="full")


if __name__ == "__main__":
    logger.info(f"Starting Gradio AI Research Mega Agent (v6.5 - Direct UI Update & Core Learning Memories, Memory: {MEMORY_STORAGE_BACKEND})...")
    app_port = int(os.getenv("GRADIO_PORT", 7860))
    app_server = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
    app_debug = os.getenv("GRADIO_DEBUG", "False").lower() == "true"
    app_share = os.getenv("GRADIO_SHARE", "False").lower() == "true"
    logger.info(f"Launching Gradio server: http://{app_server}:{app_port}. Debug: {app_debug}, Share: {app_share}")
    demo.queue().launch(server_name=app_server, server_port=app_port, debug=app_debug, share=app_share)
    logger.info("Gradio application shut down.")