app.py

import os
DEMO_MODE = False
MEMORY_STORAGE_TYPE = "RAM" 

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

# --- New Imports for Image KV Functionality ---
import io
import struct
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.exceptions import InvalidTag
# --- End New Imports ---
import random

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'}")


# --- KV to Image Functions (Constants and Implementation) ---
KEY_SIZE = 32
SALT_SIZE = 16
NONCE_SIZE = 12
TAG_SIZE = 16
PBKDF2_ITERATIONS = 480000
LENGTH_HEADER_SIZE = 4  # struct.pack('>I') uses 4 bytes
PREFERRED_FONTS = ["Arial", "Helvetica", "DejaVu Sans", "Verdana", "Calibri", "sans-serif"]
MAX_KEYS_TO_DISPLAY_OVERLAY = 15
def convert_pil_to_png_bytes(image: Image.Image) -> bytes:
    """Saves a PIL image to an in-memory buffer as PNG and returns the raw bytes."""
    with io.BytesIO() as buffer:
        image.save(buffer, format="PNG")
        return buffer.getvalue()
def _get_font(preferred_fonts, base_size):
    fp = None
    safe_base_size = int(base_size)
    if safe_base_size <= 0: safe_base_size = 10
    for n in preferred_fonts:
        try: ImageFont.truetype(n.lower()+".ttf",10); fp=n.lower()+".ttf"; break
        except IOError:
            try: ImageFont.truetype(n,10); fp=n; break
            except IOError: continue
    if fp:
        try: return ImageFont.truetype(fp, safe_base_size)
        except IOError: logger.warning(f"Font '{fp}' load failed with size {safe_base_size}. Defaulting.")
    try: return ImageFont.load_default(size=safe_base_size)
    except TypeError: return ImageFont.load_default()

def set_pil_image_format_to_png(image:Image.Image)->Image.Image:
    buf=io.BytesIO(); image.save(buf,format='PNG'); buf.seek(0)
    reloaded=Image.open(buf); reloaded.format="PNG"; return reloaded

def _derive_key(pw:str,salt:bytes)->bytes:
    kdf=PBKDF2HMAC(algorithm=hashes.SHA256(),length=KEY_SIZE,salt=salt,iterations=PBKDF2_ITERATIONS)
    return kdf.derive(pw.encode('utf-8'))

def encrypt_data(data:bytes,pw:str)->bytes:
    s=os.urandom(SALT_SIZE);k=_derive_key(pw,s);a=AESGCM(k);n=os.urandom(NONCE_SIZE)
    ct=a.encrypt(n,data,None); return s+n+ct

def decrypt_data(payload:bytes,pw:str)->bytes:
    ml=SALT_SIZE+NONCE_SIZE+TAG_SIZE;
    if len(payload)<ml: raise ValueError("Payload too short.")
    s,n,ct_tag=payload[:SALT_SIZE],payload[SALT_SIZE:SALT_SIZE+NONCE_SIZE],payload[SALT_SIZE+NONCE_SIZE:]
    k=_derive_key(pw,s);a=AESGCM(k)
    try: return a.decrypt(n,ct_tag,None)
    except InvalidTag: raise ValueError("Decryption failed: Invalid password/corrupted data.")
    except Exception as e: logger.error(f"Decrypt error: {e}",exc_info=True); raise

def _d2b(d:bytes)->str: return ''.join(format(b,'08b') for b in d)
def _b2B(b:str)->bytes:
    if len(b)%8!=0: raise ValueError("Bits not multiple of 8.")
    return bytes(int(b[i:i+8],2) for i in range(0,len(b),8))

def embed_data_in_image(img_obj:Image.Image,data:bytes)->Image.Image:
    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
    lb=struct.pack('>I',len(data));fp=lb+data;db=_d2b(fp);nb=len(db)
    if nb>len(fpx): raise ValueError(f"Data too large: {nb} bits needed, {len(fpx)} available.")
    for i in range(nb): fpx[i]=(fpx[i]&0xFE)|int(db[i])
    spx=fpx.reshape(px.shape); return Image.fromarray(spx.astype(np.uint8),'RGB')

