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"

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


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
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
import random

load_dotenv()

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

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}, MemoryBackend={MEMORY_STORAGE_BACKEND}")
logger.info(f"Startup loading: Rules from {LOAD_RULES_FILE or 'None'}, Memories from {LOAD_MEMORIES_FILE or 'None'}")


KEY_SIZE = 32
SALT_SIZE = 16
NONCE_SIZE = 12
TAG_SIZE = 16
PBKDF2_ITERATIONS = 480000
LENGTH_HEADER_SIZE = 4
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:
    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:
    center_x, center_y = w / 2, h / 2
    y_coords, x_coords = np.mgrid[0:h, 0:w]

    distance = np.sqrt((x_coords - center_x)**2 + (y_coords - center_y)**2)
    max_distance = np.sqrt(center_x**2 + center_y**2)

    distance_norm = distance / max_distance

    bg_center_color = np.array([20, 25, 40])
    bg_outer_color = np.array([0, 0, 0])

    gradient = bg_outer_color + (bg_center_color - bg_outer_color) * (1 - distance_norm[..., np.newaxis])

    img = Image.fromarray(gradient.astype(np.uint8), 'RGB')
    draw = ImageDraw.Draw(img)

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

    num_main_stars = int((w * h) / 1000)
    star_colors = [
        (255, 255, 255),
        (220, 230, 255),
        (255, 240, 220),
    ]

    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)

        size = 0.5 + (2.5 * (dist_ratio ** 2))
        brightness = 120 + (135 * (dist_ratio ** 1.5))

        color = random.choice(star_colors)

        final_color = tuple(int(c * (brightness / 255.0)) for c in color)

        glow_size = size * 3
        glow_color = tuple(int(c * 0.3) for c in final_color)
        draw.ellipse([x - glow_size, y - glow_size, x + glow_size, y + glow_size], fill=glow_color)

        if random.random() < 0.15:
            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.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:
    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


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 _generate_action_plan(
    original_query: str, provider_name: str, model_display_name: str, ui_api_key_override: str | None, chat_history: list[dict]
) -> dict:
    history_str = "\n".join([f"{msg['role']}: {msg['content'][:150]}" for msg in chat_history[-4:]])

    plan_sys_prompt = """You are a master planner AI. Your goal is to decide the most efficient path to answer a user's query. You have two choices:

1.  **fast_response**: If the query is simple, conversational, or can be answered without external tools, choose this.
2.  **multi_step_plan**: If the query requires research, data retrieval, or complex reasoning, create a plan.

Your plan can use the following tools:
- `web_search`: Use for finding current, public information. The `task` should be a clear research goal (e.g., "Find the population of Tokyo in 2023").
- `memory_search`: Use for recalling past interactions or learned facts. The `task` should be a question to ask your memory (e.g., "What did the user previously say their name was?").
- `think`: A step for internal reflection. Use it to analyze the data gathered so far and decide if the plan needs adjustment or if enough information is present to proceed to the final answer. The `task` should be a question to yourself (e.g., "Is the gathered information sufficient to answer the user's main question?").
- `respond`: This should ALWAYS be the final step in a multi_step_plan. The `task` is always "Synthesize all information from the scratchpad and provide a comprehensive final answer to the user."

**Output format MUST be a single, valid JSON object.**

**Example for a simple query:**
{"action_type": "fast_response", "reason": "The user is just saying hello."}

**Example for a complex query:**
{
  "action_type": "multi_step_plan",
  "plan": [
    {"tool": "memory_search", "task": "What has the user previously expressed interest in regarding AI topics?"},
    {"tool": "web_search", "task": "Find recent advancements in large language models since early 2023."},
    {"tool": "web_scrape", "task": "https://example.com"},
    {"tool": "think", "task": "Based on the user's interests and recent advancements, what are the key points to highlight?"},
    {"tool": "respond", "task": "Synthesize all information from the scratchpad and provide a comprehensive final answer to the user."}
  ]
}
"""
    plan_user_prompt = f"Recent Conversation History:\n---\n{history_str}\n---\n\nUser Query: \"{original_query}\"\n\nBased on the query and history, what is the best action plan? Respond with JSON only."
    plan_messages = [{"role": "system", "content": plan_sys_prompt}, {"role": "user", "content": plan_user_prompt}]

    try:
        response_chunks = list(call_model_stream(
            provider=provider_name,
            model_display_name=model_display_name,
            messages=plan_messages,
            api_key_override=ui_api_key_override,
            temperature=0.0,
            max_tokens=1000
        ))
        resp_str = "".join(response_chunks).strip()
        json_match = re.search(r"\{.*\}", resp_str, re.DOTALL)
        if json_match:
            plan_data = json.loads(json_match.group(0))
            return plan_data
    except Exception as e:
        logger.error(f"PLAN_GEN: Failed to generate or parse action plan: {e}")
    
    return {
        "action_type": "multi_step_plan",
        "plan": [
            {"tool": "web_search", "task": original_query},
            {"tool": "respond", "task": "Synthesize all information from the scratchpad and provide a comprehensive final answer to the user."}
        ]
    }