def extract_data_from_image(img_obj:Image.Image)->bytes:
    img=img_obj.convert("RGB");px=np.array(img);fpx=px.ravel()
    hbc=LENGTH_HEADER_SIZE*8
    if len(fpx)<hbc: raise ValueError("Image too small for header.")
    lb="".join(str(fpx[i]&1) for i in range(hbc))
    try: pl=struct.unpack('>I',_b2B(lb))[0]
    except Exception as e: raise ValueError(f"Header decode error: {e}")
    if pl==0: return b""
    if pl>(len(fpx)-hbc)/8: raise ValueError("Header len corrupted or > capacity.")
    tpb=pl*8; so=hbc; eo=so+tpb
    if len(fpx)<eo: raise ValueError("Image truncated or header corrupted.")
    pb="".join(str(fpx[i]&1) for i in range(so,eo)); return _b2B(pb)

def parse_kv_string_to_dict(kv_str:str)->dict:
    if not kv_str or not kv_str.strip(): return {}
    dd={};
    for ln,ol in enumerate(kv_str.splitlines(),1):
        l=ol.strip()
        if not l or l.startswith('#'): continue
        lc=l.split('#',1)[0].strip();
        if not lc: continue
        p=lc.split('=',1) if '=' in lc else lc.split(':',1) if ':' in lc else []
        if len(p)!=2: raise ValueError(f"L{ln}: Invalid format '{ol}'.")
        k,v=p[0].strip(),p[1].strip()
        if not k: raise ValueError(f"L{ln}: Empty key in '{ol}'.")
        dd[k]=v
    return dd

def generate_brain_carrier_image(w=800, h=800) -> Image.Image:
    """
    Generates a high-quality carrier image with a multi-layered, procedural starfield
    using NumPy for performance and visual appeal.
    """
    # --- 1. Create the Gradient Background with NumPy ---
    center_x, center_y = w / 2, h / 2
    y_coords, x_coords = np.mgrid[0:h, 0:w]

    # Calculate distance of each pixel from the center
    distance = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
    max_distance = np.sqrt(center_x**2 + center_y**2)
    
    # Normalize distance to a 0-1 range
    distance_norm = distance / max_distance
    
    # Define colors and create the gradient array
    bg_center_color = np.array([20, 25, 40])  # Deeper blue center
    bg_outer_color = np.array([0, 0, 0])      # Black edges
    
    # Interpolate colors across all pixels at once (vectorized)
    # The [..., np.newaxis] part is for broadcasting the color channels
    gradient = bg_outer_color + (bg_center_color - bg_outer_color) * (1 - distance_norm[..., np.newaxis])
    
    # Convert the NumPy array to a PIL Image to start drawing on it
    img = Image.fromarray(gradient.astype(np.uint8), 'RGB')
    draw = ImageDraw.Draw(img)

    # --- 2. Draw a Multi-Layered Starfield ---
    
    # Layer 1: Distant, tiny stars (for depth)
    num_distant_stars = int((w * h) / 200)
    for _ in range(num_distant_stars):
        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
        brightness = random.randint(30, 90)
        draw.point((x, y), fill=(brightness, brightness, int(brightness * 1.1))) # Slightly blue tint

    # Layer 2: Main stars with glow, size, and color variation
    num_main_stars = int((w * h) / 1000)
    star_colors = [
        (255, 255, 255),  # White
        (220, 230, 255),  # Light Blue
        (255, 240, 220),  # Faint Yellow
    ]
    
    for _ in range(num_main_stars):
        x, y = random.randint(0, w - 1), random.randint(0, h - 1)
        dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
        dist_ratio = min(dist_from_center / max_distance, 1.0)
        
        # Base size and brightness increase with distance from center
        size = 0.5 + (2.5 * (dist_ratio ** 2))
        brightness = 120 + (135 * (dist_ratio ** 1.5))
        
        # Select a random base color
        color = random.choice(star_colors)
        
        # Apply brightness to the selected color
        final_color = tuple(int(c * (brightness / 255.0)) for c in color)
        
        # Simulate a soft glow by drawing a larger, dimmer circle first
        glow_size = size * 3
        glow_color = tuple(int(c * 0.3) for c in final_color) # Much dimmer
        draw.ellipse([x - glow_size, y - glow_size, x + glow_size, y + glow_size], fill=glow_color)

        # Simulate a "twinkle" effect for some stars by drawing a cross
        if random.random() < 0.15: # 15% chance to twinkle
            draw.line([x-size, y, x+size, y], fill=final_color, width=1)
            draw.line([x, y-size, x, y+size], fill=final_color, width=1)
        else:
            # Draw the main star on top of the glow
            draw.ellipse([x - size, y - size, x + size, y + size], fill=final_color)
            
    return img


def _get_text_measurement(draw_obj, text_str, font_obj):
    if hasattr(draw_obj, 'textbbox'):
        try:
            bbox = draw_obj.textbbox((0, 0), text_str, font=font_obj)
            width = bbox[2] - bbox[0]
            height = bbox[3] - bbox[1]
            return width, height
        except Exception: pass
    try:
        if hasattr(font_obj, 'getsize'): return font_obj.getsize(text_str)
        width, height = draw_obj.textsize(text_str, font=font_obj)
        return width, height
    except AttributeError:
        try:
            char_width_approx = font_obj.size * 0.6
            char_height_approx = font_obj.size
            return int(len(text_str) * char_width_approx), int(char_height_approx)
        except: return len(text_str) * 8, 10

def draw_key_list_dropdown_overlay(image: Image.Image, keys: list[str] = None, title: str = "Data Embedded") -> Image.Image:
    """Draws overlays on the image using the 'KeyLock' style."""
    img_overlayed = image.copy().convert("RGBA")
    draw = ImageDraw.Draw(img_overlayed, "RGBA")
    width, height = img_overlayed.size

    overlay_color = (15, 23, 42, 190)
    title_color = (226, 232, 240)
    key_color = (148, 163, 184)

    font_bold = _get_font(PREFERRED_FONTS, 30)
    font_regular = _get_font(PREFERRED_FONTS, 15)

    draw.rectangle([0, 20, width, 80], fill=overlay_color)
    draw.text((width / 2, 50), title, fill=title_color, font=font_bold, anchor="ms")

    if keys:
        box_padding = 15
        line_spacing = 6
        text_start_x = 35
        lines = keys

        line_heights = [_get_text_measurement(draw, line, font_regular)[1] for line in lines]
        total_text_height = sum(line_heights) + (len(lines) - 1) * line_spacing
        box_height = total_text_height + (box_padding * 2)
        box_y0 = height - box_height - 20

        draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
        current_y = box_y0 + box_padding
        
        for i, key_text in enumerate(lines):
            draw.text((text_start_x, current_y), key_text, fill=key_color, font=font_regular)
            if i < len(line_heights):
                current_y += line_heights[i] + line_spacing

    final_image_rgb = Image.new("RGB", img_overlayed.size, (0, 0, 0))
    final_image_rgb.paste(img_overlayed, (0, 0), img_overlayed)
    
    return final_image_rgb