def process_user_interaction_gradio(
    user_input: str,
    max_research_steps: int,
    provider_name: str,
    model_display_name: str,
    chat_history: 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]}...' Max Steps: {max_research_steps}")

    yield "status", "<i>[Deciding on an action plan...]</i>"
    action_plan_data = _generate_action_plan(user_input, provider_name, model_display_name, ui_api_key_override, chat_history)

    action_type = action_plan_data.get("action_type")
    
    if action_type == "fast_response":
        yield "status", "<i>[Executing fast response...]</i>"
        yield "plan", [{"tool": "fast_response", "task": action_plan_data.get("reason", "Direct answer.")}]

        now_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        final_sys_prompt = custom_system_prompt or DEFAULT_SYSTEM_PROMPT
        final_sys_prompt = f"Current Date/Time: {now_str}.\n\n" + final_sys_prompt
        
        messages_for_llm = [{"role": "system", "content": final_sys_prompt}] + chat_history + [{"role": "user", "content": user_input}]
        
        streamed_response = ""
        try:
            for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=messages_for_llm, api_key_override=ui_api_key_override, temperature=0.7, max_tokens=3000):
                streamed_response += chunk
                yield "response_chunk", chunk
        except Exception as e:
            streamed_response = f"\n\n(Error during fast response: {str(e)[:150]})"
            yield "response_chunk", streamed_response
            
        final_bot_text = streamed_response.strip()
        yield "final_response", {"response": final_bot_text}
        return

    plan = action_plan_data.get("plan", [])
    if not plan:
        plan = [{"tool": "web_search", "task": user_input}, {"tool": "respond", "task": "Synthesize a response."}]
    
    yield "plan", plan
    
    research_scratchpad = ""
    now_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    for i, step_action in enumerate(plan):
        tool = step_action.get("tool")
        task = step_action.get("task")
        
        if tool == 'respond':
            break

        if i + 1 > max_research_steps:
            research_scratchpad += f"\n\n---NOTE: Maximum research step budget of {max_research_steps} reached. Proceeding to final response.---\n"
            logger.warning(f"PUI_GRADIO [{request_id}]: Max research steps ({max_research_steps}) reached.")
            break

        yield "status", f"<i>[Executing Step {i+1}/{len(plan)-1}: {tool} -> {task[:70]}...]</i>"
        
        step_findings = f"Step {i+1} ({tool}: '{task}'): "
        
        if tool == 'web_search':
            try:
                web_results = search_and_scrape_duckduckgo(task, num_results=2)
                scraped_content = "\n".join([f"Source:\nURL:{r.get('url','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:1500]}\n---" for r in web_results]) if web_results else "No results found."
                synthesis_prompt = f"Relevant web content for the task '{task}':\n\n{scraped_content}\n\nConcisely summarize the findings from the content."
                summary = "".join(list(call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=[{"role": "user", "content": synthesis_prompt}], api_key_override=ui_api_key_override, temperature=0.1, max_tokens=400)))
                step_findings += summary
            except Exception as e:
                try:
                    web_results = search_and_scrape_google(task, num_results=2)
                    scraped_content = "\n".join([f"Source:\nURL:{r.get('url','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:1500]}\n---" for r in web_results]) if web_results else "No results found."
                    synthesis_prompt = f"Relevant web content for the task '{task}':\n\n{scraped_content}\n\nConcisely summarize the findings from the content."
                    summary = "".join(list(call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=[{"role": "user", "content": synthesis_prompt}], api_key_override=ui_api_key_override, temperature=0.1, max_tokens=400)))
                    step_findings += summary
                except Exception as e:
                    step_findings += f"Error during web search: {e}"

        elif tool == 'web_scrape':
            try:
                web_results = scrape_url(task)
                scraped_content = "\n".join([f"Source:\nURL:{r.get('url','N/A')}\nContent:\n{(r.get('content') or r.get('error') or 'N/A')[:1500]}\n---" for r in web_results]) if web_results else "No results found."
                synthesis_prompt = f"Relevant web content for the task '{task}':\n\n{scraped_content}\n\nConcisely summarize the findings from the content."
                summary = "".join(list(call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=[{"role": "user", "content": synthesis_prompt}], api_key_override=ui_api_key_override, temperature=0.1, max_tokens=400)))
                step_findings += summary
            except Exception as e:
                step_findings += f"Error during web search: {e}"

        elif tool == 'memory_search':
            try:
                retrieved_mems = retrieve_memories_semantic(task, k=3)
                if retrieved_mems:
                    memory_context = "\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])
                    step_findings += f"Found relevant memories:\n{memory_context}"
                else:
                    step_findings += "No relevant memories found."
            except Exception as e:
                step_findings += f"Error during memory search: {e}"

        elif tool == 'think':
            try:
                think_prompt = f"Original Query: '{user_input}'\n\nResearch Scratchpad:\n```\n{research_scratchpad}\n```\n\nMy current thinking task is: '{task}'. Based on the scratchpad, what is the conclusion of this thinking step?"
                thought = "".join(list(call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=[{"role": "user", "content": think_prompt}], api_key_override=ui_api_key_override, temperature=0.3, max_tokens=500)))
                step_findings += f"Conclusion: {thought}"
            except Exception as e:
                step_findings += f"Error during thinking step: {e}"
        else:
            step_findings += "Unknown tool specified in plan."

        research_scratchpad += f"\n\n---\n{step_findings}\n---"
        yield "step_result", {"step": i + 1, "tool": tool, "task": task, "result": step_findings}

    yield "status", "<i>[Synthesizing final report...]</i>"

    final_sys_prompt = custom_system_prompt or DEFAULT_SYSTEM_PROMPT
    final_sys_prompt += f"\n\nCurrent Date/Time: {now_str}. You have just completed a research plan. Synthesize the information in the 'Research Scratchpad' into a final, comprehensive answer. Cite sources by including URLs if available."
    final_user_prompt = f"Original user query: \"{user_input}\"\n\nResearch Scratchpad:\n```\n{research_scratchpad}\n```\n\nNow, provide the final, synthesized answer to the user."
    final_messages = [{"role": "system", "content": final_sys_prompt}, {"role": "user", "content": final_user_prompt}]

    streamed_response = ""
    try:
        for chunk in call_model_stream(provider=provider_name, model_display_name=model_display_name, messages=final_messages, api_key_override=ui_api_key_override, temperature=0.6, max_tokens=3000):
            streamed_response += chunk
            yield "response_chunk", chunk
    except Exception as e:
        error_msg = f"\n\n(Error during final synthesis: {str(e)[:150]})"
        streamed_response += error_msg
        yield "response_chunk", error_msg

    final_bot_text = streamed_response.strip() or "(No response or error during synthesis.)"
    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", {"response": final_bot_text}