# --- 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.")

            if significant_learnings_summary:
                learning_digest = "SYSTEM CORE LEARNING DIGEST:\n" + "\n".join(significant_learnings_summary)
                system_metrics = {
                    "takeaway": "Core knowledge refined.",
                    "response_success_score": 1.0,
                    "future_confidence_score": 1.0,
                    "type": "SYSTEM_REFLECTION"
                }
                add_memory_entry(
                    user_input="SYSTEM_INTERNAL_REFLECTION_TRIGGER",
                    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..."
    updated_rules_text = ui_refresh_rules_display_fn()
    updated_mems_json = ui_refresh_memories_display_fn()
    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))
        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>"))
    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})
    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)

    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:
                    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)
            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..."
                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)")

                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)

                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 (cleared_input, updated_gr_hist, status_txt, def_detect_out_md, def_fmt_out_txt, def_dl_btn, updated_rules_text, updated_mems_json)

            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)

        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)
        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

    if final_bot_resp_acc and not final_bot_resp_acc.startswith("Error:"):
        current_chat_session_history.extend([{"role": "user", "content": user_msg_txt}, {"role": "assistant", "content": final_bot_resp_acc}])
        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>"
        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
    else:
        status_txt = "Processing finished; no valid response or error occurred during main phase."

    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)

    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):
    if not filepath:
        logger.info("LOAD_RULES_FILE environment variable not set. Skipping rules loading from file.")
        return 0, 0, 0

    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

    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")
        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"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

    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

    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):
    if not filepath:
        logger.info("LOAD_MEMORIES_FILE environment variable not set. Skipping memories loading from file.")
        return 0, 0, 0

    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

    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):
                 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

    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):
        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:
                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 (Text and Image) ---
def convert_kb_to_kv_string(rules: list[str], memories: list[dict], include_rules: bool, include_memories: bool) -> str:
    """Serializes rules and/or memories into a single key-value string for image embedding based on user selection."""
    lines = ["# iLearn Knowledge Base Export", f"# Exported on: {datetime.utcnow().isoformat()}Z"]
    
    if include_rules:
        lines.append("\n# --- RULES ---")
        for i, rule_text in enumerate(rules):
            lines.append(f"rule_{i+1} = {json.dumps(rule_text)}")

    if include_memories:
        lines.append("\n# --- MEMORIES ---")
        for i, mem_dict in enumerate(memories):
            lines.append(f"memory_{i+1} = {json.dumps(mem_dict)}")

    return "\n".join(lines)


def ui_refresh_rules_display_fn(): return "\n\n---\n\n".join(get_all_rules_cached()) or "No rules found."
def ui_refresh_memories_display_fn(): return get_all_memories_cached() or []

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_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()):
    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 ui_upload_kb_from_image_fn(uploaded_image_filepath: str, password: str, progress=gr.Progress()):
    if DEMO_MODE: 
        gr.Warning("Uploading is disabled in Demo Mode.")
        return "Upload disabled in Demo Mode."
    if not uploaded_image_filepath: 
        return "No image file provided or pasted."
        
    progress(0, desc="Loading image...")
    try: 
        img = Image.open(uploaded_image_filepath)
    except Exception as e: 
        logger.error(f"KB ImgUL: Open fail: {e}")
        return f"Error: Could not open image file: {e}"

    progress(0.2, desc="Extracting data from image...")
    try:
        extracted_bytes = extract_data_from_image(img)
        if not extracted_bytes: return "No data found embedded in the image."
    except ValueError as e: 
        logger.error(f"KB ImgUL: Extract fail: {e}")
        return f"Error extracting data: {e}"
    except Exception as e: 
        logger.error(f"KB ImgUL: Extract error: {e}", exc_info=True)
        return f"Unexpected extraction error: {e}"
    
    kv_string = ""
    try:
        if extracted_bytes[:20].decode('utf-8', errors='ignore').strip().startswith("# iLearn"):
             kv_string = extracted_bytes.decode('utf-8')
             progress(0.4, desc="Parsing data...")
        elif password and password.strip():
            progress(0.3, desc="Attempting decryption...")
            kv_string = decrypt_data(extracted_bytes, password.strip()).decode('utf-8')
            progress(0.4, desc="Parsing decrypted data...")
        else: return "Data appears encrypted, but no password was provided."
    except (UnicodeDecodeError, InvalidTag, ValueError) as e:
        if isinstance(e, (InvalidTag, ValueError)): return f"Decryption Failed. Check password/file. Details: {e}"
        return "Data is binary and requires a password for decryption."
    except Exception as e: 
        logger.error(f"KB ImgUL: Decrypt error: {e}", exc_info=True)
        return f"Unexpected decryption error: {e}"

    if not kv_string: return "Could not get data from image (after potential decryption)."
    try: 
        kv_dict = parse_kv_string_to_dict(kv_string)
    except Exception as e: 
        logger.error(f"KB ImgUL: Parse fail: {e}")
        return f"Error parsing data: {e}"
    if not kv_dict: return "Parsed data is empty."

    rules_to_add, memories_to_add = [], []
    for key, value in kv_dict.items():
        if key.startswith("rule_"):
            try: rules_to_add.append(json.loads(value))
            except: logger.warning(f"KB ImgUL: Bad rule format for key {key}")
        elif key.startswith("memory_"):
            try:
                mem_dict = json.loads(value)
                if isinstance(mem_dict, dict) and all(k in mem_dict for k in ['user_input', 'bot_response', 'metrics']):
                    memories_to_add.append(mem_dict)
            except: logger.warning(f"KB ImgUL: Bad memory format for key {key}")
    
    added_rules, skip_r, err_r, added_mems, err_m = 0, 0, 0, 0, 0
    total = len(rules_to_add) + len(memories_to_add)
    progress(0.5, desc=f"Adding {len(rules_to_add)} rules...")
    for i, rule in enumerate(rules_to_add):
        s, m = add_rule_entry(rule)
        if s: added_rules += 1
        elif m == "duplicate": skip_r += 1
        else: err_r += 1
        if total > 0: progress(0.5 + (0.4 * ((i+1)/total)) if total else 0)
    
    progress(0.9, desc=f"Adding {len(memories_to_add)} memories...")
    for i, mem in enumerate(memories_to_add):
        s, _ = add_memory_entry(mem['user_input'], mem['metrics'], mem['bot_response'])
        if s: added_mems += 1
        else: err_m += 1
        if total > 0: progress(0.9 + (0.1 * ((i+1)/total)) if total else 0)
        
    progress(1.0, desc="Upload complete!")
    msg = f"Upload Complete. Rules - Add: {added_rules}, Skip: {skip_r}, Err: {err_r}. Mems - Add: {added_mems}, Err: {err_m}."
    logger.info(f"Image KB Upload: {msg}")
    return msg

def app_load_fn():
    logger.info("App loading. Initializing systems...")
    initialize_memory_system()
    logger.info("Memory system initialized.")
    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)
    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."
    rules_on_load, mems_on_load = ui_refresh_rules_display_fn(), ui_refresh_memories_display_fn()
    return (final_status, rules_on_load, mems_on_load, gr.Markdown(visible=False),
            gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True),
            gr.DownloadButton(interactive=False, value=None, visible=False))


# --- Create a placeholder image for the examples ---
# This makes the script self-contained and runnable without needing a separate file.
placeholder_filename = "placeholder_image.png"
try:
    if not os.path.exists(placeholder_filename):
        img = Image.new('RGB', (200, 100), color='darkblue')
        draw = Image.Draw(img)
        try:
            font = _get_font(PREFERRED_FONTS, 14)
            draw.text((10, 45), "Placeholder KB Image", font=font, fill='white')
        except Exception:
            draw.text((10, 45), "Placeholder", fill='white')
        img.save(placeholder_filename)
        logger.info(f"Created '{placeholder_filename}' for Gradio examples.")
except Exception as e:
    logger.error(f"Could not create placeholder image. The examples may not load correctly. Error: {e}")