def perform_post_interaction_learning(user_input: str, bot_response: str, provider: str, model_disp_name: str, 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 = []

    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 [Name], 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 [Name], 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
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

                if action == "add":
                    success, status_msg = add_rule_entry(insight_text)
                    if success:
                        processed_count +=1
                        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":
                    if old_insight and 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.")

                    success, status_msg = add_rule_entry(insight_text)
                    if success:
                        processed_count +=1
                        if insight_text.upper().startswith("[CORE_RULE"):
                             significant_learnings_summary.append(f"Core Rule Updated to: {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, max_research_steps: int, 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()
    plan_md_output = gr.Markdown(visible=False)
    final_report_tb = gr.Textbox(value="*Waiting...*", interactive=True, show_copy_button=True)
    dl_report_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, plan_md_output, final_report_tb, dl_report_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, plan_md_output, final_report_tb, dl_report_btn, updated_rules_text, updated_mems_json)

    internal_hist = list(current_chat_session_history)

    final_bot_resp_acc = ""
    temp_dl_file_path = None
    
    try:
        processor_gen = process_user_interaction_gradio(
            user_input=user_msg_txt, 
            max_research_steps=max_research_steps,
            provider_name=sel_prov_name, 
            model_display_name=sel_model_disp_name, 
            chat_history=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 = ""
        full_plan = []
        step_results = {}

        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"<i>{status_txt}</i>")
            elif upd_type == "plan":
                full_plan = upd_data
                plan_md = "### Action Plan\n" + "\n".join([f"**Step {i+1} ({step.get('tool')})**: {step.get('task')}" for i, step in enumerate(full_plan)])
                plan_md_output = gr.Markdown(value=plan_md, visible=True)
            elif upd_type == "step_result":
                step_num = upd_data["step"]
                step_results[step_num] = upd_data["result"]
                results_so_far = "### Research Log\n"
                for i in range(1, len(full_plan)):
                    if i in step_results:
                        results_so_far += f"**Step {i} ({full_plan[i-1].get('tool')})**: ✅ Completed\n"
                    elif i <= len(step_results) + 1:
                        results_so_far += f"**Step {i} ({full_plan[i-1].get('tool')})**: ⏳ In progress...\n"
                    else:
                        results_so_far += f"**Step {i} ({full_plan[i-1].get('tool')})**: - Pending\n"
                
                if updated_gr_hist and updated_gr_hist[-1][0] == user_msg_txt:
                    updated_gr_hist[-1] = (user_msg_txt, results_so_far)
            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":
                final_bot_resp_acc = upd_data["response"]
                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)")
                final_report_tb = 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
                        dl_report_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)
                        dl_report_btn = gr.DownloadButton(interactive=False, value=None, visible=False, label="Download Error")
                else:
                    dl_report_btn = gr.DownloadButton(interactive=False, value=None, visible=False)

            yield (cleared_input, updated_gr_hist, status_txt, plan_md_output, final_report_tb, dl_report_btn, updated_rules_text, updated_mems_json)

            if upd_type == "final_response": 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)
        final_report_tb = gr.Textbox(value=error_message_for_chat, interactive=True)
        dl_report_btn = gr.DownloadButton(interactive=False, value=None, visible=False)
        plan_md_output = 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, plan_md_output, final_report_tb, dl_report_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}])
        
        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, plan_md_output, final_report_tb, dl_report_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,
                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."
    else:
        status_txt = "Processing finished; no valid response or error occurred."

    updated_rules_text = ui_refresh_rules_display_fn()
    updated_mems_json = ui_refresh_memories_display_fn()

    yield (cleared_input, updated_gr_hist, status_txt, plan_md_output, final_report_tb, dl_report_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}")


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


def convert_kb_to_kv_string(rules: list[str], memories: list[dict], include_rules: bool, include_memories: bool) -> str:
    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))


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_create_kb_image_fn(password: str, content_to_include: list, progress=gr.Progress()):
    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:
        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 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()):
    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."

    
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.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)
            research_steps_slider = gr.Slider(label="Max Research Steps", minimum=1, maximum=10, step=1, value=3, 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.Tabs():
        with gr.TabItem("💬 Chat & Research"):
            with gr.Row():
                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 & Plan", open=False):
                        plan_display_md = gr.Markdown(visible=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)

        with gr.TabItem("🧠 Knowledge Base"):
            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():
                        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.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")
                            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)
                    gr.Examples(
                    examples=[
                        [placeholder_filename, ""],
                    ],
                    inputs=[upload_kb_img_fobj, load_kb_password_tb],
                    label="Click an Example to Load Data"
                    )

    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_ins = [user_msg_tb, research_steps_slider, 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, plan_display_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)

    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)

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

    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
    )

    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_outputs = [agent_stat_tb, rules_disp_ta, mems_disp_json, plan_display_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() == "false"
    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.")