def ui_download_kb_as_image_fn(password: str, progress=gr.Progress()):
    """
    Generates a KB image and returns both the image object for display
    and a file path for a download button.
    """
    progress(0, desc="Fetching knowledge base...")
    rules, memories = get_all_rules_cached(), get_all_memories_cached()
    if not rules and not memories:
        gr.Warning("Knowledge base is empty. Nothing to create.")
        # Return updates to hide the components if they were previously visible
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "Knowledge base is empty."

    progress(0.2, desc="Serializing data...")
    kv_string = convert_kb_to_kv_string(rules, memories)
    data_bytes = kv_string.encode('utf-8')

    if password and password.strip():
        progress(0.4, desc="Encrypting data...")
        try:
            data_bytes = encrypt_data(data_bytes, password.strip())
            gr.Info("Data encrypted successfully.")
        except Exception as e:
            logger.error(f"KB ImgDL: Encrypt failed: {e}")
            gr.Error(f"Encryption failed: {e}")
            return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"

    progress(0.6, desc="Generating carrier image...")
    carrier_image = generate_brain_carrier_image(w=800, h=800, msg="iLearn Knowledge Base")

    try:
        progress(0.7, desc="Embedding data...")
        embedded_image = embed_data_in_image(carrier_image, data_bytes)
    except ValueError as e:
        logger.error(f"KB ImgDL: Embed failed: {e}")
        gr.Error(f"Data is too large for this image size: {e}")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"

    progress(0.8, desc="Adding visual overlay...")
    keys_for_overlay = [f"Rule Count: {len(rules)}", f"Memory Count: {len(memories)}", "---"]
    for r in rules[:5]:
        match = re.search(r"\](.*)", r, re.DOTALL)
        rule_content = match.group(1).strip() if match else r
        keys_for_overlay.append(f"Rule: {rule_content[:40]}...")
    if len(rules) > 5: keys_for_overlay.append("...")

    title_overlay = "Encrypted Data" if password and password.strip() else "Embedded Data"
    final_image = draw_key_list_dropdown_overlay(embedded_image, keys=keys_for_overlay, title=title_overlay)

    progress(0.9, desc="Preparing final image and download file...")
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
            final_image.save(tmpfile, format="PNG")
            tmp_path = tmpfile.name

        progress(1.0, desc="Image created!")
        gr.Info("Image created and is ready for download or copy.")
        # Return updates to show the components with the new data
        return gr.update(value=final_image, visible=True), gr.update(value=tmp_path, visible=True), "Success! Image created."
    except Exception as e:
        logger.error(f"KB ImgDL: Save failed: {e}")
        gr.Error(f"Failed to save final image: {e}")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"




def ui_create_kb_image_fn(password: str, content_to_include: list, progress=gr.Progress()):
    """
    Generates a KB image and returns a file path to both the display and download components.
    """
    include_rules = "Include Rules" in content_to_include
    include_memories = "Include Memories" in content_to_include
    
    if not include_rules and not include_memories:
        gr.Warning("Nothing selected to save.")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "Nothing selected to save."

    progress(0, desc="Fetching knowledge base...")
    rules = get_all_rules_cached() if include_rules else []
    memories = get_all_memories_cached() if include_memories else []

    if not rules and not memories:
        gr.Warning("Knowledge base is empty or selected content is empty.")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), "No content to save."

    progress(0.2, desc="Serializing data...")
    kv_string = convert_kb_to_kv_string(rules, memories, include_rules, include_memories)
    data_bytes = kv_string.encode('utf-8')

    if password and password.strip():
        progress(0.4, desc="Encrypting data...")
        try:
            data_bytes = encrypt_data(data_bytes, password.strip())
        except Exception as e:
            logger.error(f"KB ImgDL: Encrypt failed: {e}")
            return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"

    progress(0.6, desc="Generating carrier image...")
    carrier_image = generate_brain_carrier_image(w=800, h=800)

    try:
        progress(0.7, desc="Embedding data...")
        embedded_image = embed_data_in_image(carrier_image, data_bytes)
    except ValueError as e:
        logger.error(f"KB ImgDL: Embed failed: {e}")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"

    progress(0.8, desc="Adding visual overlay...")
    keys_for_overlay = []
    if include_rules: keys_for_overlay.append(f"Rule Count: {len(rules)}")
    if include_memories: keys_for_overlay.append(f"Memory Count: {len(memories)}")
    
    title_overlay = "Encrypted KB" if password and password.strip() else "iLearn KB"
    final_image = draw_key_list_dropdown_overlay(embedded_image, keys=keys_for_overlay, title=title_overlay)

    progress(0.9, desc="Preparing final image and download file...")
    try:
        # Create a temporary file and save the image as a PNG.
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
            final_image.save(tmpfile, format="PNG")
            tmp_path = tmpfile.name
        progress(1.0, desc="Image created!")
        # Return the FILE PATH to both components.
        return gr.update(value=tmp_path, visible=True), gr.update(value=tmp_path, visible=True), "Success! Image created."
    except Exception as e:
        logger.error(f"KB ImgDL: Save failed: {e}")
        return gr.update(value=None, visible=False), gr.update(value=None, visible=False), f"Error: {e}"

def ui_load_from_sources_fn(image_filepath: str, rules_file_obj: object, mems_file_obj: object, password: str, progress=gr.Progress()):
    """
    Loads data from one of the available sources with precedence: Image > Rules File > Memories File.
    """
    if image_filepath:
        progress(0.1, desc="Image source detected. Starting image processing...")
        return ui_upload_kb_from_image_fn(image_filepath, password, progress)
    
    if rules_file_obj:
        progress(0.1, desc="Rules file detected. Starting rules import...")
        return ui_upload_rules_action_fn(rules_file_obj, progress)
        
    if mems_file_obj:
        progress(0.1, desc="Memories file detected. Starting memories import...")
        return ui_upload_memories_action_fn(mems_file_obj, progress)
        
    return "No file or image uploaded. Please provide a source file to load."

    
# --- 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: 400px; overflow-y: auto; }") as demo:

    gr.Markdown(f"# 🤖 iLearn: An Autonomous Learning Agent {'(DEMO MODE)' if DEMO_MODE else ''}", elem_classes=["header"])
    is_sqlite, is_hf_dataset = (MEMORY_STORAGE_BACKEND == "SQLITE"), (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.Tabs():
        with gr.TabItem("💬 Chat & Research"):
            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 "groq" not in available_providers else "groq"
                        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)

                with gr.Column(scale=3):
                    gr.Markdown("### AI Chat Interface")
                    main_chat_disp = gr.Chatbot(label=None, height=450, bubble_full_width=False,avatar_images=(None, "https://huggingface.co/spaces/Space-Share/bucket/resolve/main/images/pfp.webp"), 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():
                rules_stat_tb = gr.Textbox(label="Rules Status", interactive=False, lines=1, elem_classes=["status-text"])
                mems_stat_tb = gr.Textbox(label="Memories Status", interactive=False, lines=1, elem_classes=["status-text"])

            with gr.Tabs():
                with gr.TabItem("🎛️ System"):
                    gr.Markdown("View and directly manage the current rules and memories in the system.")
                    with gr.Row(equal_height=False, variant='compact'):
                        with gr.Column():
                            gr.Markdown("### 📜 Current Rules")
                            rules_disp_ta = gr.TextArea(label=None, lines=15, placeholder="Rules will appear here.", interactive=True)
                            save_edited_rules_btn = gr.Button("💾 Save Edited Rules", variant="primary", interactive=not DEMO_MODE)
                            clear_rules_btn = gr.Button("🗑️ Clear All Rules", variant="stop", visible=not DEMO_MODE)
                        with gr.Column():
                            gr.Markdown("### 📚 Current Memories")
                            mems_disp_json = gr.JSON(label=None, value=[], scale=1)
                            clear_mems_btn = gr.Button("🗑️ Clear All Memories", variant="stop", visible=not DEMO_MODE)

                with gr.TabItem("💾 Save KB"):
                    gr.Markdown("Export the current knowledge base as text files or as a single, portable PNG image.")
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("### Text File Export")
                            dl_rules_btn = gr.DownloadButton("⬇️ Download Rules (.txt)", value=None)
                            dl_mems_btn = gr.DownloadButton("⬇️ Download Memories (.jsonl)", value=None)
                            gr.Row()
                            if MEMORY_STORAGE_BACKEND == "RAM": save_faiss_sidebar_btn = gr.Button("Save FAISS Indices", variant="secondary")


                        with gr.Column():
                            gr.Markdown("### Image Export")
                            with gr.Group():
                                save_kb_password_tb = gr.Textbox(label="Password (optional for encryption)", type="password")
                                save_kb_include_cbg = gr.CheckboxGroup(label="Content to Include", choices=["Include Rules", "Include Memories"], value=["Include Rules", "Include Memories"])
                                create_kb_img_btn = gr.Button("✨ Create KB Image", variant="secondary")
                            # Removed type="pil" to allow Gradio to handle the raw PNG bytes correctly
                            kb_image_display_output = gr.Image(label="Generated Image (Right-click to copy)", visible=False)
                            kb_image_download_output = gr.DownloadButton("⬇️ Download Image File", visible=False)

                with gr.TabItem("📂 Load KB"):
                    gr.Markdown("Import rules, memories, or a full KB from local files or a portable PNG image.")
                    load_status_tb = gr.Textbox(label="Load Operation Status", interactive=False, lines=2)
                    load_kb_password_tb = gr.Textbox(label="Password (for decrypting images)", type="password")
                    
                    with gr.Group():
                        gr.Markdown("#### Sources (Priority: Image > Rules File > Memories File)")
                        with gr.Row():
                             upload_kb_img_fobj = gr.Image(label="1. Image Source", type="filepath", sources=["upload", "clipboard"], interactive=not DEMO_MODE)
                             upload_rules_fobj = gr.File(label="2. Rules File Source (.txt/.jsonl)", file_types=[".txt", ".jsonl"], interactive=not DEMO_MODE)
                             upload_mems_fobj = gr.File(label="3. Memories File Source (.json/.jsonl)", file_types=[".jsonl", ".json"], interactive=not DEMO_MODE)
                    
                    load_master_btn = gr.Button("⬆️ Load from Sources", variant="primary", interactive=not DEMO_MODE)
                    

    # --- Event Wiring ---
    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)

    # Chat Tab
    chat_ins = [user_msg_tb, main_chat_disp, prov_sel_dd, model_sel_dd, api_key_tb, sys_prompt_tb]
    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)

    # KB Tab -> System
    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)
    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)
    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)

    # KB Tab -> Save KB
    dl_rules_btn.click(fn=ui_download_rules_action_fn, inputs=None, outputs=dl_rules_btn, show_progress=False)
    dl_mems_btn.click(fn=ui_download_memories_action_fn, inputs=None, outputs=dl_mems_btn, show_progress=False)
    create_kb_img_btn.click(
        fn=ui_create_kb_image_fn,
        inputs=[save_kb_password_tb, save_kb_include_cbg],
        outputs=[kb_image_display_output, kb_image_download_output, load_status_tb],
        show_progress="full"
    )

    # KB Tab -> Load KB
    load_master_btn.click(
        fn=ui_load_from_sources_fn,
        inputs=[upload_kb_img_fobj, upload_rules_fobj, upload_mems_fobj, load_kb_password_tb],
        outputs=[load_status_tb],
        show_progress="full"
    ).then(
        fn=ui_refresh_rules_display_fn, outputs=rules_disp_ta
    ).then(
        fn=ui_refresh_memories_display_fn, outputs=mems_disp_json
    )

    # Sidebar FAISS button
    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)

    # App Load
    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 (v9.1 - Correct 1-Click JS Download, 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.")