app.py

"""
Privacy Risk Assessment UI for LLM Interactions (v9 - Dual DP Mechanisms)

Key improvements over v8:
    * Two complementary Differential Privacy mechanisms:
        1. Retrieval DP (Exponential Mechanism on document scores) – existing,
           sensitivity now computed empirically from actual utility range rather
           than a naive hardcoded constant (Dwork & Roth 2014; Koga et al. 2024).
        2. Input DP (RANTEXT-inspired Exponential Mechanism on word embeddings) –
           new; perturbs content words in the user's message before the LLM sees
           them, satisfying ε-Local DP per token (Tong et al. 2025, InferDPT).
    * User sees original message in chat; a DP badge with hover tooltip on each
      user bubble transparently communicates how many words were perturbed and at
      what privacy level – without replacing visible text.
    * Perturbed text logged to CSV as user_prompt_perturbed for research analysis.

URL parameters:
    model               – LLM model name
    rag                 – 0 or 1  (enable External Data linkage)
    epsilon             – float   (DP privacy budget, shared by both mechanisms)
    show_risk           – 0 or 1  (risk panel)
    show_tips           – 0 or 1  (PII tooltips)
    show_rag_highlights – 0 or 1  (External Data link highlights)
    show_pii_highlights – 0 or 1  (PII highlights)
    show_settings       – 0 or 1  (settings panel)
    demo                – 0 or 1  (demo mode, NO LLM calls)
    enable_social_scraping – 0 or 1  (enable social media scraping - experimental)
    token               – str     (access token for the session)
    show_dp             – int: when to show the Privacy Settings panel
                          0 = never
                          1 = after each conversation turn (default)
                          2 = at the beginning only (before first turn)
                          3 = only when the user presses "End conversation"
    show_infr_attr_card – int: when to show the "How the AI sees you" card
                          0 = never
                          1 = after each conversation turn (default)
                          2 = only when the user presses "End conversation"


"""

import copy
import pandas as pd
import json
import math
import os
import pickle
import random
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict
import threading
import csv
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging
import sys
import datetime
import gradio as gr
from openai import OpenAI
from langchain_community.embeddings import HuggingFaceEmbeddings
from create_rag_retriever import load_retriever_components, DPRetriever, HybridRetriever
import numpy as np
import hashlib
from langchain_core.documents import Document

# Import Apify social media scraper
from social_scrape_via_apify import (
    get_twitter_user_posts,
    get_facebook_posts,
    get_facebook_page_info,
    get_linkedin_posts,
    get_web_search_results
)

# Configure logging to write to BOTH file and console

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    datefmt='%H:%M:%S',
    handlers=[
        # logging.FileHandler(log_filename, mode='w', encoding='utf-8'),
        logging.StreamHandler(sys.stdout)  # Try console too
    ]
)

logger = logging.getLogger(__name__)
logger.info("="*60)


# ============================================================
# CONFIGURATION
# All tuneable constants are grouped here for easy maintenance.
# ============================================================

PROLIFIC_ID_PATTERN_REGEX = r'[A-Za-z0-9]{24}'

# ── Access Control ───────────────────────────────────────────
# Valid access tokens (store in HF Secrets in production).
# Each entry maps a token (read from an env var) to metadata.
VALID_ACCESS_TOKENS = {
    os.environ.get("VALID_ACCESS_TOKEN_1", None): {"name": "", "expires": "2099-12-31"},
    os.environ.get("VALID_ACCESS_TOKEN_2", None): {"name": "", "expires": "2099-12-31"},
    os.environ.get("VALID_ACCESS_TOKEN_3", None): {"name": "", "expires": "2099-12-31"},
}

# ── Logging & HF Hub ─────────────────────────────────────────
INTERACTION_LOG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "interaction_log.csv")
HF_LOG_REPO_ID   = "michael2222222/app-logs"   # HF dataset repo name
HF_LOG_REPO_TYPE = "dataset"

# Column headers for the per-participant CSV log file.
LOG_COLUMNS = [
    "timestamp",
    "session_source",
    "turn_number",
    "demo_mode",
    "scenario_mode",
    "persona_attributes",
    "model",
    "epsilon",
    "rag_enabled",
    "show_risk",
    "show_rag_highlights",
    "show_tips",
    "show_pii_highlights",
    "show_settings",
    "show_dp",
    "show_infr_attr_card",
    "show_social_scraping",
    "show_upload_data",
    "rag_corpus_path",
    "access_token",
    "social_scraping_enabled",
    "corpus_source",
    "uploaded_file_path",
    "uploaded_file_preview",
    "user_prompt",
    "user_prompt_perturbed",
    "num_input_dp_substitutions",
    "user_prompt_length",
    "user_perturbed_prompt_length",
    "llm_response",
    "llm_response_length",
    "risk_score",
    "rag_user_count",
    "rag_linkages_user",
    "rag_llm_count",
    "rag_linkages_llm",
    "pii_user_count",
    "pii_detected_user",
    "pii_user_perturbed_count",  # ← ADD
    "pii_detected_user_perturbed",  # ← ADD
    "pii_llm_count",
    "pii_detected_llm",
    "inference_warning_shown",
    "inference_panel_updated",
    "num_attributes_changed",
    "inference_lifts",
    "inferential_score",
    "inferential_score_breakdown",
    "scraped_docs",
    "scraped_social_summary",
]

# ── LLM Models ───────────────────────────────────────────────
# Uncomment / comment entries to enable or disable models.
MODEL_CONFIGS = [
    # ("together", "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", "Llama-4"),
    # ("together", "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", "Llama-3.1"),
    # ("openai", "gpt-4o", "GPT-4o"),
    # ("openai", "gpt-5-mini-2025-08-07", "GPT-5-mini"),
    # ("together", "openai/gpt-oss-120b", "GPT-OSS"),
    # ("together", "meta-llama/Llama-3.3-70B-Instruct-Turbo", "Llama-3.3-70B"),
    # ("together", "deepseek-ai/DeepSeek-V3.1", "DeepSeek-V3.1"),
    ("gemini", "gemini-2.5-flash-lite", "Gemini-2.5-Flash"),
]
DEFAULT_MODEL_PROVIDER = MODEL_CONFIGS[0][0]
DEFAULT_MODEL_ID       = MODEL_CONFIGS[0][1]
DEFAULT_MODEL_NAME     = MODEL_CONFIGS[0][2]

# ── LLM Inference Parameters ─────────────────────────────────
MAX_TOKENS_PII_DETECTION             = 1400
MAX_TOKENS_LLM_RESPONSE_GENERATION   = 260
MAX_TOKENS_INFER_ATTRIBUTES          = 1400
TEMPERATURE_LLM_RESPONSE_GENERATION  = 0.3
TEMPERATURE_INTERNAL_TASKS           = 0

# ── RAG & Retrieval ──────────────────────────────────────────
EMBEDDING_MODEL       = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_RETRIEVAL_PICKLE_PATH = "./faiss_persona_sarah_chen_retriever_components.pkl" #None

# Per-scenario system corpus pickle paths.
# Set a value to None to fall back to the default RETRIEVAL_PICKLE_PATH.
SCENARIO_RETRIEVER_PATHS = {
    "real":     DEFAULT_RETRIEVAL_PICKLE_PATH,
    # "persona": "./faiss_panorama_retriever_components.pkl",
    "persona2": "./faiss_persona_sarah_chen_retriever_components.pkl",
    # Add more scenarios here following the same pattern.
}

# Base directory used for relative file paths (e.g. fallback tweet JSON files).
_BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Per-scenario fallback tweet data (used when live scraping fails / returns empty).
# Each scenario_mode maps to its own JSON file of fallback tweets.
SCENARIO_FALLBACK_TWEET_PATHS = {
    "real":     None,                                      # no fallback for free-style
    "persona": None,
    "persona2": os.path.join(_BASE_DIR, "sarah_chen_tweets_scraped.json"),
    # Add more scenarios here following the same pattern.
}

# Default fallback tweet dataset (empty; populated per-scenario at runtime).
FALLBACK_TWEET_DATA = []

RAG_K       = 50
RAG_FETCH_K = 10000
RAG_LINKAGE_STORED_EXCERPT_LENGTH             = 500
ATTR_INFERENCE_EVIDENCE_STORED_EXCERPT_LENGTH = 500

# ── RAG Linkage Similarity Thresholds ────────────────────────
ROUGE_L_THRESHOLD    = 0.16   # Minimum ROUGE-L score (0.0-1.0)
COSINE_SIM_THRESHOLD = 0.16   # Minimum cosine similarity (0.0-1.0)
MIN_COMBINED_SCORE   = 0.2    # Minimum combined similarity score
ROUGE_WEIGHT         = 0.3    # Weight for ROUGE-L in combined score
COSINE_WEIGHT        = 0.7    # Weight for cosine similarity in combined score
MIN_NGRAM_LENGTH     = 4      # Minimum characters for n-gram matching
MIN_NGRAM_WORDS      = 4      # Minimum words in a phrase to highlight
MAX_NGRAM_WORDS      = 7      # Maximum words in a phrase to highlight

RAG_MIN_SIMILARITY_THRESHOLD = 0.25   # baseline threshold when DP is off
RAG_DP_THRESHOLD_SCALE       = 1.5   # how fast the threshold rises as ε falls

# ── Privacy / Differential Privacy ───────────────────────────
MAX_POSSIBLE_EPS         = 100.0
INFERENCE_LIFT_THRESHOLD = 0

DP_FLOOR = 0.25  # minimum retained fraction even under maximum perturbation
# ── Tunable constants ─────────────────────────────────────────────────────
INFER_MAX = 55.0  # hard cap on inference contribution
INFER_SCALE = 4.0  # controls how fast the curve saturates
PII_MAX = 40.0  # raised cap to accommodate new baseline
PII_BASE = 15.0  # flat baseline added whenever ≥1 PII is detected
PII_PER_HIT = 4.0  # points per detected PII entity (on top of baseline)
RAG_MAX = 15.0  # cap on RAG-linkage contribution
RAG_PER_HIT = 0.75  # points per RAG linkage

# ── UI & App Defaults ─────────────────────────────────────────
DEFAULT_USER_TEXT = (
    "Hi, I am Sarah Chen. I want to take a more preventive approach to my health. "
    "What routine tests or screenings should I consider that I could afford, and "
    "what is the nearest location where I can undergo these tests?"
)

# Human-readable label for documents retrieved from the default background corpus.
DEFAULT_CORPUS_SOURCE_LABEL = "From Social Profiles (Internet)"

# ── Locale & Attribute Mappings ───────────────────────────────
# Note: ATTRIBUTE_VALUES_MAP references LOCALE_TO_LOCATION and must come after it.
# PII_TYPE_TO_CATEGORY and PII_COLORS are defined later (in SECTION 1 / SECTION 2)
# because they depend on the PIICategory enum.
LOCALE_TO_LOCATION = {
    "en_PH": "Philippines",
    "en_CA": "Canada",
    "en_US": "United States",
    "en_IE": "Ireland",
    "en_NZ": "New Zealand",
    "en_IN": "India",
    "en_AU": "Australia",
    "en_GB": "United Kingdom",
    "en_IL": "Israel",
    "en_DE": "Germany",
    "en_IT": "Italy",
    "en_FR": "France",
}

SENSITIVE_ATTRIBUTES = ["Age bin", "Gender", "Marital Status", "Finance Status", "Education", "Locale"]

ATTRIBUTE_VALUES_MAP = {
    "Gender": ["Female", "Male"],
    "Age bin": ["0-17", "18-29", "30-44", "45-59", "60+"],
    "Marital Status": ["Single", "Married", "Divorced", "Widowed"],
    "Finance Status": ["Low", "Medium", "High"],
    "Locale": [LOCALE_TO_LOCATION["en_PH"], LOCALE_TO_LOCATION["en_CA"], LOCALE_TO_LOCATION["en_US"],
               LOCALE_TO_LOCATION["en_IE"], LOCALE_TO_LOCATION["en_NZ"], LOCALE_TO_LOCATION["en_IN"],
               LOCALE_TO_LOCATION["en_AU"], LOCALE_TO_LOCATION["en_GB"], LOCALE_TO_LOCATION["en_IL"],
                LOCALE_TO_LOCATION["en_DE"], LOCALE_TO_LOCATION["en_IT"], LOCALE_TO_LOCATION["en_FR"]],
    "Education": ["High School", "Bachelor's", "Master's", "PhD"]
}

# ── Population Priors for Inferential Privacy ────────────────
# Baseline population probabilities for each attribute value, used as the
# denominator in the inferential-privacy lift calculation.  The lift measures
# how much the system's available information (message text and/or RAG
# context) raises the probability of inferring an attribute above this
# base rate.
#
# Values within each attribute should sum to ≈1.0.  Calibrate these to
# the study's participant pool or the general population.
POPULATION_PRIORS = {
    "Gender": {
        "Female": 0.50,
        "Male":   0.50,
    },
    "Age bin": {
        "0-17":  0.15,
        "18-29": 0.15,
        "30-44": 0.38,
        "45-59": 0.15,
        "60+":   0.15,
    },
    "Marital Status": {
        "Single":   0.39,
        "Married":  0.51,
        "Divorced": 0.05,
        "Widowed":  0.05,
    },
    "Finance Status": {
        "Low":    0.23,
        "Medium": 0.53,
        "High":   0.23,
    },
    "Education": {
        "High School": 0.19,
        "Bachelor's":  0.43,
        "Master's":    0.19,
        "PhD":         0.19,
    },
    "Locale": {
        "Philippines":    0.09,
        "Canada":         0.05,
        "United States":  0.40,
        "Ireland":        0.05,
        "New Zealand":    0.05,
        "India":          0.05,
        "Australia":      0.05,
        "United Kingdom": 0.05,
        "Israel":         0.05,
        "Germany":        0.05,
        "Italy":          0.05,
        "France":         0.05,
    },
}

# Steepness parameter for the log-based inferential privacy saturation
# function: f(x; k) = k·ln(1+x) / (1 + k·ln(1+x)).
# Higher k → the curve reaches closer to 1 faster for the same lift.
INFERENTIAL_SCORE_STEEPNESS_K = 3.0

# ── Demo Prompts ──────────────────────────────────────────────
DATASET_PROMPTS_PANORAMA = {
    "health": (
        "Hi, I am {name}. I want to take a more preventive approach to my "
        "health. What routine tests or screenings should I consider that I "
        "could afford, and what is the nearest location where I can undergo "
        "these tests?"
    ),
    "job": (
        "I am {name}. I am searching for new job. Do you have any "
        "suggestions that would work for me?"
    ),
    "general": (
        "Hi, I'm {name}. Can you tell me about hobbies I might be "
        "interested?"
    ),
}

# ── Personas ──────────────────────────────────────────────────
# "real" means the participant uses their own profile (free-style).
# Each "personaX" entry contains the ground-truth attributes used in the study
# and a placeholder prompt shown in the text box when the app loads.
PERSONAS = {
    "real": {
        "name": "",
        "description": None,   # no placeholder injected
        "attributes":  {},
    },
    "persona1": {
        "name": "Raymond Phillips",
        "description": (
            "You are Raymond Phillips, a 52-year-old widowed male from the Philippines. "
            "You work as a fisherman and have a low income. "
            "Write your messages as Raymond would."
        ),
        "attributes": {
            "Gender":         "Male",
            "Age bin":        "45-59",
            "Marital Status": "Widowed",
            "Finance Status": "Low",
            "Locale":         "Philippines",
        },
    },
    "persona2": {
        "name": "Sarah Chen",
        "description": (
            # "You are Sarah Chen, a 38-year-old married female from the United States. "
            # "You work as a nurse and have a medium income. "
            """
            You will take on the role of Sarah Chen, a 38-year-old nurse. <br>
            Sarah recently moved from Seattle to Austin for a new ICU nursing job at St. David's Medical Center. Her spouse, Daniel, is still in Seattle finishing up his own job. 
            She has been struggling with anxiety and sleep problems since the move, and her doctor back in Seattle had prescribed her Lexapro, which she has been taking for about two years. She hasn't found a new doctor in Austin yet. 
            With long shifts, an empty apartment, and mounting credit card debt from the move, she has decided it's time to find a therapist or psychiatrist in Austin who accepts her Blue Shield insurance. <br>
            <br>
            She plans to ask for recommendations in a local community forum for nurses.
            """
        ),
        "attributes": {
            "Gender":         "Female",
            "Age bin":        "30-44",
            "Marital Status": "Married",
            "Finance Status": "Medium",
            "Locale":         "United States",
            "Education":       "Bachelor\'s"
        },
    },
}

# ============================================================
# END OF CONFIGURATION
# ============================================================


def _extract_twitter_username(url):
    """Extract Twitter/X username from a tweet URL.

    e.g. https://x.com/realDonaldTrump/status/123  →  'realDonaldTrump'
    """
    try:
        from urllib.parse import urlparse
        path = urlparse(url).path          # '/realDonaldTrump/status/123'
        parts = path.strip("/").split("/") # ['realDonaldTrump', 'status', '123']
        if parts and parts[0] not in ("", "search", "i", "intent", "hashtag"):
            return parts[0]
    except Exception:
        pass
    return ""


def _tweet_items_to_documents(tweet_items, fallback=False):
    """Convert raw tweet JSON items to langchain Documents.

    Each document's page_content is prefixed with the Twitter username so that
    it surfaces naturally in RAG linkage tooltips and evidence attribution.

    Args:
        tweet_items: list of dicts with at least 'url' and 'text' keys
        fallback: True when these come from the fallback dataset (affects log only)

    Returns:
        list of Document objects
    """
    from langchain_core.documents import Document
    documents = []
    source_tag = "fallback_tweet_data" if fallback else "social_media_twitter"

    for idx, item in enumerate(tweet_items):
        if not isinstance(item, dict):
            continue
        raw_text = item.get("text", "").strip()
        if not raw_text:
            continue

        url = item.get("url", "") or item.get("twitterUrl", "")
        username = _extract_twitter_username(url)

        # Prefix every post with the author so it appears in RAG context
        content = f"[{username}] {raw_text}" if username else raw_text

        doc = Document(
            page_content=content,
            metadata={
                "source": source_tag,
                "twitter_username": username,
                "full_name": username,          # used by RAG source tooltip
                "platform": "twitter",
                "type": "social_media_scrape",
                "post_index": idx,
                "tweet_url": url,
                "created_at": item.get("createdAt", ""),
                "is_retweet": item.get("isRetweet", False),
            }
        )
        documents.append(doc)

    label = "fallback" if fallback else "live"
    logger.info(f"  ✓ Converted {len(documents)} {label} tweets to Documents")
    return documents

# ============================================================
# SECTION 2.5 – INTERACTION LOGGER
# ============================================================

logger.info(f"Interaction will be recorded to file at: {INTERACTION_LOG_PATH}")


# Known API error prefixes returned by call_generate_response's except clause
_LLM_ERROR_PREFIXES = (
    "Error:",
    "You exceeded your current quota",
    "Rate limit",
    "insufficient_quota",
    "Connection error",
    "Timeout",
    "APIError",
    "AuthenticationError",
)


def _is_llm_error(response_text):
    """Return True when the LLM response string is actually an error message."""
    if not response_text:
        return True
    t = response_text.strip()
    return any(t.startswith(p) for p in _LLM_ERROR_PREFIXES)

def _participant_id(access_token, session_source):
    """Return a safe filename-compatible participant identifier.
    Prefers access_token; falls back to a short hash of session_source."""
    raw = (access_token or "").strip()
    if not raw:
        raw = "anon_" + hashlib.sha1((session_source or "").encode()).hexdigest()[:10]
    # Strip anything that is not alphanumeric, dash, or underscore
    return re.sub(r"[^a-zA-Z0-9_\-]", "_", raw)[:64]


def _participant_log_path(participant_id):
    """Local filesystem path for this participant's CSV."""
    log_dir = os.path.dirname(INTERACTION_LOG_PATH)
    return os.path.join(log_dir, f"log_{participant_id}.csv")


def _get_file_lock(participant_id):
    """Return (creating if needed) a per-participant threading lock."""
    with _log_lock:
        if participant_id not in _log_locks:
            _log_locks[participant_id] = threading.Lock()
        return _log_locks[participant_id]

# ── Add near the top with other imports ────────────────────────────────
from huggingface_hub import HfApi


def _push_log_to_hub(local_path, repo_filename):
    """Push a single participant CSV to the HF dataset repo in a background thread.

    The actual upload is dispatched to a daemon thread so this function returns
    immediately and never blocks the HTTP response path.  Failures are logged
    but otherwise swallowed — logging is best-effort.
    """
    def _do_upload():
        try:
            from huggingface_hub import HfApi
            token = os.environ.get("HF_TOKEN", None)
            if not token or not os.path.exists(local_path):
                return
            HfApi(token=token).upload_file(
                path_or_fileobj=local_path,
                path_in_repo=f"logs/{repo_filename}",
                repo_id=HF_LOG_REPO_ID,
                repo_type=HF_LOG_REPO_TYPE,
                commit_message=f"log update: {repo_filename}",
            )
            logger.info("📤 Log pushed to HF Hub: logs/%s", repo_filename)
        except Exception as e:
            logger.warning("Log push to Hub failed (non-fatal): %s", e)

    threading.Thread(target=_do_upload, daemon=True).start()

_LOG_FIELDS = [
    "timestamp",
    "session_source",
    "turn_number",
    "demo_mode",
    "scenario_mode",  # ← ADD: "real" | "persona1" | "persona2" | …
    "persona_attributes",  # ← ADD: JSON dict of persona ground-truth attrs, or "{}"
    "model",
    "epsilon",
    "rag_enabled",
    "show_risk",
    "show_rag_highlights",
    "show_tips",
    "show_pii_highlights",
    "show_settings",
    "access_token",
    "social_scraping_enabled",
    "corpus_source",
    "uploaded_file_path",
    "uploaded_file_preview",        # first 100 000 chars of the uploaded CSV
    "user_prompt",
    "user_prompt_perturbed",        # DP-perturbed version sent to LLM (same as user_prompt when DP off)
    "num_input_dp_substitutions",   # number of words perturbed by Input DP (0 when DP off)
    "user_prompt_length",           # character length of user_prompt
    "llm_response",
    "llm_response_length",          # character length of llm_response
    "risk_score",
    # ── RAG: split by user input vs LLM response ──────────────
    "num_rag_links_user",
    "rag_linkages_user",            # JSON: [{linked_text, source, similarity, ...}]
    "num_rag_links_llm",
    "rag_linkages_llm",             # JSON: [{linked_text, source, similarity, ...}]
    # ── PII: split by user input vs LLM response ──────────────
    "num_pii_detected_user",
    "pii_detected_user",            # JSON: [{text, type, confidence}]
    "num_pii_detected_llm",
    "pii_detected_llm",             # JSON: [{text, type, confidence}]
    # ── Inference ─────────────────────────────────────────────
    "inference_warning_shown",      # bool – was the warning banner displayed?
    "inference_lifts",              # JSON: {attr: {top_value, confidence, lift,
                                    #   prob_rag, prob_no_rag, evidence_type}}
    "inferential_score",            # float [0,1): 1−exp(−max_lift)
    "inferential_score_breakdown",  # JSON: {score, max_lift, max_attr, p_post, p_pop, p_no_rag}
    # ── Social scraping ───────────────────────────────────────
    "scraped_social_data",          # JSON: list of scraped post objects (if enabled)
    "scraped_social_summary",       # JSON: per-platform summary [{platform, post_count, posts:[{text,url}]}]
]

# _log_lock = threading.Lock()
_log_lock       = threading.Lock()          # guards _log_locks dict itself
_log_locks      = {}                        # per-participant file locks


def _ensure_log_file():
    """Create the CSV with header if it does not exist yet. Safe for concurrent startup."""
    if not os.path.exists(INTERACTION_LOG_PATH):
        with _log_lock:
            if not os.path.exists(INTERACTION_LOG_PATH):   # double-check after acquiring
                with open(INTERACTION_LOG_PATH, "w", newline="", encoding="utf-8") as f:
                    csv.writer(f).writerow(_LOG_FIELDS)
                logger.info(f"📋 Interaction log created: {INTERACTION_LOG_PATH}")


def _sanitize_cell(text):
    """Replace raw newlines and tabs in free-text fields so they cannot
    break CSV row boundaries, while remaining human-readable in Excel."""
    if not text:
        return text or ""
    return (
        str(text)
        .replace("\r\n", " \\n ")
        .replace("\r",   " \\n ")
        .replace("\n",   " \\n ")
        .replace("\t",   " \\t ")
    )

def append_interaction_log(
    session_source,
    turn_number,
    demo_mode,
    scenario_mode,
    persona_attributes,
    model,
    epsilon,
    rag_enabled,
    show_risk,
    show_rag_highlights,
    show_tips,
    show_pii_highlights,
    show_settings,
    access_token,
    social_scraping_enabled,
    corpus_source,
    uploaded_file_path,
    user_prompt,
    llm_response,
    risk_score,
    u_rag,           # RAGLink list for user input only
    r_rag,           # RAGLink list for LLM response only
    u_pii,           # PIIMatch list for user input only
    r_pii,           # PIIMatch list for LLM response only
    u_pii_perturbed,
    inference_metrics,
    inference_warning_shown,
    scraped_docs=None,          # list of Document objects from social scraping
    user_prompt_perturbed=None, # DP-perturbed version sent to LLM
    num_input_dp_substitutions=0,  # number of words perturbed by Input DP
    num_attributes_changed=0,
    show_dp=1,
    show_infr_attr_card=1,
    show_social_scraping=False,
    show_upload_data=False,
    rag_corpus_path="",
):
    """Append one interaction row to the persistent CSV log."""
    try:
        # ── Resolve participant identity and paths ────────────
        pid = _participant_id(access_token, session_source)
        log_path = _participant_log_path(pid)
        log_file = f"log_{pid}.csv"
        file_lock = _get_file_lock(pid)

        # ── Uploaded file info ────────────────────────────────
        uploaded_file_preview = ""
        if uploaded_file_path:
            try:
                with open(uploaded_file_path, "r", encoding="utf-8", errors="replace") as _f:
                    raw_preview = _f.read(100000)
                # json.dumps produces a single-line string with all special
                # characters escaped (\n, \r, \", \t, etc.), making it
                # completely safe to embed in any CSV cell.
                uploaded_file_preview = json.dumps(raw_preview)
            except Exception:
                uploaded_file_preview = "<unreadable>"

        # Create the file with headers if it does not exist yet
        if not os.path.exists(log_path):
            os.makedirs(os.path.dirname(log_path), exist_ok=True)
            with file_lock:
                if not os.path.exists(log_path):  # double-check after acquiring lock
                    with open(log_path, "w", newline="", encoding="utf-8") as f:
                        csv.writer(f, quoting=csv.QUOTE_ALL).writerow(LOG_COLUMNS)  # your existing header list

        # ── RAG linkages – user input ─────────────────────────
        def _serialise_rag(rag_list):
            return json.dumps([
                {
                    "linked_text":     lk.text,
                    "source":          getattr(lk, "source", ""),
                    "similarity":      round(float(lk.top_similarity), 4),
                    "start":           lk.start,
                    "end":             lk.end,
                    "corpus_snippets": lk.corpus_snippets,
                    "top_doc_score":   lk.top_doc_score,
                    "overlap_keywords": lk.overlap_keywords,
                }
                for lk in (rag_list or [])
            ], ensure_ascii=False)

        rag_linkages_user = _serialise_rag(u_rag)
        rag_linkages_llm  = _serialise_rag(r_rag)

        # ── PII detected – split by source ───────────────────
        def _serialise_pii(pii_list):
            rows = []
            for m in (pii_list or []):
                if isinstance(m, dict):
                    rows.append({
                        "text":       m.get("value", m.get("text", "")),
                        "type":       m.get("type", "unknown"),
                        "confidence": round(float(m.get("confidence", 0.0)), 4),
                    })
                else:
                    rows.append({
                        "text":       m.text,
                        "type":       m.fine_type,
                        "confidence": round(float(m.confidence), 4),
                    })
            return json.dumps(rows, ensure_ascii=False)

        pii_detected_user = _serialise_pii(u_pii)
        pii_detected_llm  = _serialise_pii(r_pii)

        # ── Inference lifts ───────────────────────────────────
        lifts = {}
        for attr, m in (inference_metrics or {}).items():
            lift = m.get("lift", 0)
            if not (math.isinf(lift) or math.isnan(lift)):
                lifts[attr] = {
                    "top_value":     m.get("top_value", ""),
                    "confidence":    round(float(m.get("confidence", m.get("probability", 0))), 4),
                    "lift":          round(float(lift), 4),
                    "prob_rag":      round(float(m.get("prob_rag", 0)), 4),
                    "prob_no_rag":   round(float(m.get("prob_no_rag", 0)), 4),
                    "p_pop":         round(float(m.get("p_pop", 0)), 4),
                    "evidence_type": m.get("evidence_type", "unknown"),
                }
        lifts_json = json.dumps(lifts, ensure_ascii=False)

        # ── Inferential privacy score and breakdown ───────────
        infer_score_info = calculate_inferential_privacy_score(inference_metrics or {})
        inferential_score = round(float(infer_score_info.get("mean_score", 0.0)), 4)
        infer_breakdown = {
            "score": inferential_score,
            "max_score":   round(float(infer_score_info.get("max_score", 0.0)), 4),
            "max_lift":    round(float(infer_score_info.get("max_lift", 0.0)), 4),
            "max_attr":    infer_score_info.get("max_attr", ""),
            "p_post":      round(float(infer_score_info.get("p_rag", 0.0)), 4),
            "p_pop":       round(float(infer_score_info.get("p_pop", 0.0)), 4),
            "p_no_rag":    round(float(infer_score_info.get("p_no_rag", 0.0)), 4),
            "mean_score": round(float(infer_score_info.get("mean_score", 0.0)), 4),
            "mean_lift": round(float(infer_score_info.get("mean_lift", 0.0)), 4),
            "median_score": round(float(infer_score_info.get("median_score", 0.0)), 4),
            "median_lift": round(float(infer_score_info.get("median_lift", 0.0)), 4),
        }
        inferential_score_breakdown_json = json.dumps(infer_breakdown, ensure_ascii=False)

        # ── Scraped social data ───────────────────────────────
        scraped_json = "[]"
        scraped_summary_json = "[]"
        if scraped_docs:
            try:
                scraped_json = json.dumps([
                    {
                        "text":     doc.page_content,
                        "metadata": doc.metadata if hasattr(doc, "metadata") else {},
                    }
                    for doc in scraped_docs
                ], ensure_ascii=False)
            except Exception:
                scraped_json = "[]"

            # Build a human-readable per-platform summary for analysis
            try:
                platform_groups = {}
                for doc in scraped_docs:
                    meta = doc.metadata if hasattr(doc, "metadata") else {}
                    platform = meta.get("platform", "unknown").lower()
                    source_url = (
                        meta.get("source_url") or
                        meta.get("tweet_url") or
                        meta.get("url") or
                        ""
                    )
                    text = doc.page_content or ""
                    if platform not in platform_groups:
                        platform_groups[platform] = []
                    platform_groups[platform].append({
                        "text": text,
                        "url":  source_url,
                    })

                summary_entries = []
                for platform, posts in platform_groups.items():
                    summary_entries.append({
                        "platform":    platform,
                        "post_count":  len(posts),
                        "posts": [
                            {
                                "text": p["text"],
                                "url":  p["url"],
                            }
                            for p in posts
                        ],
                    })
                scraped_summary_json = json.dumps(summary_entries, ensure_ascii=False)
            except Exception:
                scraped_summary_json = "[]"

        persona_attributes_json = json.dumps(persona_attributes or {}, ensure_ascii=False)

        row = [
            datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
            str(session_source),
            int(turn_number),
            str(demo_mode),
            str(scenario_mode),
            persona_attributes_json,
            str(model),
            str(epsilon),
            str(rag_enabled),
            str(show_risk),
            str(show_rag_highlights),
            str(show_tips),
            str(show_pii_highlights),
            str(show_settings),
            str(show_dp),
            str(show_infr_attr_card),
            str(show_social_scraping),
            str(show_upload_data),
            str(rag_corpus_path),
            str(access_token),
            str(social_scraping_enabled),
            str(corpus_source),
            str(uploaded_file_path or ""),
            uploaded_file_preview,
            _sanitize_cell(user_prompt),
            _sanitize_cell(user_prompt_perturbed if user_prompt_perturbed is not None else user_prompt),
            int(num_input_dp_substitutions),
            len(user_prompt or ""),
            len(user_prompt_perturbed or ""),
            _sanitize_cell(llm_response),
            len(llm_response or ""),
            round(float(risk_score), 2),
            len(u_rag or []),
            rag_linkages_user,
            len(r_rag or []),
            rag_linkages_llm,
            len(u_pii or []),
            pii_detected_user,
            len(u_pii_perturbed or []),
            _serialise_pii(u_pii_perturbed or []),
            len(r_pii or []),
            pii_detected_llm,
            str(bool(inference_warning_shown)),
            str(num_attributes_changed > 0),
            int(num_attributes_changed),
            lifts_json,
            inferential_score,
            inferential_score_breakdown_json,
            scraped_json,
            scraped_summary_json,
        ]

        with file_lock:
            with open(log_path, "a", newline="", encoding="utf-8") as f:
                csv.writer(f, quoting=csv.QUOTE_ALL).writerow(row)

        _push_log_to_hub(log_path, log_file)
        logger.info(
            f"📋 Logged turn {turn_number} for participant '{pid}' "
            f"(risk={risk_score:.0f}, "
            f"rag_user={len(u_rag or [])}, rag_llm={len(r_rag or [])}, "
            f"pii_user={len(u_pii or [])}, pii_llm={len(r_pii or [])}, "
            f"lifts={len(lifts)}, warning_shown={inference_warning_shown})"
        )

    except Exception as exc:
        logger.error(f"⚠️ Failed to write interaction log: {exc}")


def append_session_end_log(state):
    """Write a sentinel END row to the participant's CSV when they end the conversation."""
    try:
        access_token   = getattr(state, "_access_token", "")
        session_source = getattr(state, "_session_source", "unknown")
        pid      = _participant_id(access_token, session_source)
        log_path = _participant_log_path(pid)
        log_file = f"log_{pid}.csv"
        file_lock = _get_file_lock(pid)

        end_ts = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")

        # Columns whose values are meaningful for the END row
        config_values = {
            "timestamp":               end_ts,
            "session_source":          str(session_source),
            "turn_number":             "END",
            "demo_mode":               str(getattr(state, "_demo_mode", False)),
            "scenario_mode":           str(getattr(state, "_scenario_mode", "real")),
            "persona_attributes":      json.dumps(getattr(state, "_persona_attributes", {}), ensure_ascii=False),
            "model":                   str(getattr(state, "_last_model", "")),
            "epsilon":                 str(getattr(state, "_last_epsilon", float("inf"))),
            "rag_enabled":             str(getattr(state, "_last_rag_enabled", False)),
            "show_risk":               str(getattr(state, "_show_risk", True)),
            "show_rag_highlights":     str(getattr(state, "_show_rag_highlights", False)),
            "show_tips":               str(getattr(state, "_show_tips", False)),
            "show_pii_highlights":     str(getattr(state, "_show_pii_hl", True)),
            "show_settings":           str(getattr(state, "_show_settings", False)),
            "access_token":            str(access_token),
            "social_scraping_enabled": str(getattr(state, "_last_social_scraping", False)),
            "corpus_source":           str(getattr(state, "_last_corpus_source", "system")),
            "uploaded_file_path":      str(getattr(state, "_uploaded_file_path", "") or ""),
            "show_dp": str(getattr(state, "_show_dp", 1)),
            "show_infr_attr_card":  str(getattr(state, "_show_infr_attr_card", 1)),
            "show_social_scraping": str(getattr(state, "_show_social_scraping", False)),
            "show_upload_data": str(getattr(state, "_show_upload_data", False)),
            "rag_corpus_path": str(getattr(state, "_rag_corpus_path", "")),
        }

        # Build the row: config values where available, "END" for all per-turn columns
        row = [
            config_values.get(col, "END")
            for col in LOG_COLUMNS
        ]

        with file_lock:
            with open(log_path, "a", newline="", encoding="utf-8") as f:
                csv.writer(f, quoting=csv.QUOTE_ALL).writerow(row)

        _push_log_to_hub(log_path, log_file)
        logger.info(f"📋 Logged END-OF-CONVERSATION for participant '{pid}' at {end_ts}")
    except Exception as exc:
        logger.error(f"⚠️ Failed to write session-end log: {exc}")


os.makedirs("logs", exist_ok=True)

# ────────────────────────────────────────────────────────────
# Apify Configuration for Social Media Scraping
# ────────────────────────────────────────────────────────────

def validate_access_token(token):
    """Check if access token is valid and not expired.

    A Prolific token (exactly 24 alphanumeric characters) is always accepted.
    Any other token must appear in VALID_ACCESS_TOKENS and must not be expired.
    """
    if not token:
        return False

    # Accept valid Prolific tokens (exactly 24 alphanumeric characters).
    if re.fullmatch(PROLIFIC_ID_PATTERN_REGEX, token):
        return True

    # For non-Prolific tokens, check the VALID_ACCESS_TOKENS dictionary.
    token_info = VALID_ACCESS_TOKENS.get(token)
    if not token_info:
        return False

    # Check expiration date.
    from datetime import datetime
    expires = datetime.strptime(token_info["expires"], "%Y-%m-%d")
    if datetime.now() > expires:
        return False

    return True


def _load_fallback_tweets(scenario_mode: str):
    """Load and return fallback tweet items for the given scenario_mode.
    Returns [] if the file is missing, unmapped, or set to None."""
    key = str(scenario_mode).strip().lower()
    path = SCENARIO_FALLBACK_TWEET_PATHS.get(key)
    if not path:
        return []
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        logger.info(f"✓ Loaded {len(data)} fallback tweets for scenario '{key}' from {path}")
        return data
    except FileNotFoundError:
        logger.warning(f"⚠️ Fallback tweet file not found for scenario '{key}': {path}")
        return []
    except Exception as e:
        logger.error(f"✗ Failed to load fallback tweets for scenario '{key}': {e}")
        return []


# ============================================================
# SECTION 1 – DATA STRUCTURES
# ============================================================


def _fill_missing_with_uniform(probs):
    """Fill any missing SENSITIVE_ATTRIBUTES with a uniform distribution.

    This ensures the inference panel always has something to display even when
    an LLM call fails or returns a partial response.  A uniform distribution
    signals maximum uncertainty — no strong prediction — which is honest and
    safe to show.
    """
    filled = dict(probs)  # shallow copy — don't mutate the original
    for attr in SENSITIVE_ATTRIBUTES:
        if attr not in filled or not filled[attr]:
            values = ATTRIBUTE_VALUES_MAP.get(attr, [])
            if values:
                p = round(1.0 / len(values), 4)
                filled[attr] = {v: p for v in values}
    return filled


def get_persona(scenario_mode: str) -> dict:
    """Return the persona dict for a given scenario_mode string (case-insensitive).
    Falls back to 'real' if the key is not found."""
    return PERSONAS.get(str(scenario_mode).strip().lower(), PERSONAS["real"])


class PIICategory(Enum):
    IDENTITY  = "identity"   # IDs, names, SSN, usernames, etc.
    CONTACT   = "contact"    # Phone, emails
    LOCATION  = "location"   # Addresses, organizations
    SENSITIVE = "sensitive"  # Medical, financial, and other sensitive data


# Mapping from specific PII fine-types (as returned by the LLM) to the four
# broad UI categories.  Add new fine-types here; the rest of the code adapts.
PII_TYPE_TO_CATEGORY = {
    # ── Identity ─────────────────────────────────────────────────────────────
    "name":           PIICategory.IDENTITY,
    "full name":      PIICategory.IDENTITY,
    "first name":     PIICategory.IDENTITY,
    "last name":      PIICategory.IDENTITY,
    "ssn":            PIICategory.IDENTITY,
    "social security number": PIICategory.IDENTITY,
    "username":       PIICategory.IDENTITY,
    "id":             PIICategory.IDENTITY,
    "identifier":     PIICategory.IDENTITY,
    "passport":       PIICategory.IDENTITY,
    "license":        PIICategory.IDENTITY,
    "date of birth":  PIICategory.IDENTITY,
    "dob":            PIICategory.IDENTITY,
    "age":            PIICategory.IDENTITY,
    "ip address":     PIICategory.IDENTITY,
    "ip":             PIICategory.IDENTITY,
    # ── Contact ───────────────────────────────────────────────────────────────
    "email":          PIICategory.CONTACT,
    "email address":  PIICategory.CONTACT,
    "phone":          PIICategory.CONTACT,
    "phone number":   PIICategory.CONTACT,
    "mobile":         PIICategory.CONTACT,
    "fax":            PIICategory.CONTACT,
    # ── Location ──────────────────────────────────────────────────────────────
    "location":       PIICategory.LOCATION,
    "address":        PIICategory.LOCATION,
    "street address": PIICategory.LOCATION,
    "city":           PIICategory.LOCATION,
    "state":          PIICategory.LOCATION,
    "country":        PIICategory.LOCATION,
    "zip code":       PIICategory.LOCATION,
    "postal code":    PIICategory.LOCATION,
    "organization":   PIICategory.LOCATION,
    "workplace":      PIICategory.LOCATION,
    "school":         PIICategory.LOCATION,
    # ── Sensitive (medical, financial, and other sensitive data) ──────────────
    "medical":        PIICategory.SENSITIVE,
    "medical record": PIICategory.SENSITIVE,
    "diagnosis":      PIICategory.SENSITIVE,
    "condition":      PIICategory.SENSITIVE,
    "medication":     PIICategory.SENSITIVE,
    "prescription":   PIICategory.SENSITIVE,
    "treatment":      PIICategory.SENSITIVE,
    "health":         PIICategory.SENSITIVE,
    "disability":     PIICategory.SENSITIVE,
    "mental health":  PIICategory.SENSITIVE,
    "insurance":      PIICategory.SENSITIVE,
    "financial":      PIICategory.SENSITIVE,
    "bank account":   PIICategory.SENSITIVE,
    "credit card":    PIICategory.SENSITIVE,
    "credit card number": PIICategory.SENSITIVE,
    "account number": PIICategory.SENSITIVE,
    "salary":         PIICategory.SENSITIVE,
    "income":         PIICategory.SENSITIVE,
    "debt":           PIICategory.SENSITIVE,
    "loan":           PIICategory.SENSITIVE,
    "mortgage":       PIICategory.SENSITIVE,
    "tax id":         PIICategory.SENSITIVE,
    "political":      PIICategory.SENSITIVE,
    "religion":       PIICategory.SENSITIVE,
    "sexual orientation": PIICategory.SENSITIVE,
    "ethnicity":      PIICategory.SENSITIVE,
    "race":           PIICategory.SENSITIVE,
    "biometric":      PIICategory.SENSITIVE,
}


@dataclass
class PIIMatch:
    text: str
    start: int
    end: int
    fine_type: str
    category: PIICategory
    confidence: float


@dataclass
class RAGLink:
    """A span of user/assistant text linked to RAG corpus documents."""
    text: str
    start: int
    end: int
    corpus_snippets: list
    top_similarity: float
    top_doc_text: str = ""
    top_doc_score: float = 0.0
    overlap_keywords: list = None
    source: str = "From system data"  # Source of the matched text
    url: str = ""                      # Source URL of the matched document (if available)


@dataclass
class Message:
    role: str
    content: str
    pii_matches: list = None
    pii_card_matches: list = None  # PII detected on perturbed text (for the right-panel card)
    rag_links: list = None
    dp_metadata: dict = None  # Input DP info: {epsilon, num_substitutions, substitutions}


class ConversationState:
    """Simple in-memory conversation container."""
    def __init__(self):
        self.messages         = []
        self.last_probs_rag   = {}
        self.last_probs_no_rag = {}
        self.last_evidence_rag = {}
        # Logging metadata – set once per session
        self._session_source    = "unknown"
        self._access_token      = ""
        self._show_risk         = True
        self._show_rag_highlights = False
        self._show_tips         = False
        self._show_pii_hl       = True
        self._show_settings     = False
        self._show_dp            = 1
        self._show_infr_attr_card = 1
        self._show_social_scraping = False
        self._show_upload_data   = False
        self._rag_corpus_path    = ""
        self._uploaded_file_path = None
        self._turn_count        = 0          # incremented per user message
        self._peak_inference_metrics = {}
        self._prev_inference_attrs = {}
        self._scenario_mode      = "real"
        self._last_model = ""
        self._last_epsilon = float("inf")
        self._last_rag_enabled = False
        self._last_social_scraping = False
        self._last_corpus_source = "system"
        self._demo_mode = False
        self._persona_attributes = {}
        self._last_avatar_html   = ""        # cached inference card HTML for end-reveal
        self._last_privacy_html  = ""        # cached privacy settings HTML for end-reveal
        self._scraped_docs = None            # None = not yet scraped; [] = scraped, nothing found
        self._best_probs_rag = {}  # attr → best prob distribution seen so far
        self._best_evidence_rag = {}  # attr → evidence from that best turn


    def add(self, role, content, pii_matches=None, rag_links=None, dp_metadata=None, pii_card_matches=None):
        self.messages.append(Message(role=role, content=content,
                                     pii_matches=pii_matches,
                                     pii_card_matches=pii_card_matches,
                                     rag_links=rag_links,
                                     dp_metadata=dp_metadata))

    def clear(self):
        self.messages          = []
        self.last_probs_rag    = {}
        self.last_probs_no_rag = {}
        self.last_evidence_rag = {}
        self._turn_count       = 0
        self._peak_inference_metrics = {}
        self._last_model = ""
        self._last_epsilon = float("inf")
        self._last_rag_enabled = False
        self._last_social_scraping = False
        self._last_corpus_source = "system"
        self._demo_mode = False
        self._show_dp            = 1
        self._show_infr_attr_card = 1
        self._show_social_scraping = False
        self._show_upload_data   = False
        self._rag_corpus_path    = ""
        self._prev_inference_attrs = {}
        self._last_avatar_html   = ""
        self._last_privacy_html  = ""
        self._scraped_docs = None
        self._best_probs_rag = {}  # attr → best prob distribution seen so far
        self._best_evidence_rag = {}  # attr → evidence from that best turn

# ============================================================
# SECTION 2 – COLOUR CONFIGURATION
# ============================================================

PII_COLORS = {
    PIICategory.IDENTITY:  {"bg": "#FFE4E1", "border": "#E74C3C",
                            "label": "Identity (IDs, Names)"},
    PIICategory.CONTACT:   {"bg": "#E1F0FF", "border": "#2980B9",
                            "label": "Contact (Phone, Email)"},
    PIICategory.LOCATION:  {"bg": "#E8F5E9", "border": "#27AE60",
                            "label": "Location"},
    PIICategory.SENSITIVE: {"bg": "#FFF3E0", "border": "#E65100",
                            "label": "Sensitive (Medical / Financial)"},
}

RAG_LINK_COLOR = {"bg": "#E8EAF6", "border": "#3F51B5",
                  "label": "Linked to External Data"}


# ============================================================
# SECTION 3 – LLM CLIENT INITIALIZATION
# ============================================================

# Global LLM clients dictionary
LLM_CLIENTS = {}

def initialize_llm_clients():
    """Initialize API clients for each provider."""
    global LLM_CLIENTS

    try:
        # OpenAI
        openai_key = os.environ.get("OPENAI_API_KEY", None)
        if openai_key:
            os.environ["OPENAI_API_KEY"] = openai_key
            LLM_CLIENTS["openai"] = OpenAI(api_key=openai_key, timeout=60.0)
            logger.info("✓ OpenAI client initialized")

        # Together AI (OpenAI-compatible)
        together_key = os.environ.get("TOGETHER_API_KEY", None)
        if together_key:
            os.environ["TOGETHER_API_KEY"] = together_key
            LLM_CLIENTS["together"] = OpenAI(
                api_key=together_key,
                base_url="https://api.together.xyz/v1",
                timeout=60.0
            )
            logger.info("✓ Together AI client initialized")

        # Gemini (OpenAI-compatible)
        gemini_key = os.environ.get("GEMINI_API_KEY", None)
        if gemini_key:
            os.environ["GEMINI_API_KEY"] = gemini_key
            LLM_CLIENTS["gemini"] = OpenAI(
                api_key=gemini_key,
                base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
                timeout=60.0,
                max_retries=5
            )
            logger.info("✓ Gemini client initialized")


    except Exception as e:
        logger.info(f"Warning: Could not initialize LLM clients: {e}")
        logger.info("Demo mode will still work without LLM clients.")


# ============================================================
# SECTION 3.5 – SOCIAL MEDIA SCRAPING & CUSTOM CORPUS
# ============================================================

class SocialMediaScraper:
    """
    Social media scraper for retrieving user data from public sources using Apify.
    Integrates with Apify API to scrape posts from Twitter/X.
    """

    def __init__(self, enabled=False):
        self.enabled = enabled
        self.api_token = os.environ.get("APIFY_TOKEN", None)
        # Twitter/X actor (existing)
        self.actor_id = os.environ.get("APIFY_X_ACTOR_ID", None)
        # Facebook / LinkedIn actor IDs (read from env; the scraper module
        # also has built-in defaults so these can be left unset)
        self.fb_posts_actor = os.environ.get("APIFY_FB_POSTS_ACTOR", None)
        self.fb_pages_actor = os.environ.get("APIFY_FB_PAGES_ACTOR", None)
        self.li_posts_actor = os.environ.get("APIFY_LI_POSTS_ACTOR", None)
        self.web_scraper_actor = os.environ.get("APIFY_WEB_SCRAPER_ACTOR", None)
        logger.info(f"SocialMediaScraper initialized (enabled={enabled})")
        if enabled and not self.api_token:
            logger.warning("⚠️ Apify scraping enabled but APIFY_TOKEN not set!")

    def scrape_user_data(self, user_name, platforms=None, twitter_handle=None,
                         facebook_handle=None, linkedin_username=None,
                         fallback_data=None):
        """
        Scrape user data from one or more social media platforms using Apify.
        All platform scrapes are executed in parallel via ThreadPoolExecutor.
        """
        if not self.enabled:
            logger.info("Social media scraping is disabled")
            return []

        if platforms is None:
            platforms = ["Twitter"]

        normalised = [p.strip().lower() for p in platforms]

        # Build the list of (callable, args) tasks to run in parallel
        tasks = []
        if any(p in ("twitter", "x") for p in normalised):
            tasks.append((self._scrape_twitter,
                          (user_name, twitter_handle),
                          {"fallback_data": fallback_data}))

        if "facebook" in normalised:
            tasks.append((self._scrape_facebook_posts,
                          (user_name, facebook_handle), {}))
            tasks.append((self._scrape_facebook_page_info,
                          (user_name, facebook_handle), {}))

        if "linkedin" in normalised:
            tasks.append((self._scrape_linkedin_posts,
                          (user_name, linkedin_username), {}))

        if "web" in normalised:
            tasks.append((self._scrape_web, (user_name,), {}))

        if not tasks:
            return []

        documents = []
        with ThreadPoolExecutor(max_workers=len(tasks)) as executor:
            futures = {
                executor.submit(fn, *args, **kwargs): fn.__name__
                for fn, args, kwargs in tasks
            }
            for future in as_completed(futures):
                fn_name = futures[future]
                try:
                    result = future.result()
                    logger.info(f"  ✓ {fn_name} completed: {len(result)} documents")
                    documents.extend(result)
                except Exception as exc:
                    logger.error(f"  ✗ {fn_name} raised an exception: {exc}")

        logger.info(f"Social media scraping completed: {len(documents)} total documents")
        return documents

    def _scrape_twitter(self, user_name, twitter_handle=None, fallback_data=None):
        """Scrape Twitter/X; falls back to embedded JSON on failure."""
        documents = []
        live_ok = False

        if self.api_token:
            logger.info(f"  🔍 Live scraping Twitter for '{user_name}' via Apify...")
            try:
                posts = get_twitter_user_posts(
                    full_name=user_name,
                    api_token=self.api_token,
                    twitter_handle=twitter_handle,
                    actor_id=self.actor_id,
                    general_search_actor_id=self.web_scraper_actor or None,
                    max_items=30,
                    tweet_language="en",
                    sort="Latest",
                    log_csv_path="apify_requests_log.csv",
                )
                if posts:
                    logger.info(f"  ✓ Twitter live scrape returned {len(posts)} posts")
                    for idx, item in enumerate(posts):
                        if isinstance(item, dict):
                            post_text = item.get("text", "").strip()
                            post_url = item.get("url", "") or ""
                        else:
                            post_text = str(item).strip()
                            post_url = ""
                        if not post_text:
                            continue
                        documents.append(Document(
                            page_content=post_text,
                            metadata={
                                "source": "social_media_twitter",
                                "user": user_name,
                                "full_name": user_name,
                                "scraped_at": datetime.datetime.now().isoformat(),
                                "platform": "twitter",
                                "type": "social_media_scrape",
                                "post_index": idx,
                                "twitter_handle": twitter_handle or "",
                                "source_url": post_url,
                            },
                        ))
                    live_ok = True
                else:
                    logger.warning(f"  ⚠️ Twitter live scrape returned no posts for '{user_name}'")
            except Exception as exc:
                logger.error(f"  ✗ Twitter Apify scrape failed: {exc}")
                import traceback
                logger.error(traceback.format_exc())
        else:
            logger.warning("  ⚠️ APIFY_TOKEN not set – skipping Twitter live scrape")

        if not live_ok:
            data = fallback_data if fallback_data is not None else FALLBACK_TWEET_DATA
            logger.info(f"  🔄 Using fallback tweet dataset ({len(data)} items)...")
            documents = _tweet_items_to_documents(data, fallback=True)

        return documents


    def _scrape_facebook_posts(self, user_name, facebook_handle=None):
        """Scrape Facebook posts via Apify facebook-posts-scraper Actor."""
        documents = []
    
        if not self.api_token:
            logger.warning("  ⚠️ APIFY_TOKEN not set – skipping Facebook posts scrape")
            return documents
    
        logger.info(f"  🔍 Scraping Facebook posts for '{user_name}'...")
        try:
            posts = get_facebook_posts(
                full_name=user_name,
                api_token=self.api_token,
                facebook_handle=facebook_handle,
                max_items=30,
                actor_id=self.fb_posts_actor or None,
                general_search_actor_id=self.web_scraper_actor or None,
            )
            for idx, item in enumerate(posts):
                if isinstance(item, dict):
                    post_text = item.get("text", "").strip()
                    post_url = item.get("url", "") or ""
                else:
                    post_text = str(item).strip()
                    post_url = ""
                if not post_text:
                    continue
                documents.append(Document(
                    page_content=post_text,
                    metadata={
                        "source": "social_media_facebook",
                        "user": user_name,
                        "full_name": user_name,
                        "scraped_at": datetime.datetime.now().isoformat(),
                        "platform": "facebook",
                        "type": "social_media_scrape",
                        "post_index": idx,
                        "facebook_handle": facebook_handle or "",
                        "source_url": post_url,
                    },
                ))
            logger.info(f"  ✓ Facebook posts: {len(documents)} documents")
        except Exception as exc:
            logger.error(f"  ✗ Facebook posts scrape failed: {exc}")
            import traceback
            logger.error(traceback.format_exc())
    
        return documents


    def _scrape_facebook_page_info(self, user_name, facebook_handle=None):
        """Scrape Facebook page / profile bio via Apify facebook-pages-scraper Actor."""
        documents = []
    
        if not self.api_token:
            logger.warning("  ⚠️ APIFY_TOKEN not set – skipping Facebook page info scrape")
            return documents
    
        logger.info(f"  🔍 Scraping Facebook page info for '{user_name}'...")
        try:
            texts = get_facebook_page_info(
                full_name=user_name,
                api_token=self.api_token,
                facebook_handle=facebook_handle,
                actor_id=self.fb_pages_actor or None,
                general_search_actor_id=self.web_scraper_actor or None,
            )
            for idx, item in enumerate(texts):
                if isinstance(item, dict):
                    page_text = item.get("text", "").strip()
                    page_url = item.get("url", "") or ""
                else:
                    page_text = str(item).strip()
                    page_url = ""
                if not page_text:
                    continue
                documents.append(Document(
                    page_content=page_text,
                    metadata={
                        "source": "social_media_facebook",
                        "user": user_name,
                        "full_name": user_name,
                        "scraped_at": datetime.datetime.now().isoformat(),
                        "platform": "facebook",
                        "type": "social_media_scrape",
                        "post_index": idx,
                        "facebook_handle": facebook_handle or "",
                        "content_type": "page_info",
                        "source_url": page_url,
                    },
                ))
            logger.info(f"  ✓ Facebook page info: {len(documents)} documents")
        except Exception as exc:
            logger.error(f"  ✗ Facebook page info scrape failed: {exc}")
            import traceback
            logger.error(traceback.format_exc())
    
        return documents


    def _scrape_linkedin_posts(self, user_name, linkedin_username=None):
        """Scrape LinkedIn posts via Apify apimaestro/linkedin-profile-posts Actor."""
        documents = []
    
        if not self.api_token:
            logger.warning("  ⚠️ APIFY_TOKEN not set – skipping LinkedIn posts scrape")
            return documents
    
        logger.info(f"  🔍 Scraping LinkedIn posts for '{user_name}'...")
        try:
            posts = get_linkedin_posts(
                full_name=user_name,
                api_token=self.api_token,
                linkedin_username=linkedin_username,
                max_items=30,
                actor_id=self.li_posts_actor or None,
                general_search_actor_id=self.web_scraper_actor or None,
            )
            for idx, item in enumerate(posts):
                if isinstance(item, dict):
                    post_text = item.get("text", "").strip()
                    post_url = item.get("url", "") or ""
                else:
                    post_text = str(item).strip()
                    post_url = ""
                if not post_text:
                    continue
                documents.append(Document(
                    page_content=post_text,
                    metadata={
                        "source": "social_media_linkedin",
                        "user": user_name,
                        "full_name": user_name,
                        "scraped_at": datetime.datetime.now().isoformat(),
                        "platform": "linkedin",
                        "type": "social_media_scrape",
                        "post_index": idx,
                        "linkedin_username": linkedin_username or "",
                        "source_url": post_url,
                    },
                ))
            logger.info(f"  ✓ LinkedIn posts: {len(documents)} documents")
        except Exception as exc:
            logger.error(f"  ✗ LinkedIn posts scrape failed: {exc}")
            import traceback
            logger.error(traceback.format_exc())
    
        return documents

    def _scrape_web(self, user_name, search_query=None, extra_urls=None,
                    max_pages=3):
        """Scrape top web search results about the person via apify/web-scraper."""
        documents = []

        if not self.api_token:
            logger.warning("  ⚠️ APIFY_TOKEN not set – skipping web scrape")
            return documents

        logger.info(f"  🔍 Scraping web results for '{user_name}'...")
        try:
            texts = get_web_search_results(
                full_name=user_name,
                api_token=self.api_token,
                search_query=search_query,
                extra_urls=extra_urls,
                max_pages=max_pages,
                actor_id=self.web_scraper_actor or None,
            )
            for idx, item in enumerate(texts):
                if isinstance(item, dict):
                    web_text = item.get("text", "").strip()
                    web_url = item.get("url", "") or ""
                else:
                    web_text = str(item).strip()
                    web_url = ""
                if not web_text:
                    continue
                documents.append(Document(
                    page_content=web_text,
                    metadata={
                        "source": "social_media_web",
                        "user": user_name,
                        "full_name": user_name,
                        "scraped_at": datetime.datetime.now().isoformat(),
                        "platform": "web",
                        "type": "web_search_scrape",
                        "post_index": idx,
                        "source_url": web_url,
                    },
                ))
            logger.info(f"  ✓ Web scrape: {len(documents)} documents")
        except Exception as exc:
            logger.error(f"  ✗ Web scrape failed: {exc}")
            import traceback
            logger.error(traceback.format_exc())

        return documents

def build_retriever_from_csv(csv_file_path, progress=None, user_id_col="user id", text_col="text",
                              first_name_col="First Name", last_name_col="Last Name"):
    """
    Build a new HybridRetriever from uploaded CSV data.

    Args:
        csv_file_path: Path to CSV file
        progress: Gradio Progress object for tracking progress
        user_id_col: Column name for user IDs
        text_col: Column name for text content
        first_name_col: Column name for first names
        last_name_col: Column name for last names

    Returns:
        HybridRetriever instance or None if failed
    """
    try:
        logger.info(f"Building retriever from CSV: {csv_file_path}")
        if progress:
            progress(0.25, desc="Loading CSV data...")
            time.sleep(0.05)  # Brief pause for UI update

        # Read CSV
        df = pd.read_csv(csv_file_path)
        logger.info(f"  ✓ CSV loaded: {len(df)} rows, {len(df.columns)} columns")

        # Validate required column
        if text_col not in df.columns:
            available = ", ".join(f"'{c}'" for c in df.columns)
            raise ValueError(
                f"Required column '{text_col}' not found. "
                f"Available columns: {available}. "
                f"Please ensure your CSV has a column named '{text_col}'."
            )

        if progress:
            progress(0.35, desc=f"Processing {len(df)} documents...")
            time.sleep(0.05)

        # Build documents
        documents = []
        total_rows = len(df)
        last_progress_update = 0

        for idx, row in df.iterrows():
            # Update progress more frequently (every 5% or at least every 100 rows)
            current_progress = idx / total_rows if total_rows > 0 else 0
            if progress and (current_progress - last_progress_update >= 0.05 or idx % 100 == 0):
                progress(0.35 + 0.35 * current_progress, desc=f"Processing documents: {idx+1}/{total_rows}")
                last_progress_update = current_progress

            text = str(row.get(text_col, "")).strip()
            if len(text) < 10:  # Skip very short texts
                continue

            # Build metadata
            metadata = {
                "user_id": str(row.get(user_id_col, f"user_{idx}")),
                "source": "uploaded_csv",
                "row_index": idx
            }

            # Add optional fields
            if first_name_col in df.columns:
                metadata["first_name"] = str(row.get(first_name_col, ""))
            if last_name_col in df.columns:
                metadata["last_name"] = str(row.get(last_name_col, ""))

            # Add all other columns as metadata
            for col in df.columns:
                if col not in [user_id_col, text_col, first_name_col, last_name_col]:
                    metadata[col] = row.get(col)

            doc = Document(page_content=text, metadata=metadata)
            documents.append(doc)

        logger.info(f"  ✓ Created {len(documents)} documents")

        if progress:
            progress(0.70, desc=f"Building retriever from {len(documents)} documents (this may take a while)...")
            time.sleep(0.05)

        # Build retriever
        retriever = HybridRetriever(documents)
        logger.info("  ✓ Retriever built successfully")

        if progress:
            progress(0.88, desc="Retriever built successfully!")
            time.sleep(0.05)

        return retriever

    except Exception as e:
        logger.error(f"  ✗ Failed to build retriever from CSV: {str(e)}")
        raise  # re-raise so _upload_corpus's except block shows the message


# Global social media scraper instance
SOCIAL_SCRAPER = None


def initialize_social_scraper(enabled=False):
    """Initialize the social media scraper."""
    global SOCIAL_SCRAPER
    SOCIAL_SCRAPER = SocialMediaScraper(enabled=enabled)
    return SOCIAL_SCRAPER


# ============================================================
# SECTION 4 – PII DETECTION (LLM-BASED ONLY)
# ============================================================

def _mk(text, start, end, ft, conf):
    """Create PIIMatch with category mapping."""
    ft_lower = ft.lower().strip()

    # Map to category – exact lookup first, then keyword-based fallback
    if ft_lower in PII_TYPE_TO_CATEGORY:
        category = PII_TYPE_TO_CATEGORY[ft_lower]
    else:
        if any(kw in ft_lower for kw in ["name", "id", "ssn", "username",
                                          "passport", "license", "birth", "dob"]):
            category = PIICategory.IDENTITY
        elif any(kw in ft_lower for kw in ["email", "phone", "mobile", "fax"]):
            category = PIICategory.CONTACT
        elif any(kw in ft_lower for kw in ["address", "location", "organization",
                                            "city", "country", "zip"]):
            category = PIICategory.LOCATION
        elif any(kw in ft_lower for kw in ["medical", "health", "diagnosis",
                                            "medication", "prescription", "condition",
                                            "financial", "bank", "credit", "salary",
                                            "income", "insurance", "loan", "mortgage",
                                            "political", "religion", "ethnicity",
                                            "race", "biometric", "sexual"]):
            category = PIICategory.SENSITIVE
        else:
            category = PIICategory.IDENTITY  # safest default

    return PIIMatch(text=text, start=start, end=end, fine_type=ft_lower,
                    category=category, confidence=conf)


def _extract_pii_spans_from_values(text, items):
    """Convert LLM-returned {value,type} items to PIIMatch spans.

    Searches are case-insensitive so that minor capitalisation differences
    between the LLM's extraction and the original text do not cause misses.
    The span is anchored to the *original* casing in `text`.
    """
    matches = []
    used = [False] * (len(text) + 1)
    text_lower = text.lower()

    for it in items:
        if not isinstance(it, dict):
            continue
        val = str(it.get("value", "")).strip()
        ft  = str(it.get("type",  "")).strip().lower()
        conf = float(it.get("confidence", 0.85))
        conf = max(0.0, min(1.0, conf))
        if not val:
            continue

        val_lower = val.lower()
        start = 0
        while True:
            idx = text_lower.find(val_lower, start)
            if idx == -1:
                break
            j = idx + len(val)

            if any(used[idx:j]):
                start = idx + 1
                continue

            for k in range(idx, j):
                used[k] = True
            # Use the original-case slice from `text` for the PIIMatch text
            original_slice = text[idx:j]
            matches.append(_mk(original_slice, idx, j, ft, conf))
            # print("Detected PIIs >>>>", _mk(original_slice, idx, j, ft, conf))
            start = j

    matches.sort(key=lambda x: (x.start, x.end))
    return matches


def detect_pii_combined(user_text, response_text, provider, model_id):
    """
    Detect PII in BOTH user input and response in a SINGLE LLM call.
    This is more efficient than separate calls.

    Returns: (user_pii_items, response_pii_items)
    """
    logger.info("  🔍 [LLM CALL] Detecting PII in both user input and response (COMBINED)...")
    start_time = time.time()
    pii_cats = list(PII_TYPE_TO_CATEGORY.keys())
    client = LLM_CLIENTS.get(provider)
    if not client or (not user_text and not response_text):
        return [], []

    detection_prompt = f"""Identify PII in the two texts below. Only flag information that directly identifies or could be linked back to THIS specific user.

    CATEGORIES (use only these type values):
      "name"      – person's name
      "contact"   – email, phone, username, social handle
      "address"   – specific street, city, postal code the user states as their own
      "workplace" – specific named employer or school the user states as their own
      "health"    – medical conditions, medications, diagnoses
      "financial" – salary, account numbers, credit cards
      "sensitive" – ethnicity, religion, sexual orientation, political views

    NEVER flag:
      - Academic disciplines or research fields (e.g. "Machine Learning", "Industrial Engineering", "Computer Science")
      - Job titles or roles (e.g. "software engineer", "data analyst")
      - Generic sectors or industries (e.g. "healthcare sector", "financial institutions")
      - Places mentioned in general advice or background context (only flag when user says "I live/work/study at X")
      - Technology names or product names
      - Pronouns (e.g., "you, "I", "we")
      - Never flag any of the following strings as PII: {_DP_EMAIL_PLACEHOLDER}, {_DP_PHONE_PLACEHOLDER}, {_DP_SSN_PLACEHOLDER}, {_DP_ID_PLACEHOLDER}

    EXAMPLES
    --------
    Text: "I am Sarah Chen. I study Industrial Engineering at MIT and I work part-time at Google."
    Correct: [{{"type":"name","value":"Sarah Chen","confidence":0.99}},{{"type":"workplace","value":"MIT","confidence":0.95}},{{"type":"workplace","value":"Google","confidence":0.95}}]
    Wrong:   [{{"type":"workplace","value":"Industrial Engineering","confidence":0.80}}]  ← discipline, NOT a workplace

    Text: "I have a Master's degree in Machine Learning and I'm looking for jobs in AI."
    Correct: []
    Wrong:   [{{"type":"workplace","value":"Machine Learning","confidence":0.75}}]  ← field of study, never PII

    Text: "I'm Raymond Phillips. I live at 42 Oak Street, Boston. I take metformin for my diabetes."
    Correct: [{{"type":"name","value":"Raymond Phillips","confidence":0.99}},{{"type":"address","value":"42 Oak Street, Boston","confidence":0.97}},{{"type":"health","value":"metformin","confidence":0.93}},{{"type":"health","value":"diabetes","confidence":0.96}}]

    Text: "What are the best hospitals in New York for cardiology treatment?"
    Correct: []
    Wrong:   [{{"type":"address","value":"New York","confidence":0.80}}]  ← general context, not the user's location

    TEXT 1 (user): {user_text[:2000]}

    TEXT 2 (assistant): {response_text[:2000]}

    Respond ONLY with JSON, no other text.
    For example: {{"text1_pii": [{{"type": "name", "value": "Sarah Chen", "confidence": 0.98}}], "text2_pii": [{{"type": "address", "value": "Boston", "confidence": 0.99}}]}}"""

    try:
        # response = client.responses.create(
        response = client.chat.completions.create(
            model=model_id,
            messages=[{"role": "user", "content": detection_prompt}],
            # max_output_tokens=MAX_TOKENS_PII_DETECTION,
            max_tokens=MAX_TOKENS_PII_DETECTION,
            temperature=TEMPERATURE_INTERNAL_TASKS,
            timeout=25
        )

        response_text_llm = response.choices[0].message.content.strip()

        # Extract JSON
        json_match = re.search(r'```json\s*(\{.*?\})\s*```', response_text_llm, re.DOTALL)
        if json_match:
            json_str = json_match.group(1)
        else:
            json_match = re.search(r'\{.*\}', response_text_llm, re.DOTALL)
            if json_match:
                json_str = json_match.group(0)
            else:
                json_str = response_text_llm

        result = json.loads(json_str)

        user_pii_items = result.get('text1_pii', [])
        response_pii_items = result.get('text2_pii', [])

        elapsed = time.time() - start_time
        logger.info(f"  ✓ PII detection completed in {elapsed:.2f}s (found {len(user_pii_items)} in user, {len(response_pii_items)} in response)")

        return user_pii_items, response_pii_items

    except Exception as e:
        logger.info(f"  ✗ Error detecting PII: {e}")
        return [], []


# ============================================================
# SECTION 5 – RETRIEVER CLASSES
# ============================================================

@dataclass
class RetrievalResult:
    """Container for retrieval results."""
    document: object  # LangChain Document
    score: float
    sources: dict = field(default_factory=dict)


def strip_special_chars(text):
    """Strip digits and special characters."""
    return re.sub(r'[^a-zA-Z]', '', text.lower())

def epsilon_adjusted_threshold(epsilon, base=RAG_MIN_SIMILARITY_THRESHOLD, scale=RAG_DP_THRESHOLD_SCALE):
    """
    Compute a similarity threshold that tightens as epsilon falls.

    At epsilon=inf (no DP)  → threshold = base            (e.g. 0.30)
    At epsilon=1            → threshold ≈ base + 0.12     (e.g. 0.42)
    At epsilon=0.1          → threshold ≈ base + 0.29     (e.g. 0.59)

    The log1p curve means the boost is largest at very small epsilon
    (strong privacy) and negligible at large epsilon (no privacy).
    The cap at 0.92 prevents the threshold from becoming unreachable.
    """
    if epsilon == float("inf"):
        return base
    boost = scale * math.log1p(1.0 / max(epsilon, 1e-6))
    return min(base + boost, 0.92)


# ============================================================
# SECTION 5b – INPUT DIFFERENTIAL PRIVACY (WORD-LEVEL LDP)
# ============================================================
#
# Literature basis
# ----------------
# Primary:
#   Tong et al. (2025). "InferDPT: Privacy-Preserving Inference for
#   Closed-Box Large Language Models." IEEE Transactions on Dependable
#   and Secure Computing 22(5), 4625-4640.
#   → RANTEXT mechanism: for each content token, sample a replacement
#     from the vocabulary ∝ exp(ε·sim(w,w')/2Δ) using the Exponential
#     Mechanism on word embeddings.  Satisfies ε-LDP per token.
#
# Supporting:
#   Yue et al. (2021). "Differential Privacy for Text Analytics via
#   Natural Text Sanitization." ACL Findings.
#   → SanText/CusText: original LDP-on-embeddings paradigm.
#
# Novel contribution vs. Tong et al.
# ------------------------------------
# Rather than replacing the visible text shown to the user (which would
# confuse non-expert users and be unrelated to our tool's purpose), we:
#   • Send the PERTURBED text silently to the LLM.
#   • Display the ORIGINAL text in the chat bubble.
#   • Attach a visible DP badge + hover tooltip to each user bubble,
#     communicating (a) that perturbation occurred, (b) how many words
#     were changed, and (c) the ε guarantee it satisfied.
# This preserves transparency without exposing the noisy output, and is
# more aligned with usable-privacy principles for non-expert audiences.


# ── DP vocabulary and stopwords: loaded from external text files ─────────────
# Place dp_vocab.txt and dp_stopwords.txt in the same directory as this script.
# Each file: one word per line, lines starting with # are comments.
# Fallback to minimal built-in lists if the files are missing.

def _load_dp_word_file(filename, fallback):
    """Load a word list from a .txt file; fall back to *fallback* if missing."""
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
    try:
        with open(path, "r", encoding="utf-8") as f:
            words = [
                line.strip().lower()
                for line in f
                if line.strip() and not line.startswith("#")
            ]
        words = list(dict.fromkeys(words))   # deduplicate, preserve order
        logger.info("Loaded %d words from %s", len(words), filename)
        return words
    except FileNotFoundError:
        logger.warning(
            "%s not found – using built-in fallback (%d words). "
            "Create the file next to the script for an extended list.",
            filename, len(fallback),
        )
        return list(fallback)
    except Exception as exc:
        logger.error("Failed to load %s: %s", filename, exc)
        return list(fallback)

_DP_INPUT_VOCAB = _load_dp_word_file("dp_vocab.txt", [])
_DP_STOPWORDS   = set(_load_dp_word_file("dp_stopwords.txt", []))

# ── Structured-token DP redaction ─────────────────────────────────────────────
_DP_EMAIL_PLACEHOLDER   = "someone@example.com"
_DP_PHONE_PLACEHOLDER   = "555-000-0000"
_DP_SSN_PLACEHOLDER     = "000-00-0000"
_DP_ID_PLACEHOLDER      = "000000"
_DP_NUMBER_PLACEHOLDER  = "some"

_RE_EMAIL  = re.compile(r'[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}')
_RE_PHONE  = re.compile(
    r'(?<!\w)(?:\+?1[\s\-.]?)?(?:\(?\d{3}\)?[\s\-.]?)\d{3}[\s\-.]?\d{4}(?!\w)'
)
_RE_SSN = re.compile(
    r'(?<!\w)'
    r'(?:\d{3}-\d{2}-\d{4}'   # formatted:   123-45-6789
    r'|\d{9})'                 # unformatted: 123456789
    r'(?!\w)'
)
_RE_ID_NUMBER = re.compile(
    r'(?<!\w)\d{6,}(?!\w)'    # 6+ consecutive digits (passport, national ID, etc.)
)
_RE_NUMBER = re.compile(r'(?<!\w)\d+(?:[.,]\d+)*(?!\w)')


class DPInputPerturber:
    """Word-level Local DP perturbation of user input using the Exponential Mechanism.

    Each content word w is replaced by a vocabulary word w' sampled with
    probability proportional to exp(ε · cos_sim(w, w') / 2Δ), where Δ=2
    is the global sensitivity of cosine similarity (range [-1, 1]).

    This satisfies ε-Local Differential Privacy per token (Definition 1 /
    Equation 2, Tong et al. 2025).  Across independently perturbed tokens
    the parallel composition theorem (Dwork & Roth 2014) keeps the per-token
    budget constant rather than compounding.

    Vocabulary embeddings are cached at the class level so they are computed
    only once per process regardless of how many turns are processed.
    """

    # Class-level embedding cache (computed once)
    _vocab_matrix = None   # np.ndarray, shape (V, d), L2-normalised
    _vocab_words  = None   # list[str], aligned with _vocab_matrix rows
    _SENSITIVITY  = 2.0    # global sensitivity of cosine similarity
    _STRUCTURED_TOKEN_DECAY   = 0.15  # controls replacement rate for emails/numbers

    def __init__(self, epsilon=1.0, embedding_model=None, min_word_len=3):
        self.epsilon = epsilon
        self.embedding_model = embedding_model
        self.min_word_len = min_word_len
        self._build_vocab_cache()

    # ------------------------------------------------------------------
    # Vocabulary cache – built once per process
    # ------------------------------------------------------------------
    def _build_vocab_cache(self):
        if DPInputPerturber._vocab_matrix is not None:
            return
        if self.embedding_model is None:
            logger.warning("DPInputPerturber: no embedding model – perturbation disabled.")
            return
        try:
            logger.info("DPInputPerturber: embedding vocabulary (%d words)…", len(_DP_INPUT_VOCAB))
            raw = self.embedding_model.embed_documents(_DP_INPUT_VOCAB)
            mat = np.array(raw, dtype=np.float32)
            norms = np.linalg.norm(mat, axis=1, keepdims=True)
            DPInputPerturber._vocab_matrix = mat / np.clip(norms, 1e-10, None)
            DPInputPerturber._vocab_words  = list(_DP_INPUT_VOCAB)
            logger.info("DPInputPerturber: vocabulary ready (%d words).", len(_DP_INPUT_VOCAB))
        except Exception as exc:
            logger.error("DPInputPerturber: vocabulary embedding failed: %s", exc)

    def is_ready(self):
        return DPInputPerturber._vocab_matrix is not None

    # ------------------------------------------------------------------
    # Exponential mechanism for a single word vector
    # ------------------------------------------------------------------
    def _sample_replacement(self, word_vec_normed, original_word):
        V = DPInputPerturber._vocab_matrix
        words = DPInputPerturber._vocab_words

        sims_vocab = V @ word_vec_normed  # (V,)
        sims_all = np.append(sims_vocab, 1.0)  # (V+1,)  ← original appended

        logits = (self.epsilon * sims_all) / (2.0 * self._SENSITIVITY)
        logits -= logits.max()
        probs = np.exp(logits) / np.exp(logits).sum()

        idx = np.random.choice(len(sims_all), p=probs)
        if idx == len(words):  # last slot = original
            return original_word, float(probs[idx])
        return words[idx], float(probs[idx])

    # ------------------------------------------------------------------
    # Token helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _strip_punct(word):
        return word.strip(".,!?;:\"'()[]{}—–-")

    @staticmethod
    def _split_affixes(token):
        """Return (prefix_punct, core, suffix_punct)."""
        prefix, suffix, s = "", "", token
        while s and s[0] in "\"'([{":
            prefix += s[0]; s = s[1:]
        while s and s[-1] in ".,!?;:\"'()[]{}—–":
            suffix = s[-1] + suffix; s = s[:-1]
        return prefix, s, suffix

    def _redact_structured_tokens(self, tokens):
        """Probabilistically redact emails, phone numbers, SSNs, IDs, and numbers.

        Uses exponential decay: p_replace = exp(-ε · DECAY), so:
          - small ε (strong privacy) → p approaches 1 (almost always replaced)
          - large ε (weak privacy)   → p approaches 0 (almost always kept)
          - ε = ∞ is already handled by the early return in perturb()
        """
        p_replace = float(np.exp(-self.epsilon * self._STRUCTURED_TOKEN_DECAY))

        substitutions = []
        for i, tok in enumerate(tokens):
            placeholder = None

            if _RE_EMAIL.search(tok):
                placeholder = _RE_EMAIL.sub(_DP_EMAIL_PLACEHOLDER, tok)
            elif _RE_PHONE.search(tok):
                placeholder = _RE_PHONE.sub(_DP_PHONE_PLACEHOLDER, tok)
            elif _RE_SSN.search(tok):
                placeholder = _RE_SSN.sub(_DP_SSN_PLACEHOLDER, tok)
            elif _RE_ID_NUMBER.search(tok):
                placeholder = _RE_ID_NUMBER.sub(_DP_ID_PLACEHOLDER, tok)
            elif _RE_NUMBER.search(tok):  # ← was fullmatch
                placeholder = _RE_NUMBER.sub(_DP_NUMBER_PLACEHOLDER, tok)  # ← was constant

            if placeholder is None or placeholder == tok:
                continue

            if np.random.random() < p_replace:
                substitutions.append({
                    "original": tok,
                    "replacement": placeholder,
                    "position": i,
                    "prob": p_replace,
                })
                tokens[i] = placeholder

        return substitutions

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def perturb(self, text):
        if self.epsilon == float("inf") or not self.is_ready():
            return text, []

        tokens = text.split()
        if not tokens:
            return text, []

        new_tokens = list(tokens)

        # ── Step 1: structured-token redaction (emails, phones, numbers) ──────
        substitutions = self._redact_structured_tokens(new_tokens)

        # ── Step 2: exponential-mechanism perturbation of content words ────────
        content_idx = [
            i for i, tok in enumerate(new_tokens)
            if self._strip_punct(tok).lower() not in _DP_STOPWORDS
               and len(self._strip_punct(tok)) >= self.min_word_len
               and self._strip_punct(tok).replace("'", "").isalpha()
        ]
        if not content_idx:
            return " ".join(new_tokens), substitutions

        content_words = [self._strip_punct(new_tokens[i]) for i in content_idx]
        try:
            raw_embs = self.embedding_model.embed_documents(content_words)
        except Exception as exc:
            logger.warning("DPInputPerturber: embed call failed (%s) – skipping.", exc)
            return " ".join(new_tokens), substitutions

        embs = np.array(raw_embs, dtype=np.float32)
        norms = np.linalg.norm(embs, axis=1, keepdims=True)
        embs = embs / np.clip(norms, 1e-10, None)

        for local_i, global_i in enumerate(content_idx):
            original_core = content_words[local_i]
            replacement, prob = self._sample_replacement(embs[local_i], original_core.lower())

            if replacement.lower() == original_core.lower():
                continue

            prefix, _, suffix = self._split_affixes(new_tokens[global_i])
            if new_tokens[global_i][len(prefix):len(prefix) + 1].isupper():
                replacement = replacement.capitalize()

            new_tokens[global_i] = prefix + replacement + suffix
            substitutions.append({
                "original": original_core,
                "replacement": replacement,
                "position": global_i,
                "prob": prob,
            })

        return " ".join(new_tokens), substitutions


def _dp_badge_html(dp_metadata):
    """Build the inline DP badge appended to a user message bubble.

    The badge pill is right-aligned via its wrapper div (handled in
    fmt_conversation). The tooltip shows the full perturbed sentence
    with substituted words highlighted (green = replacement; the
    original word is shown in red strikethrough before it).

    Parameters
    ----------
    dp_metadata : dict
        epsilon           – float  privacy budget
        num_substitutions – int    number of words changed
        substitutions     – list   [{original, replacement, position, prob}]
        original_text     – str    what the user typed
        perturbed_text    – str    what was sent to the LLM
    """
    if not dp_metadata:
        return ""

    eps     = dp_metadata.get("epsilon", float("inf"))
    n_sub   = dp_metadata.get("num_substitutions", 0)
    subs    = dp_metadata.get("substitutions", [])
    orig    = dp_metadata.get("original_text", "")
    pert    = dp_metadata.get("perturbed_text", "")
    eps_str = f"{eps:.1f}" if eps != float("inf") else "∞"

    # ── Build the highlighted perturbed sentence ──────────────────────────
    # Map token position → substitution entry
    pos_to_sub = {s["position"]: s for s in subs}

    if pert and subs:
        tokens = pert.split()
        rendered_tokens = []
        for i, tok in enumerate(tokens):
            if i in pos_to_sub:
                sub   = pos_to_sub[i]
                orig_w = _esc(sub["original"])
                new_w  = _esc(tok)   # the replacement word (may be capitalised)
                rendered_tokens.append(
                    f'<del style="color:#c0392b;font-family:inherit;">{orig_w}</del>'
                    f'&nbsp;'
                    f'<span style="background:#d4edda;color:#155724;border-radius:3px;'
                    f'padding:1px 4px;font-weight:600;">{new_w}</span>'
                )
            else:
                rendered_tokens.append(_esc(tok))
        sentence_html = " ".join(rendered_tokens)
    elif pert:
        sentence_html = _esc(pert)
    else:
        sentence_html = '<em style="color:#888;">—</em>'

    # ── Tooltip content ───────────────────────────────────────────────────
    tooltip_content = (
        f'<div style="font-weight:700;font-size:1.05em;margin-bottom:6px;color:#111;">'
        f'🔒 Privacy Protection Applied</div>'

        f'<div style="margin-bottom:8px;color:#444;font-size:0.85em;line-height:1.5;">'
        f'<strong>{n_sub} word{"s" if n_sub != 1 else ""}</strong> in your message '
        f'{"were" if n_sub != 1 else "was"} automatically replaced by similar words '
        f'before the AI read it. ' # (ε&nbsp;=&nbsp;{eps_str}). '
        f'Your original message is shown in the chat above.'
        f'</div>'

        f'<div style="height:1px;background:#ddd;margin:6px 0;"></div>'

        f'<div style="font-weight:600;color:#2c3e50;margin-bottom:5px;font-size:0.9em;">'
        f'What the AI actually received:</div>'
        f'<div style="background:#f8f8f8;border:1px solid #ddd;border-radius:6px;'
        f'padding:8px 10px;font-size:0.92em;line-height:1.6;color:#222;">'
        f'{sentence_html}'
        f'</div>'

        f'<div style="margin-top:6px;font-size:0.82em;color:#666;">'
        f'<del style="color:#c0392b;">red</del> = original &nbsp;|&nbsp; '
        f'<span style="background:#d4edda;color:#155724;border-radius:3px;'
        f'padding:1px 4px;">green</span> = replacement'
        f'</div>'
    )

    # ── Badge pill + tooltip wrapper ──────────────────────────────────────
    badge_html = (
        f'<span class="tt-anchor" style="display:inline-block;">'
        f'<span style="display:inline-flex;align-items:center;gap:4px;'
        f'background:#e8f4fd;border:1px solid #7ec8e3;border-radius:12px;'
        f'padding:2px 10px;font-size:0.78em;color:#0a5275;cursor:help;'
        f'font-weight:600;user-select:none;white-space:nowrap;">'
        f'🔒&nbsp;DP-protected &middot; ' # ε={eps_str} &middot;
        f'{n_sub}&nbsp;word{"s" if n_sub!=1 else ""}&nbsp;changed'
        f'</span>'
        f'<span class="tt-popup" style="width:600px;font-size:0.82em!important;">'
        f'{tooltip_content}'
        f'</span>'
        f'</span>'
    )
    return badge_html


# ============================================================
# SECTION 6 – CORPUS LOADING (GENERIC)
# ============================================================

_GLOBAL_RETRIEVER = None
_DP_INPUT_PERTURBER = None   # DPInputPerturber instance, initialised lazily


# Process-level cache: scenario_mode → loaded retriever object.
# Populated lazily on first request for each scenario.
_SCENARIO_RETRIEVERS = {}
_SCENARIO_RETRIEVERS_LOCK = threading.Lock()


def get_scenario_retriever(scenario_mode: str):
    """Return (and lazily load) the system retriever for the given scenario_mode.

    Thread-safe. Falls back to the default RETRIEVAL_PICKLE_PATH if the
    scenario has no dedicated entry or its file cannot be loaded.
    """
    key = str(scenario_mode).strip().lower()
    if key in _SCENARIO_RETRIEVERS:
        return _SCENARIO_RETRIEVERS[key]

    with _SCENARIO_RETRIEVERS_LOCK:
        # Double-checked locking
        if key in _SCENARIO_RETRIEVERS:
            return _SCENARIO_RETRIEVERS[key]

        path = SCENARIO_RETRIEVER_PATHS.get(key) or DEFAULT_RETRIEVAL_PICKLE_PATH
        logger.info(f"🔄 Loading retriever for scenario '{key}' from: {path}")
        if path is None:
            logger.error(f"  Retriever for scenario '{key}' is None and will be ignored!")
            retriever = _GLOBAL_RETRIEVER
        else:
            try:
                retriever = load_retriever_components(path)
                logger.info(f"  ✓ Retriever loaded for scenario '{key}'")
            except Exception as e:
                logger.error(f"  ✗ Failed to load retriever for scenario '{key}': {e}")
                # Fall back to whatever the global default retriever is
                retriever = _GLOBAL_RETRIEVER
        _SCENARIO_RETRIEVERS[key] = retriever
        return retriever


def _init_dp_input_perturber(retriever):
    """Initialise the global DPInputPerturber, reusing the retriever embedding model.

    LangChain FAISS stores the embedding function as vectorstore.embeddings.
    Reusing it avoids loading a second copy of the sentence-transformer weights.
    Falls back to creating a fresh HuggingFaceEmbeddings if that attribute is
    absent (e.g. when the retriever was pickled without it).
    """
    global _DP_INPUT_PERTURBER
    if _DP_INPUT_PERTURBER is not None:
        return  # already initialised

    embed_fn = None
    try:
        embed_fn = retriever.vectorstore.embeddings
        logger.info("DPInputPerturber: reusing vectorstore embedding model.")
    except AttributeError:
        pass

    if embed_fn is None:
        try:
            embed_fn = HuggingFaceEmbeddings(
                model_name=EMBEDDING_MODEL,
                model_kwargs={"device": "cpu"},
                encode_kwargs={"normalize_embeddings": True},
            )
            logger.info("DPInputPerturber: created standalone HuggingFaceEmbeddings.")
        except Exception as exc:
            logger.error("DPInputPerturber: could not create embedding model: %s", exc)
            return

    _DP_INPUT_PERTURBER = DPInputPerturber(epsilon=1.0, embedding_model=embed_fn) # epsilon is placeholder; overridden per-request
    logger.info("DPInputPerturber: initialised and ready.")


def load_retriever(pickle_path=None):
    """Load retriever from component-based pickle (FAISS-compatible)."""
    global _GLOBAL_RETRIEVER

    path = pickle_path or DEFAULT_RETRIEVAL_PICKLE_PATH
    if path is None:
        _init_dp_input_perturber(None)
        return None
    try:
        _GLOBAL_RETRIEVER = load_retriever_components(path)
        logger.info(f"✓ Retriever loaded from {path}")
        _init_dp_input_perturber(_GLOBAL_RETRIEVER)  # initialise Input DP perturber
        return _GLOBAL_RETRIEVER
    except Exception as e:
        logger.error(f"✗ Failed to load retriever from {path}: {e}")
        logger.info("  System will work without RAG.")
        # Still initialise the Input-DP perturber using a standalone embedding
        # model so word-level DP perturbation works even without a RAG retriever.
        _init_dp_input_perturber(None)
        return None


def get_retriever_lazy():
    """Lazy load retriever on first use."""
    global _GLOBAL_RETRIEVER

    if _GLOBAL_RETRIEVER is None:
        logger.info("🔄 Loading retriever for first time...")
        load_retriever()
        logger.info("✓ Retriever loaded")

    return _GLOBAL_RETRIEVER

def load_retriever_async():
    """Load retriever in background thread."""
    thread = threading.Thread(target=load_retriever, daemon=True)
    thread.start()
    return thread

def get_retriever():
    """Get the global retriever instance."""
    return _GLOBAL_RETRIEVER

def _set_system_retriever(new_retriever):
    """Set the global retriever instance and (re-)initialise Input DP perturber."""
    global _GLOBAL_RETRIEVER, _DP_INPUT_PERTURBER
    _GLOBAL_RETRIEVER = new_retriever
    # Reset perturber so it is re-initialised with the new retriever's embedding model
    _DP_INPUT_PERTURBER = None
    if new_retriever is not None:
        _init_dp_input_perturber(new_retriever)

# ============================================================
# SECTION 7 – DEMO MODE (ZERO API CALLS)
# ============================================================

def fabricate_demo_pii(user_text):
    """Fabricate PII detections for demo mode - NO LLM calls."""
    # Simple heuristic: look for common name patterns
    pii_list = []

    # Look for "I am" or "my name is" patterns
    name_patterns = [
        (r"(?:I am|I'm|my name is|call me)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)", "name"),
        (r"\b([A-Z][a-z]+\s+[A-Z][a-z]+)\b", "name"),
    ]

    for pattern, pii_type in name_patterns:
        for match in re.finditer(pattern, user_text, re.IGNORECASE):
            pii_list.append({
                "text": match.group(1),
                "start": match.start(1),
                "end": match.end(1),
                "type": pii_type
            })
            break  # Only first match

    # If no name found, fabricate one from the text
    if not pii_list:
        words = user_text.split()
        if len(words) >= 2:
            # Use first two capitalized words as fake name
            capitalized = [w for w in words if w and w[0].isupper() and len(w) > 2]
            if len(capitalized) >= 2:
                fake_name = f"{capitalized[0]} {capitalized[1]}"
                idx = user_text.find(fake_name)
                if idx != -1:
                    pii_list.append({
                        "text": fake_name,
                        "start": idx,
                        "end": idx + len(fake_name),
                        "type": "name"
                    })

    return pii_list


def fabricate_demo_response(user_text):
    """Fabricate a response for demo mode - NO LLM calls."""
    # Generate a generic but contextual response
    keywords = ["job", "work", "career", "health", "medical", "hobby", "interest"]

    context = "general"
    for kw in keywords:
        if kw in user_text.lower():
            if kw in ["job", "work", "career"]:
                context = "job"
            elif kw in ["health", "medical"]:
                context = "health"
            elif kw in ["hobby", "interest"]:
                context = "hobby"
            break

    responses = {
        "job": "Based on your background and interests, here are some suggestions:\n\n1. Consider roles in the technology sector - your skills would be valuable there.\n2. Look into positions at medium-sized companies (50-200 employees) where you can make an impact.\n3. Check job boards like LinkedIn and Indeed for opportunities in your area.\n\nWould you like me to help you refine your search based on specific criteria?",
        "health": "Here are some preventive health screenings to consider:\n\n1. Annual physical exam with bloodwork (cholesterol, glucose, etc.)\n2. Blood pressure monitoring\n3. Age-appropriate cancer screenings\n4. Dental check-up every 6 months\n\nFor affordable options, check with local community health centers or your insurance provider's network. Would you like more specific recommendations?",
        "hobby": "Based on common interests, here are some hobbies you might enjoy:\n\n1. Outdoor activities like hiking or cycling\n2. Creative pursuits such as photography or painting\n3. Learning new skills through online courses\n4. Community involvement through volunteering\n\nWhat type of activities do you find most appealing?",
        "general": "I'd be happy to help! To provide more personalized recommendations, could you tell me a bit more about:\n\n1. Your current interests or goals\n2. How much time you have available\n3. Any specific areas you'd like to explore\n\nThis will help me give you more relevant suggestions."
    }

    return responses.get(context, responses["general"])


def fabricate_demo_probs(user_text, has_rag):
    """Fabricate probability distributions for demo mode - NO LLM calls."""
    # Generate probabilities based on text analysis
    base_probs = {
        "Gender": {"Male": 0.55, "Female": 0.45},
        "Age bin": {"0-17": 0.05, "18-29": 0.25, "30-44": 0.45, "45-59": 0.20, "60+": 0.05},
        "Marital Status": {"Single": 0.45, "Married": 0.40, "Divorced": 0.10, "Widowed": 0.05},
        "Finance Status": {"Low": 0.25, "Medium": 0.50, "High": 0.25},
        "Education": {"High School": 0.15, "Bachelor's": 0.50, "Master's": 0.25, "PhD": 0.10},
        "Locale": {
            "United States": 0.60, "Canada": 0.15, "United Kingdom": 0.10,
            "Australia": 0.05, "India": 0.04, "Ireland": 0.02,
            "New Zealand": 0.02, "Philippines": 0.02, "Israel": 0.01, "Germany": 0.01, "Italy":0.01,"France":0.01,
        }
    }

    if has_rag:
        # RAG adds confidence - shift probabilities more extremely
        rag_probs = {
            "Gender": {"Male": 0.65, "Female": 0.35},
            "Age bin": {"0-17": 0.02, "18-29": 0.18, "30-44": 0.55, "45-59": 0.20, "60+": 0.05},
            "Marital Status": {"Single": 0.55, "Married": 0.30, "Divorced": 0.10, "Widowed": 0.05},
            "Finance Status": {"Low": 0.15, "Medium": 0.50, "High": 0.35},
            "Education": {"High School": 0.08, "Bachelor's": 0.52, "Master's": 0.30, "PhD": 0.10},
            "Locale": {
                "United States": 0.70, "Canada": 0.08, "United Kingdom": 0.08,
                "Australia": 0.04, "India": 0.02, "Ireland": 0.01,
                "New Zealand": 0.02, "Philippines": 0.01, "Israel": 0.01, "Germany": 0.01, "Italy":0.01,"France":0.01,
            }
        }
        return rag_probs, base_probs
    else:
        return base_probs, base_probs


def fabricate_demo_rag_docs(user_text):
    """Fabricate RAG documents for demo mode - NO retrieval."""
    # Generate fake documents that seem relevant
    templates = [
        "User profile indicates experience in software development with 5 years in web technologies. Previous roles include positions at tech startups in the Bay Area.",
        "Background shows interest in preventive healthcare and fitness. Regular physical activity noted, with preference for outdoor activities.",
        "Career history includes work at medium-sized technology companies. Skills include project management and technical leadership.",
        "Education background: Bachelor's degree in Computer Science. Additional certifications in cloud computing and data analysis.",
        "Location history shows residence in major metropolitan areas. Current location indicates access to urban amenities and services."
    ]

    # Select relevant templates based on keywords
    selected = random.sample(templates, min(3, len(templates)))

    return [(doc, 0.85 - i*0.05) for i, doc in enumerate(selected)]


def process_demo_message(user_input, show_tips, show_rag_hl, show_pii_hl, state, use_rag):
    """Process message in demo mode - ZERO API calls, all data fabricated."""

    # Step 1: Fabricate PII in user input (NO LLM call)
    user_pii_data = fabricate_demo_pii(user_input)
    user_pii_matches = []
    for pii_item in user_pii_data:
        user_pii_matches.append(PIIMatch(
            text=pii_item['text'],
            start=pii_item['start'],
            end=pii_item['end'],
            fine_type=pii_item['type'],
            category=PII_TYPE_TO_CATEGORY.get(pii_item['type'], PIICategory.IDENTITY),
            confidence=0.9
        ))

    # Step 2: Fabricate RAG docs (NO retrieval)
    rag_docs = fabricate_demo_rag_docs(user_input) if use_rag else []

    # Step 3: Fabricate RAG links in user input
    user_rag_links = []
    if use_rag and rag_docs:
        query_words = user_input.lower().split()
        for doc_text, score in rag_docs[:1]:
            doc_words = set(doc_text.lower().split())
            overlap = [w for w in query_words if w in doc_words and len(w) > 3]

            if overlap:
                for word in overlap[:1]:
                    idx = user_input.lower().find(word)
                    if idx != -1:
                        user_rag_links.append(RAGLink(
                            text=word,
                            start=idx,
                            end=idx + len(word),
                            corpus_snippets=[(doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH], score)],
                            top_similarity=score,
                            top_doc_text=doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH],
                            top_doc_score=score,
                            overlap_keywords=overlap[:3]
                        ))
                        break

    demo_evidence = {}
    for attr in SENSITIVE_ATTRIBUTES:
        demo_evidence[attr] = [
            f"Based on context about {attr.lower()} from user input",
            f"Relevant information suggesting {attr.lower()} from documents"
        ]
    state.last_evidence_rag = demo_evidence

    # Add user message
    state.add("user", user_input, user_pii_matches, user_rag_links)

    # Step 4: Fabricate response (NO LLM call)
    llm_response = fabricate_demo_response(user_input)

    # Step 5: Fabricate PII in response (NO LLM call)
    response_pii_data = []  # Typically no PII in response

    # Step 6: Fabricate RAG links in response
    response_rag_links = []
    if use_rag and rag_docs:
        response_words = llm_response.lower().split()
        for doc_text, score in rag_docs:
            doc_words = set(doc_text.lower().split())
            overlap = [w for w in response_words if w in doc_words and len(w) > 4]

            if overlap:
                for word in overlap[:1]:
                    idx = llm_response.lower().find(word)
                    if idx != -1:
                        response_rag_links.append(RAGLink(
                            text=word,
                            start=idx,
                            end=idx + len(word),
                            corpus_snippets=[(doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH], score)],
                            top_similarity=score,
                            top_doc_text=doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH],
                            top_doc_score=score,
                            overlap_keywords=overlap[:3]
                        ))
                        break
                if len(response_rag_links) >= 2:
                    break

    # Add assistant message
    state.add("assistant", llm_response, [], response_rag_links)

    # Step 7: Fabricate probability distributions (NO LLM call)
    probs_rag, probs_no_rag = fabricate_demo_probs(user_input, use_rag)

    state.last_probs_rag = probs_rag
    state.last_probs_no_rag = probs_no_rag

    # Step 8: Calculate risk score
    all_pii = user_pii_matches
    all_rag = user_rag_links + response_rag_links
    risk_score, risk_breakdown = calculate_privacy_risk(all_pii, all_rag, float('inf'))

    # Build outputs
    conv_html = fmt_conversation(state.messages, show_tips, show_rag_hl, show_pii_hl)
    warning_html, demo_inference_metrics, demo_warning_shown = build_inference_warning(
        user_input, probs_rag, probs_no_rag, use_rag, demo_evidence
    )

    # ── Log demo interaction ──────────────────────────────────
    state._turn_count += 1
    state._last_model = "demo"
    state._last_epsilon = float("inf")
    state._last_rag_enabled = use_rag
    state._last_social_scraping = False
    state._last_corpus_source = "demo"
    state._demo_mode = True
    append_interaction_log(
        session_source=getattr(state, "_session_source", "demo"),
        turn_number=state._turn_count,
        demo_mode=True,
        scenario_mode=getattr(state, "_scenario_mode", "real"),
        persona_attributes={},
        model="demo",
        epsilon=float('inf'),
        rag_enabled=use_rag,
        show_risk=getattr(state, "_show_risk", True),
        show_rag_highlights=show_rag_hl,
        show_tips=show_tips,
        show_pii_highlights=show_pii_hl,
        show_settings=getattr(state, "_show_settings", False),
        access_token=getattr(state, "_access_token", ""),
        social_scraping_enabled=False,
        corpus_source="demo",
        uploaded_file_path=None,
        user_prompt=user_input,
        llm_response=llm_response,
        risk_score=risk_score,
        u_rag=user_rag_links,
        r_rag=response_rag_links,
        u_pii=user_pii_matches,
        u_pii_perturbed=0,
        r_pii=[],
        inference_metrics=demo_inference_metrics,
        inference_warning_shown=demo_warning_shown,
        scraped_docs=None,
        num_attributes_changed=0,
        show_dp=getattr(state, "_show_dp", 1),
        show_infr_attr_card=getattr(state, "_show_infr_attr_card", 1),
        show_social_scraping=getattr(state, "_show_social_scraping", False),
        show_upload_data=getattr(state, "_show_upload_data", False),
        rag_corpus_path=getattr(state, "_rag_corpus_path", ""),
    )

    analysis = "### Demo Mode (NO API Calls)\n\nThis is a demonstration with fabricated data. "
    analysis += f"**Detected {len(user_pii_matches)} PII(s)** in user prompt. "
    analysis += f"**{len(user_rag_links)} RAG link(s)** in user text, "
    analysis += f"**{len(response_rag_links)} RAG link(s)** in response.\n\n"
    analysis += "⚠️ All data is fabricated for demonstration purposes.\n\n"
    analysis += "Enable live mode by removing the `demo=1` parameter from the URL."

    return conv_html, create_risk_display(risk_score, risk_breakdown), "", analysis, ""


# ============================================================
# SECTION 8 – ATTRIBUTE INFERENCE & METRICS
# ============================================================

class AttributeInferenceEngine:
    """Extract probabilities from LLM inference responses."""

    # def extract_probabilities_from_response(self, response_text):
    #     """Extract probability distributions from JSON response."""
    #     try:
    #         # Try to extract JSON from the response
    #         json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
    #         if json_match:
    #             result = json.loads(json_match.group(0))
    #             return result
    #     except:
    #         pass
    #     return {}
    def extract_probabilities_from_response(self, response_text):
        """Extract probability distributions AND evidence from JSON response.

        Handles both evidence formats:
          - new: [{"quote": "...", "type": "explicit"}, ...]
          - legacy: ["...", ...]   (plain strings, assumed "implicit")
        """
        try:
            # print("In extract_probabilities_from_response", response_text)
            # Try to extract JSON from the response
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            # print("In extract_probabilities_from_response", json_match)
            if json_match:
                result = json.loads(json_match.group(0))
                # print("********In extract_probabilities_from_response", result)
                # Restructure to separate probabilities and evidence
                probabilities = {}
                evidence = {}      # {attr: [{"quote": str, "type": "explicit"|"implicit"}]}

                for attr, data in result.items():
                    if isinstance(data, dict):
                        # New format with probabilities and evidence
                        if 'probabilities' in data:
                            probabilities[attr] = data['probabilities']
                            raw_ev = data.get('evidence', [])
                            # Normalise evidence items to dicts
                            normalised = []
                            for item in raw_ev:
                                if isinstance(item, dict):
                                    normalised.append({
                                        "quote": str(item.get("quote", item.get("text", ""))).strip(),
                                        "type":  item.get("type", "implicit"),
                                    })
                                else:
                                    # Legacy plain-string evidence
                                    normalised.append({
                                        "quote": str(item).strip(),
                                        "type":  "implicit",
                                    })
                            evidence[attr] = normalised
                        else:
                            # Old format (backward compatibility)
                            probabilities[attr] = data
                            evidence[attr] = []
                    else:
                        # Fallback for unexpected format
                        probabilities[attr] = {}
                        evidence[attr] = []
                probabilities = _fill_missing_with_uniform(probabilities)
                # print('probabilities', probabilities, 'evidence', evidence)
                return {'probabilities': probabilities, 'evidence': evidence}
        except:
            pass
        logger.info("Returning nothing after extracting probabilities from response.....")
        return {'probabilities': _fill_missing_with_uniform({}), 'evidence': {}}

    def get_predictions(self, probs):
        """Get argmax predictions."""
        preds = {}
        for attr, dist in probs.items():
            if dist:
                preds[attr] = max(dist, key=dist.get)
        return preds


class InferentialPrivacyMetrics:
    """Compute privacy metrics."""

    @staticmethod
    def compute_z(p_rag, p_no_rag):
        """Compute log-ratio."""
        if p_no_rag <= 0:
            return float('inf')
        return math.log(p_rag / p_no_rag) if p_rag > 0 else -float('inf')

    @staticmethod
    def compute_lift(z):
        """Compute lift from z."""
        return math.exp(z) - 1.0 if z != float('inf') else float('inf')

    @staticmethod
    def compute_inferential_privacy(probs_rag, probs_no_rag, ground_truth):
        """Compute inferential privacy metrics with robust error handling."""
        metrics = {}

        # Validate inputs are dictionaries
        if not isinstance(probs_rag, dict):
            probs_rag = {}
        if not isinstance(probs_no_rag, dict):
            probs_no_rag = {}
        if not isinstance(ground_truth, dict):
            ground_truth = {}

        for attr in SENSITIVE_ATTRIBUTES:
            try:
                # Check if attribute exists in both probability dictionaries
                if attr not in probs_rag or attr not in probs_no_rag:
                    continue

                # Validate that values are dictionaries
                if not isinstance(probs_rag[attr], dict) or not isinstance(probs_no_rag[attr], dict):
                    continue

                # Get ground truth value
                gt_val = ground_truth.get(attr)
                if not gt_val:
                    continue

                # Get probabilities with default 0.0
                p_rag = probs_rag[attr].get(gt_val, 0.0)
                p_no_rag = probs_no_rag[attr].get(gt_val, 0.0)

                # Validate probabilities are numeric
                if not isinstance(p_rag, (int, float)) or not isinstance(p_no_rag, (int, float)):
                    continue

                # Handle NaN or negative values
                if math.isnan(p_rag) or math.isnan(p_no_rag) or p_rag < 0 or p_no_rag < 0:
                    continue

                # Compute metrics
                z = InferentialPrivacyMetrics.compute_z(p_rag, p_no_rag)
                lift = InferentialPrivacyMetrics.compute_lift(z)

                metrics[attr] = {
                    'z': z,
                    'lift': lift,
                    'p_rag': p_rag,
                    'p_no_rag': p_no_rag
                }

            except (TypeError, ValueError, KeyError, AttributeError):
                # Silently skip this attribute if any error occurs
                continue

        return metrics


# ============================================================
# SECTION 9 – PROMPT BUILDERS
# ============================================================

# def build_inference_prompt(user_prompt, retrieved_docs=None):
#     """Build prompt for attribute inference."""
#     ctx_note = ""
#     if retrieved_docs:
#         doc_snippets = "\n".join(
#             f"  [{i}] {doc.page_content[:200]}..."
#             for i, doc in enumerate(retrieved_docs[:5], 1)
#         )
#         ctx_note = (
#             f"Note: Retrieved {len(retrieved_docs)} documents about the user:\n"
#             f"{doc_snippets}\n\n"
#         )
#
#     return f"""As an AI safety researcher, analyze the following to estimate probability distributions for sensitive user attributes.
#
# {ctx_note}User Query: {user_prompt}
#
# Based on the {'retrieved documents and ' if retrieved_docs else ''}user query, estimate probability distributions for:
#
# {json.dumps(ATTRIBUTE_VALUES_MAP, indent=2)}
#
# Respond with ONLY a JSON object with probability estimates.
# Each attribute maps to {{value: probability}}, probabilities sum to 1.0.
#
# Example format:
# {{
#   "Gender": {{"Male": 0.7, "Female": 0.3}},
#   "Age bin": {{"18-29": 0.2, "30-44": 0.5, "45-59": 0.3, "60+": 0.0, "0-17": 0.0}},
#   ...
# }}
#
# Your JSON response:"""

def build_inference_prompt(user_prompt, retrieved_docs=None, conversation_state=None):
    """Build prompt for attribute inference with evidence extraction.

    Including prior conversation turns ensures that attribute lifts remain
    stable across the session: even when the current turn is benign, the LLM
    can still see sensitive information disclosed in earlier turns, so the
    inferred probabilities (and therefore lifts) reflect the full accumulated
    disclosure rather than only the isolated current message.
    """
    ctx_note = ""
    if retrieved_docs:
        doc_snippets = "\n".join(
            f"  [{i}] {doc.page_content[:200]}..."
            for i, doc in enumerate(retrieved_docs[:5], 1)
        )
        ctx_note = (
            f"Note: Retrieved {len(retrieved_docs)} documents about the user:\n"
            f"{doc_snippets}\n\n"
        )

    # Include prior turns so inference is cumulative across the conversation.
    # Only user messages are included (assistant replies may introduce bias).
    # IMPORTANT: when DP was active for a prior turn, use the PERTURBED text
    # (what the LLM actually saw) rather than the original unperturbed text.
    history_note = ""
    if conversation_state and conversation_state.messages:
        prior = []
        for i, msg in enumerate(conversation_state.messages):
            if msg.role != "user":
                continue
            dp_meta = getattr(msg, "dp_metadata", None)
            if dp_meta and dp_meta.get("perturbed_text"):
                turn_text = dp_meta["perturbed_text"]   # what the LLM saw
            else:
                turn_text = msg.content                  # no DP, use as-is
            prior.append(f"  Turn {i+1}: {turn_text[:300]}")
        if prior:
            history_note = (
                "Prior conversation turns (earlier user messages):\n"
                + "\n".join(prior)
                + "\n\n"
            )

    prior_best_note = ""
    if conversation_state and conversation_state._best_probs_rag:
        lines = []
        for attr, dist in conversation_state._best_probs_rag.items():
            if not dist:
                continue
            top_val = max(dist, key=dist.get)
            top_p = dist[top_val]
            evs = conversation_state._best_evidence_rag.get(attr, [])
            ev_type = "explicit" if any(
                (e.get("type") if isinstance(e, dict) else "implicit") == "explicit"
                for e in evs
            ) else "implicit"
            lines.append(f"  {attr}: {top_val} (confidence {top_p:.0%}, {ev_type} evidence)")
        if lines:
            prior_best_note = (
                    "Previously inferred with highest confidence — treat as prior, "
                    "only revise if strong new contradicting evidence exists:\n"
                    + "\n".join(lines)
                    + "\n\n"
            )

    return f"""As an AI safety researcher, analyze the following to estimate probability distributions for sensitive user attributes.

    {ctx_note}{history_note}{prior_best_note}Current User Query: {user_prompt}

    Based on the {'retrieved documents, prior conversation, and ' if (retrieved_docs or history_note) else ''}current user query, estimate probability distributions for ALL of the following attributes:

    {json.dumps(ATTRIBUTE_VALUES_MAP, indent=2)}

    IMPORTANT: You MUST include every attribute listed above in your response - do not omit any.

    For each attribute, also provide 1 evidence item: a short direct quotation (max 300 chars) supporting your prediction, classified as:
    - "explicit": the attribute is directly stated (e.g., "I am 35 years old")
    - "implicit": inferred from indirect hints (e.g., job or hobbies suggest age)

    Respond with ONLY a valid JSON object. Each attribute must follow this structure:
    {{
      "<AttributeName>": {{
        "probabilities": {{"<Value1>": <float>, "<Value2>": <float>, ...}},
        "evidence": [{{"quote": "<short supporting quote>", "type": "explicit|implicit"}}]
      }},
      ...
    }}
    
    Example (showing one attribute — repeat this structure for all {len(ATTRIBUTE_VALUES_MAP)} attributes):

    {{
      "Gender": {{
        "probabilities": {{"Male": 0.7, "Female": 0.3}},
        "evidence": [{{"quote": "my wife and I went fishing", "type": "implicit"}}]
      }},
      ...
    }}
    
    Rules:
    - Include ALL 6 attributes. Missing any is an error.
    - For each attribute, "probabilities" must contain every value listed in the attribute map above, and they must sum to 1.0.
    - The estimated probabilities should be accurate as possible, rather than rounded up. For example, if you deduce a value of some attribute at probability 0.38, do not round it to 0.4, but report the precise probability 0.38.
    - Use only the exact value strings listed above for each attribute
    - "type" must be exactly "explicit" or "implicit"

Your JSON response:"""


def build_user_response_prompt(user_prompt, retrieved_docs=None, conversation_state=None):
    """
    Build messages list for generating user-visible response with conversation history.

    Returns a list of message dicts suitable for OpenAI-style chat completion API.
    """
    messages = []

    # System message with instructions
    system_content = "You are a helpful assistant. Please provide helpful, personalized responses to the user's queries."
    if retrieved_docs:
        system_content += " Use the relevant context provided to enhance your responses when applicable."
    system_content += f" Keep your response concise (up to {MAX_TOKENS_LLM_RESPONSE_GENERATION} words) and always end with a complete sentence."

    messages.append({"role": "system", "content": system_content})

    # Add conversation history from state (all previous turns)
    if conversation_state and conversation_state.messages:
        for msg in conversation_state.messages:
            # Only include role and content, strip PII/RAG metadata
            messages.append({
                "role": msg.role,
                "content": msg.content
            })

    # Build context from retrieved documents
    ctx = ""
    if retrieved_docs:
        parts = []
        for i, doc in enumerate(retrieved_docs[:10], 1):
            parts.append(f"[{i}] {doc.page_content[:500]}")
        ctx = "\n\nRelevant Context:\n" + "\n".join(parts)

    # Add current user query with optional context
    current_query = user_prompt
    if ctx:
        current_query = f"{user_prompt}\n{ctx}"

    messages.append({"role": "user", "content": current_query})

    return messages


# ============================================================
# SECTION 10 – RAG LINKAGE DETECTION
# ============================================================

# def find_rag_linkage(text, retrieved_docs):
#     """Find spans in text that link to retrieved documents."""
#     if not retrieved_docs:
#         return []
#
#     links = []
#     text_lower = text.lower()
#
#     # Simple keyword-based matching
#     for doc in retrieved_docs:
#         doc_text = doc.page_content if hasattr(doc, 'page_content') else str(doc)
#         doc_words = set(w.lower() for w in doc_text.split() if len(w) > 4)
#
#         # Find overlapping keywords
#         for word in text.split():
#             if len(word) <= 4:
#                 continue
#             word_lower = word.lower().strip('.,!?;:')
#             if word_lower in doc_words:
#                 idx = text_lower.find(word_lower)
#                 if idx != -1:
#                     links.append(RAGLink(
#                         text=word,
#                         start=idx,
#                         end=idx + len(word),
#                         corpus_snippets=[(doc_text[:200], 0.9)],
#                         top_similarity=0.9,
#                         top_doc_text=doc_text[:200],
#                         top_doc_score=0.9,
#                         overlap_keywords=[word_lower]
#                     ))
#                     break
#
#     return links[:5]  # Limit to top 5

def calculate_rouge_l(text1, text2):
    """
    Calculate ROUGE-L (Longest Common Subsequence) similarity between two texts.

    Returns:
        float: ROUGE-L F1 score (0.0-1.0)
    """

    def lcs_length(s1, s2):
        """Calculate longest common subsequence length."""
        m, n = len(s1), len(s2)
        dp = [[0] * (n + 1) for _ in range(m + 1)]

        for i in range(1, m + 1):
            for j in range(1, n + 1):
                if s1[i - 1] == s2[j - 1]:
                    dp[i][j] = dp[i - 1][j - 1] + 1
                else:
                    dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

        return dp[m][n]

    # Tokenize
    tokens1 = text1.lower().split()
    tokens2 = text2.lower().split()

    if not tokens1 or not tokens2:
        return 0.0

    lcs_len = lcs_length(tokens1, tokens2)

    # Calculate ROUGE-L using F1 score
    if lcs_len == 0:
        return 0.0

    precision = lcs_len / len(tokens2)
    recall = lcs_len / len(tokens1)

    if precision + recall == 0:
        return 0.0

    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def calculate_cosine_similarity(text1, text2, doc_embedding=None):
    """
    Calculate cosine similarity between two texts.

    If doc_embedding is provided (from retrieval), use it directly.
    Otherwise, fall back to simple word overlap similarity.

    Args:
        text1: First text string
        text2: Second text string
        doc_embedding: Pre-computed embedding vector for text2 (optional)

    Returns:
        float: Similarity score (0.0-1.0)
    """
    # If we have pre-computed embeddings from the retriever, use them
    if doc_embedding is not None:
        # This would require computing embedding for text1
        # For now, we'll use the fallback method
        pass

    # Fallback: Use word overlap as proxy for cosine similarity
    # This is faster but less accurate than actual embeddings
    words1 = set(re.sub(r'[^\w\s]', '', text1.lower()).split())
    words2 = set(re.sub(r'[^\w\s]', '', text2.lower()).split())

    # Remove stop words
    words1 = words1 - _DP_STOPWORDS
    words2 = words2 - _DP_STOPWORDS

    if not words1 or not words2:
        return 0.0

    intersection = len(words1 & words2)
    union = len(words1 | words2)

    # Jaccard similarity as proxy for cosine similarity
    jaccard = intersection / union if union > 0 else 0.0

    # Scale to approximate cosine similarity behavior
    # Jaccard tends to be lower than cosine, so we adjust
    return min(1.0, jaccard * 1.3)


def combine_similarity_scores(rouge_l, cosine_sim):
    """
    Combine ROUGE-L and cosine similarity into a single score.

    Args:
        rouge_l: ROUGE-L score (0.0-1.0)
        cosine_sim: Cosine similarity score (0.0-1.0)

    Returns:
        float: Combined score (0.0-1.0)
    """
    # Weighted average
    combined = (ROUGE_WEIGHT * rouge_l) + (COSINE_WEIGHT * cosine_sim)
    return combined


def extract_ngrams(text, n):
    """
    Extract meaningful n-grams (sequences of n words) from text.

    Filters out:
    - N-grams that cross sentence boundaries (. ! ?)
    - N-grams with too many stop words (must have at least 50% content words)
    - N-grams that are too short

    Args:
        text: Input text string
        n: Number of words in each n-gram (1, 2, or 3)

    Returns:
        list: List of dicts with 'text', 'start', 'end', 'words'
    """
    # Use regex to find word boundaries
    words_with_positions = []
    for match in re.finditer(r'\b\w+\b', text):
        words_with_positions.append({
            'word': match.group(0),
            'start': match.start(),
            'end': match.end()
        })

    ngrams = []
    for i in range(len(words_with_positions) - n + 1):
        # Get n consecutive words
        ngram_parts = words_with_positions[i:i + n]

        # Build ngram text
        start_pos = ngram_parts[0]['start']
        end_pos = ngram_parts[-1]['end']
        ngram_text = text[start_pos:end_pos]

        # ────────────────────────────────────────────────────────────
        # Filter 1: Skip if too short
        # ────────────────────────────────────────────────────────────
        if len(ngram_text) < MIN_NGRAM_LENGTH:
            continue

        # ────────────────────────────────────────────────────────────
        # Filter 2: Skip if crosses sentence boundaries
        # Check for period, exclamation, or question mark followed by space
        # ────────────────────────────────────────────────────────────
        if any(punct in ngram_text for punct in ['. ', '! ', '? ', '.\n', '!\n', '?\n']):
            continue

        # ────────────────────────────────────────────────────────────
        # Filter 3: Skip if ngram has too many stop words
        # Require at least 50% content words (not stop words)
        # ────────────────────────────────────────────────────────────
        words = [p['word'].lower() for p in ngram_parts]

        # Count content words (non-stop words)
        content_words = [w for w in words if w not in _DP_STOPWORDS]
        content_ratio = len(content_words) / len(words) if words else 0

        # Require at least 50% content words for n-grams with 3+ words
        # For 2-word n-grams, require at least 1 content word
        if n >= 3 and content_ratio < 0.5:
            continue
        elif n == 2 and len(content_words) < 1:
            continue

        # Skip if ALL words are stop words (redundant but explicit)
        if all(w in _DP_STOPWORDS for w in words):
            continue

        ngrams.append({
            'text': ngram_text,
            'start': start_pos,
            'end': end_pos,
            'words': words
        })

    return ngrams


def find_rag_linkage(text, retrieved_docs):
    """
    Find spans in text that link to retrieved documents using similarity metrics.

    This function:
    1. Extracts n-grams (1-3 word sequences) from the text
    2. Compares each n-gram to retrieved documents using ROUGE-L and cosine similarity
    3. Highlights n-grams that exceed similarity thresholds
    4. Filters out stop words and ensures whole word matching

    Args:
        text: Input text to find linkages in
        retrieved_docs: List of retrieved documents from RAG system

    Returns:
        list: List of RAGLink objects with similarity scores
    """
    if not retrieved_docs or not text:
        return []

    links = []
    seen_positions = set()  # Track already highlighted positions to avoid overlaps

    # Build document corpus for comparison
    doc_texts = []
    doc_scores = {}  # Store retrieval scores if available
    doc_sources = {}  # Store document sources
    doc_urls = {}    # Store document source URLs

    for idx, doc in enumerate(retrieved_docs):
        doc_text = doc.page_content if hasattr(doc, 'page_content') else str(doc)
        doc_texts.append(doc_text)

        # Try to extract retrieval score from metadata
        if hasattr(doc, 'metadata') and isinstance(doc.metadata, dict):
            score = doc.metadata.get('score', doc.metadata.get('similarity', None))
            if score is not None:
                doc_scores[idx] = float(score)

            # Extract source information from metadata
            source_type = doc.metadata.get('source', '')
            # print("In find_rag_linkage", source_type)
            if source_type == 'uploaded_csv':
                doc_sources[idx] = "From uploaded data"
            # elif source_type.startswith('social_media_'):
            elif "tweet" in source_type.lower() or "twitter" in source_type.lower():
                doc_sources[idx] = "From Twitter (Internet)"
            elif "facebook" in source_type.lower():
                content_type = doc.metadata.get("content_type", "")
                if content_type == "page_info":
                    doc_sources[idx] = "From Facebook Profile Page (Internet)"
                else:
                    doc_sources[idx] = "From Facebook (Internet)"
            elif "linkedin" in source_type.lower():
                doc_sources[idx] = "From LinkedIn (Internet)"
            elif "web" in source_type.lower():
                doc_sources[idx] = "From the web (Internet)"
            else:
                # Default background corpus (e.g. PANORAMA synthetic social profiles)
                doc_sources[idx] = DEFAULT_CORPUS_SOURCE_LABEL

            # Extract URL from metadata (try several common field names)
            raw_url = (
                doc.metadata.get("source_url") or
                doc.metadata.get("tweet_url") or
                doc.metadata.get("url") or
                ""
            )
            # Exclude sentinel values and non-URL strings
            if raw_url and raw_url not in ("unknown", ""):
                doc_urls[idx] = str(raw_url)
            else:
                doc_urls[idx] = ""

    # Extract n-grams of different lengths (prioritize longer phrases)
    # Only extract X to MAX_NGRAM_WORDS word phrases (no single words)
    candidates = []

    # Extract from longest to shortest (MAX_NGRAM_WORDS down to X)
    for n in range(min(MAX_NGRAM_WORDS, 7), MIN_NGRAM_WORDS - 1, -1):
        for ngram in extract_ngrams(text, n):
            candidates.append((ngram, n))  # (ngram_dict, priority=n)

    # Already sorted by priority (longer phrases extracted first) and position
    candidates.sort(key=lambda x: (-x[1], x[0]['start']))

    # Evaluate each candidate n-gram
    for ngram, priority in candidates:
        ngram_text = ngram['text']
        start = ngram['start']
        end = ngram['end']
        words = ngram['words']

        # Skip if position already highlighted
        if any(pos in seen_positions for pos in range(start, end)):
            continue

        # Skip if all words are stop words
        if all(w in _DP_STOPWORDS for w in words):
            continue

        # Find best matching document
        best_doc_idx = -1
        best_doc_text = ""
        best_rouge = 0.0
        best_cosine = 0.0
        best_combined = 0.0
        best_retrieval_score = 0.0
        best_source = DEFAULT_CORPUS_SOURCE_LABEL  # Default: background corpus
        best_url = ""  # URL of the best matching document

        for doc_idx, doc_text in enumerate(doc_texts):
            # Calculate ROUGE-L similarity
            rouge_l = calculate_rouge_l(ngram_text, doc_text)

            # Calculate cosine similarity (word overlap proxy)
            cosine_sim = calculate_cosine_similarity(ngram_text, doc_text)

            # Combine scores
            combined_score = combine_similarity_scores(rouge_l, cosine_sim)

            # Track best match
            if combined_score > best_combined:
                best_combined = combined_score
                best_rouge = rouge_l
                best_cosine = cosine_sim
                best_doc_idx = doc_idx
                best_doc_text = doc_text
                best_retrieval_score = doc_scores.get(doc_idx, 0.0)
                best_source = doc_sources.get(doc_idx, DEFAULT_CORPUS_SOURCE_LABEL)
                best_url = doc_urls.get(doc_idx, "")

        # Check if similarity exceeds thresholds
        passes_rouge = best_rouge >= ROUGE_L_THRESHOLD
        passes_cosine = best_cosine >= COSINE_SIM_THRESHOLD
        passes_combined = best_combined >= MIN_COMBINED_SCORE

        # Accept if either individual threshold is met OR combined threshold is met
        if passes_combined or (passes_rouge and passes_cosine):
            # Mark positions as used
            for pos in range(start, end):
                seen_positions.add(pos)

            # Create RAGLink with actual similarity scores
            links.append(RAGLink(
                text=ngram_text,
                start=start,
                end=end,
                corpus_snippets=[(best_doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH], best_combined)],
                top_similarity=best_combined,  # Combined score
                top_doc_text=best_doc_text[:RAG_LINKAGE_STORED_EXCERPT_LENGTH],
                top_doc_score=best_retrieval_score if best_retrieval_score > 0 else best_combined,
                overlap_keywords=words,
                source=best_source,
                url=best_url,
            ))

    # Sort by position in text
    links.sort(key=lambda x: x.start)

    return links

# ============================================================
# SECTION 11 – PRIVACY RISK CALCULATION
# ============================================================

def calculate_inferential_privacy_score(inference_metrics):
    """Calculate the inferential privacy score based on the maximum attribute lift.

    The lift for each attribute is computed as:
        lift = p_post / p_pop − 1
    where p_post is the posterior probability given the system's available
    information and p_pop is the population base rate (from POPULATION_PRIORS).

    The score applies a log-based saturation function to the maximum lift:
        score = k·ln(1 + max_lift) / (1 + k·ln(1 + max_lift))
    where k = INFERENTIAL_SCORE_STEEPNESS_K.

    This produces non-zero scores for ALL experimental conditions (including
    those without RAG), enabling meaningful cross-group comparison.

    Args:
        inference_metrics: dict {attr: {lift, prob_rag, prob_no_rag, p_pop, ...}}
                           as produced by build_inference_warning.

    Returns:
        dict with keys:
          score    – float in [0, 1)
          max_lift – float (raw maximum lift value, where 0 = no information gain)
          max_attr – str   (attribute with the highest lift)
          p_rag    – float (posterior probability for max_attr)
          p_pop    – float (population prior for max_attr)
          p_no_rag – float (message-only probability, kept for logging)
    """
    empty = {"score": 0.0, "max_score":0.0, "max_lift": 0.0, "max_attr": "",
             "p_rag": 0.0, "p_pop": 0.0, "p_no_rag": 0.0,
             "mean_score": 0.0, "mean_lift": 0.0, "median_score": 0.0, "median_lift": 0.0}
    if not inference_metrics:
        return empty

    max_lift = 0.0
    max_attr = ""
    max_p_rag = 0.0
    max_p_pop = 0.0
    max_p_no_rag = 0.0

    # ── Collect all positive lifts for mean calculation ────────
    positive_lifts = []

    for attr, m in inference_metrics.items():
        lift = m.get("lift", 0.0)
        if math.isfinite(lift) and lift > 0:
            positive_lifts.append(lift)
        if math.isfinite(lift) and lift > max_lift:
            max_lift = lift
            max_attr = attr
            max_p_rag = m.get("prob_rag", m.get("p_rag", 0.0))
            max_p_pop = m.get("p_pop", 0.0)
            max_p_no_rag = m.get("prob_no_rag", m.get("p_no_rag", 0.0))

    k = INFERENTIAL_SCORE_STEEPNESS_K

    # ── Max-lift score (existing) ─────────────────────────────
    adjusted_max = max(0.0, max_lift)
    if adjusted_max == 0:
        max_score = 0.0
    else:
        u = k * math.log(1.0 + adjusted_max)
        max_score = u / (1.0 + u)

    # ── Mean-lift score (new) ─────────────────────────────────
    if positive_lifts:
        mean_lift = sum(positive_lifts) / len(positive_lifts)
        u_mean = k * math.log(1.0 + mean_lift)
        mean_score = u_mean / (1.0 + u_mean)
    else:
        mean_lift = 0.0
        mean_score = 0.0

    # ── Median-lift score ─────────────────────────────────────
    if positive_lifts:
        sorted_lifts = sorted(positive_lifts)
        n = len(sorted_lifts)
        if n % 2 == 1:
            median_lift = sorted_lifts[n // 2]
        else:
            median_lift = (sorted_lifts[n // 2 - 1] + sorted_lifts[n // 2]) / 2.0
        u_med = k * math.log(1.0 + median_lift)
        median_score = u_med / (1.0 + u_med)
    else:
        median_lift = 0.0
        median_score = 0.0

    return {
        "score": round(mean_score, 4),
        "max_score":      round(max_score, 4),
        "max_lift":   round(max_lift, 4),
        "max_attr":   max_attr,
        "p_rag":      round(max_p_rag, 4),
        "p_pop":      round(max_p_pop, 4),
        "p_no_rag":   round(max_p_no_rag, 4),
        "mean_score": round(mean_score, 4),
        "mean_lift":  round(mean_lift, 4),
        "median_score": round(median_score, 4),
        "median_lift": round(median_lift, 4),
    }


def calculate_privacy_risk(pii_matches, rag_links, epsilon, inference_metrics=None, use_rag=True):
    """Calculate overall privacy risk score (0-100) with monotonic growth across turns.

    Score budget:
      Inference component : 0-55 pts  soft-saturating (never instantly maxes out)
      PII component       : 0-30 pts  cumulative – grows with each turn
      RAG linkage         : 0-15 pts  cumulative – grows with each turn
      DP discount         : 0-15 pts  subtracted when DP is enabled
      ─────────────────────────────────────────────────────
      Total maximum       : 100 pts  (only reachable after several high-risk turns)

    The inference component uses a global exponential-saturation over the
    SUM of above-threshold lifts:

        inference = INFER_MAX * (1 − exp(−total_lift / INFER_SCALE))

    Returns:
        (score, breakdown)
        score     – int 0-100
        breakdown – dict with keys: inference, pii, rag, dp_discount (all floats)
    """
    # ── Inference component: global soft saturation ───────────────────────────
    inference_score = 0.0
    if inference_metrics:
        total_lift = sum(
            min(m.get('lift', 0.0), 10.0)
            for m in inference_metrics.values()
            if (m.get('lift', 0.0) > INFERENCE_LIFT_THRESHOLD
                and not math.isinf(m.get('lift', 0.0)))
        )
        inference_score = INFER_MAX * (1.0 - math.exp(-total_lift / INFER_SCALE))

    # ── PII component ─────────────────────────────────────────────────────────
    if pii_matches:
        pii_score = min(PII_BASE + len(pii_matches) * PII_PER_HIT, PII_MAX)
    else:
        pii_score = 0.0

    # ── RAG linkage component ─────────────────────────────────────────────────
    rag_score = min(len(rag_links) * RAG_PER_HIT, RAG_MAX)

    # ── Differential Privacy discount ─────────────────────────────────────────
    dp_factor = 1.0
    dp_discount = 0.0
    if epsilon != float('inf'):
        dp_factor = DP_FLOOR + (1.0 - DP_FLOOR) * min(1.0, epsilon / MAX_POSSIBLE_EPS)
        # Express the discount as the points removed, for the breakdown display
        dp_discount = (inference_score + rag_score) * (1.0 - dp_factor)

    inference_score_adj = inference_score * dp_factor
    rag_score_adj = rag_score * dp_factor

    total = inference_score_adj + pii_score + rag_score_adj
    score = max(0, min(100, int(round(total))))

    breakdown = {
        "inference": round(inference_score_adj, 1),
        "pii": round(pii_score, 1),
        "rag": round(rag_score_adj, 1),
        "dp_discount": round(dp_discount, 1),
        "inferential_score_info": calculate_inferential_privacy_score(inference_metrics or {}),
    }
    return score, breakdown


# def create_risk_display(risk_score):
#     """Create HTML for risk meter."""
#     if risk_score == 0:
#         label=""
#         color="#4CAF50"
#     elif risk_score < 30:
#         color = "#4CAF50"
#         label = "Low Risk"
#     elif risk_score < 70:
#         color = "#FF9800"
#         label = "Medium Risk"
#     else:
#         color = "#F44336"
#         label = "High Risk"
#
#     return f"""
#     <div style="text-align:center;padding:20px;">
#         <div style="font-size:3em;font-weight:bold;color:{color};">{risk_score}</div>
#         <div style="font-size:1.2em;margin-top:10px;color:#666;">{label}</div>
#         <div style="background:#eee;height:20px;border-radius:10px;margin-top:15px;overflow:hidden;">
#             <div style="background:{color};height:100%;width:{risk_score}%;"></div>
#         </div>
#     </div>
#     """

def create_risk_display(risk_score, breakdown=None):
    """Create HTML for risk meter with shield badge and Why? breakdown tooltip.

    Args:
        risk_score – int 0-100
        breakdown  – optional dict from calculate_privacy_risk:
                     {"inference", "pii", "rag", "dp_discount"}
                     When provided the Why? anchor shows a data-driven explanation.
    """
    if risk_score == 0:
        label = "No Risk"
        color = "#4CAF50"
        icon  = "✓"
    elif risk_score < 30:
        label = "Low Risk"
        color = "#4CAF50"
        icon  = "✓"
    elif risk_score < 70:
        label = "Medium Risk"
        color = "#FF9800"
        icon  = "⚠"
    else:
        label = "High Risk"
        color = "#F44336"
        icon  = "⚠"

    # ── Build the Why? tooltip content ────────────────────────────────────
    if breakdown:
        infer_pts = breakdown.get("inference",   0.0)
        pii_pts   = breakdown.get("pii",         0.0)
        rag_pts   = breakdown.get("rag",         0.0)
        dp_pts    = breakdown.get("dp_discount", 0.0)
        use_rag   = breakdown.get("use_rag",     True)
        infer_info = breakdown.get("inferential_score_info")

        # Dominant driver sentence
        drivers = sorted(
            [("Inference risk",          infer_pts),
             ("PII exposure",            pii_pts),
             ("External data linkages",  rag_pts)],
            key=lambda x: x[1], reverse=True
        )
        top_driver, top_pts = drivers[0]
        if top_pts > 0:
            driver_sentence = (
                f"The primary driver is <b>{top_driver}</b> "
                f"({top_pts:.0f} pts)."
            )
        else:
            driver_sentence = "No significant risk factors were detected yet."

        # Component rows – only show non-zero values
        rows = []
        if infer_pts > 0:
            rows.append(
                f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:3px;">'
                f'<span>🔍 Inference risk</span>'
                f'<span style="font-weight:600;color:#b71c1c;">+{infer_pts:.0f} pts</span>'
                f'</div>'
            )
        if pii_pts > 0:
            rows.append(
                f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:3px;">'
                f'<span>🏷 PII exposure</span>'
                f'<span style="font-weight:600;color:#e65100;">+{pii_pts:.0f} pts</span>'
                f'</div>'
            )
        if rag_pts > 0:
            rows.append(
                f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:3px;">'
                f'<span>🔗 External data linkages</span>'
                f'<span style="font-weight:600;color:#1565c0;">+{rag_pts:.0f} pts</span>'
                f'</div>'
            )
        if dp_pts > 0:
            rows.append(
                f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:3px;">'
                f'<span>🔒 Differential Privacy</span>'
                f'<span style="font-weight:600;color:#2e7d32;">−{dp_pts:.0f} pts</span>'
                f'</div>'
            )

        divider = '<div class="tt-divider"></div>'
        total_row = (
            f'<div style="display:flex;justify-content:space-between;font-weight:700;font-size:1.2em">'
            f'<span class="tt-title">Total</span>'
            f'<span class="tt-title">{risk_score}/100</span>'
            f'</div>'
        )

        # ── Inferential privacy score section ──────────────────────────────────
        inferential_section = ""
        if infer_info:
            i_score    = infer_info["max_score"]
            i_attr     = infer_info["max_attr"]
            i_p_rag    = infer_info["p_rag"]
            i_p_pop    = infer_info.get("p_pop", infer_info.get("p_no_rag", 0.0))

            if i_attr:
                # We have at least one attribute with positive lift
                attr_row = (
                    f'<div class="tt-sub" style="margin-bottom:3px;">'
                    f'Highest-lift attribute: <strong>{_esc(i_attr)}</strong></div>'
                    f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:2px;">'
                    f'<span>P (population prior)</span>'
                    f'<span style="font-weight:600;">{i_p_pop:.1%}</span>'
                    f'</div>'
                    f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:2px;">'
                    f'<span>P (system estimate)</span>'
                    f'<span style="font-weight:600;color:#b71c1c;">{i_p_rag:.1%}</span>'
                    f'</div>'
                    f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:2px;">'
                    f'<span>Inferential score (max)</span>'
                    f'<span style="font-weight:700;color:#b71c1c;">{i_score:.3f}</span>'
                    f'</div>'
                    f'<div class="tt-sub" style="display:flex;justify-content:space-between;margin-bottom:4px;">'
                    f'<span>Inferential score (mean)</span>'
                    f'<span style="font-weight:700;color:#b71c1c;">'
                    f'{infer_info.get("mean_score", 0.0):.3f}</span>'
                    f'</div>'
                )
            else:
                # No positive lift yet (first turn or no RAG)
                attr_row = (
                    f'<div style="font-size:0.82em;color:#888;font-style:italic;line-height:1.4;">'
                    f'No significant external-data boost detected yet. '
                    f'The score will update as more attributes are inferred.'
                    f'</div>'
                )

            inferential_section = (
                f'<div class="tt-divider"></div>'
                f'<div class="tt-title" style="margin-bottom:4px;">🔬 Inferential Privacy Score'
                f'&ensp;<span style="font-weight:700;color:{"#b71c1c" if i_score > 0 else "#888"};">'
                f'{i_score:.3f}</span></div>'
                f'{attr_row}'
            )

        tooltip_body = (
            f'<div class="tt-title">Score breakdown</div>'
            f'{"".join(rows)}'
            f'<div class="tt-divider"></div>'
            f'{total_row}'
            f'<div class="tt-bottom" style="margin-top:8px;line-height:1.4;font-size:1.2em">'
            # f'{driver_sentence}'
            f'{inferential_section}'
            f'</div>'
        )
    else:
        tooltip_body = (
            'Score is calculated from PII exposure, external data linkages, '
            'differential privacy settings, and inference risk from sensitive attributes.'
        )

    why_anchor = (
        f'<span class="tt-anchor" style="position:relative;display:inline-block;'
        f'cursor:help;color:#888;border-bottom:1px dotted #aaa;font-size:1.2em;">'
        f'Risk Score: {int(risk_score)}/100 — Why?'
        f'<span class="tt-popup tt-above" style="width:320px;text-align:left;">'
        f'{tooltip_body}'
        f'</span>'
        f'<span class="tt-popup tt-below" style="width:320px;text-align:left;">'
        f'{tooltip_body}'
        f'</span>'
        f'</span>'
    )

    return f"""
    <div style="text-align:center;padding:20px;">
        <!-- Shield Badge – larger render (150×150) for bigger visible number -->
        <div style="position:relative;display:inline-block;margin-bottom:16px;">
            <svg width="150" height="150" viewBox="0 0 120 140"
                 style="filter:drop-shadow(0 2px 4px rgba(0,0,0,0.2));">
                <!-- Shield shape -->
                <path d="M60 10 L100 30 L100 70 Q100 110 60 130 Q20 110 20 70 L20 30 Z"
                      fill="{color}" stroke="#fff" stroke-width="3"/>
                <!-- Icon -->
                <text x="60" y="50" text-anchor="middle"
                      font-size="22" fill="#fff" font-weight="bold">{icon}</text>
                <!-- Score – larger font fills the shield -->
                <text x="60" y="98" text-anchor="middle"
                      font-size="72" fill="#fff" font-weight="bold">{risk_score}</text>
            </svg>
        </div>

        <!-- Label -->
        <div style="font-size:1.3em;font-weight:bold;margin-bottom:12px;color:{color};">
            {label}
        </div>

        <!-- Progress bar -->
        <div style="background:#eee;height:20px;border-radius:10px;overflow:hidden;
                    max-width:300px;margin:0 auto;">
            <div style="background:{color};height:100%;width:{risk_score}%;
                        transition:width 0.3s ease;"></div>
        </div>

        <!-- Why? anchor with score breakdown tooltip -->
        <div style="margin-top:14px;">
            {why_anchor}
        </div>
    </div>
    """


# ============================================================
# SECTION 12 – CONCURRENT LLM CALLS
# ============================================================

def make_concurrent_llm_calls_optimized(user_input, provider, model_id, use_rag, retrieved, conversation_state=None, original_user_input=None):
    """
    Make MAXIMUM 3 concurrent LLM calls for optimal speed:
    1. Generate response (with conversation history support)
    2. Infer attributes (combined - single call for all attributes)
    3. Detect PII (combined - single call for both user input and response)

    `user_input` is what is sent to the LLM (may be DP-perturbed).
    `original_user_input`, if provided, is used for PII span anchoring so that
    PII detected in the perturbed text is correctly located in the original.

    Args:
        conversation_state: ConversationState object containing message history

    Returns: (llm_response, u_pii, u_pii_perturbed, r_pii, probs_rag, probs_no_rag, evidence_rag)
    """
    logger.info("\n" + "="*60)
    logger.info("⚡ STARTING CONCURRENT LLM CALLS")
    logger.info("="*60)

    client = LLM_CLIENTS.get(provider)
    if not client:
        logger.info("✗ No LLM client available")
        return "", [], [], [], {}, {}, {}, [], []

    engine = AttributeInferenceEngine()
    total_start = time.time()

    # ────────────────────────────────────────────────────────────
    # CALL 1: Generate response
    # ────────────────────────────────────────────────────────────
    def call_generate_response():
        """LLM CALL 1/3: Generate user-visible response with conversation history."""
        logger.info("  🤖 [LLM CALL] Generating response...")
        start_time = time.time()
        try:
            # Build messages list including conversation history
            messages = build_user_response_prompt(
                user_input,
                retrieved if use_rag else None,
                conversation_state=conversation_state
            )

            # resp = client.responses.create(
            resp = client.chat.completions.create(
                model=model_id,
                messages=messages,  # Now a list of message dicts with history
                # max_tokens=MAX_TOKENS_LLM_RESPONSE_GENERATION,
                max_completion_tokens=MAX_TOKENS_LLM_RESPONSE_GENERATION,
                temperature=TEMPERATURE_LLM_RESPONSE_GENERATION,
                timeout=30,
            )
            result = resp.choices[0].message.content # .output_text
            # If the response was cut off at the token limit, trim to last complete sentence
            if resp.choices[0].finish_reason == "length":
                last_end = max(result.rfind('.'), result.rfind('!'), result.rfind('?'))
                if last_end > len(result) // 2:  # only trim if a sentence end exists in the second half
                    result = result[:last_end + 1] + " [...]"
            elapsed = time.time() - start_time
            logger.info(f"  ✓ Response generated in {elapsed:.2f}s (with {len(messages)} messages in context)")
            return result
        except Exception as e:
            logger.info(f"  ✗ Error generating response: {e}")
            return f"Error: {e}"

    # ────────────────────────────────────────────────────────────
    # CALL 2: Attribute inference (COMBINED - single call)
    # ────────────────────────────────────────────────────────────
    # def call_attribute_inference():
    #     """LLM CALL 2/3: Infer ALL attributes in a single call."""
    #     logger.info("  📊 [LLM CALL 2/3] Inferring all attributes (combined)...")
    #     start_time = time.time()
    #     try:
    #         # With RAG if enabled
    #         if use_rag and retrieved:
    #             inf_prompt = build_inference_prompt(user_input, retrieved_docs=retrieved)
    #             resp = client.chat.completions.create(
    #                 model=model_id,
    #                 messages=[{"role": "user", "content": inf_prompt}],
    #                 max_tokens=MAX_TOKENS_INFER_ATTRIBUTES,
    #                 temperature=TEMPERATURE_INTERNAL_TASKS
    #             )
    #             probs_with_rag = engine.extract_probabilities_from_response(resp.choices[0].message.content)
    #         else:
    #             probs_with_rag = {}
    #
    #         # Without RAG (baseline)
    #         inf_prompt_no_rag = build_inference_prompt(user_input, retrieved_docs=None)
    #         resp_no_rag = client.chat.completions.create(
    #             model=model_id,
    #             messages=[{"role": "user", "content": inf_prompt_no_rag}],
    #             max_tokens=MAX_TOKENS_INFER_ATTRIBUTES,
    #             temperature=TEMPERATURE_INTERNAL_TASKS
    #         )
    #         probs_without_rag = engine.extract_probabilities_from_response(resp_no_rag.choices[0].message.content)
    #
    #         # If RAG probs empty, use no_rag probs
    #         if not probs_with_rag:
    #             probs_with_rag = probs_without_rag
    #
    #         elapsed = time.time() - start_time
    #         logger.info(f"  ✓ Attribute inference completed in {elapsed:.2f}s")
    #         return probs_with_rag, probs_without_rag
    #     except Exception as e:
    #         logger.info(f"  ✗ Error in attribute inference: {e}")
    #         return {}, {}

    def call_inference_with_rag():
        """LLM CALL 2a/4: Infer attributes WITH RAG context."""
        # Skip immediately if RAG is not in use — saves a full LLM round-trip.
        if not (use_rag and retrieved):
            return {}, {}
        logger.info("  📊 [LLM CALL] Inferring attributes WITH RAG...")
        start_time = time.time()
        try:
            inf_prompt = build_inference_prompt(
                user_input,
                retrieved_docs=retrieved,
                conversation_state=conversation_state,
            )
            # resp = client.responses.create(
            resp = client.chat.completions.create(
                model=model_id,
                messages=[{"role": "user", "content": inf_prompt}],
                max_tokens=MAX_TOKENS_INFER_ATTRIBUTES,
                # temperature=TEMPERATURE_INTERNAL_TASKS,
                timeout=20,
            )
            result = engine.extract_probabilities_from_response(resp.choices[0].message.content)
            elapsed = time.time() - start_time
            logger.info(f"  ✓ Inference WITH RAG completed in {elapsed:.2f}s")
            return result.get("probabilities", {}), result.get("evidence", {})
        except Exception as e:
            logger.warning(f"  ✗ Inference WITH RAG failed: {e}")
            return {}, {}

    def call_inference_no_rag():
        """LLM CALL 2b/4: Infer attributes WITHOUT RAG (baseline)."""
        logger.info("  📊 [LLM CALL] Inferring attributes WITHOUT RAG (baseline)...")
        start_time = time.time()
        try:
            inf_prompt_no_rag = build_inference_prompt(
                user_input,
                retrieved_docs=None,
                conversation_state=conversation_state,
            )
            # resp_no_rag = client.responses.create(
            resp_no_rag = client.chat.completions.create(
                model=model_id,
                messages=[{"role": "user", "content": inf_prompt_no_rag}],
                max_tokens=MAX_TOKENS_INFER_ATTRIBUTES,
                # temperature=TEMPERATURE_INTERNAL_TASKS,
                timeout=20,
            )
            result = engine.extract_probabilities_from_response(resp_no_rag.choices[0].message.content)
            elapsed = time.time() - start_time
            logger.info(f"  ✓ Inference WITHOUT RAG completed in {elapsed:.2f}s")
            return result.get("probabilities", {}), result.get("evidence", {})
        except Exception as e:
            logger.warning(f"  ✗ Inference WITHOUT RAG failed: {e}")
            return {}, {}

    # ────────────────────────────────────────────────────────────
    # Submit CALL 1, 2a, and 2b concurrently.
    # CALL 3 (PII) requires the LLM response text, so it is submitted
    # as soon as future_response resolves — still overlapping with the
    # two inference futures that are still in-flight.
    # ────────────────────────────────────────────────────────────
    results = {}
    pii_source_text = original_user_input if original_user_input is not None else user_input
    pii_source_text_perturbed = user_input

    with ThreadPoolExecutor(max_workers=6) as executor:
        future_response  = executor.submit(call_generate_response)
        future_rag_inf   = executor.submit(call_inference_with_rag)
        future_norag_inf = executor.submit(call_inference_no_rag)

        # u_rag only needs user_input and retrieved — both available now,
        # so submit immediately and let it overlap with the LLM calls.
        u_rag_source_text = original_user_input if original_user_input is not None else user_input
        future_u_rag = (
            executor.submit(find_rag_linkage, u_rag_source_text, retrieved)
            if retrieved else None
        )

        # Collect the response first (needed to launch PII detection and r_rag).
        try:
            results['response'] = future_response.result(timeout=40)
        except Exception as e:
            logger.warning(f"  ✗ LLM response call timed out or failed: {e}")
            results['response'] = "Error: request timed out. Please try again."

        # As soon as the response is available, submit PII detection and r_rag
        # so they overlap with any still-running inference futures.
        llm_response = results['response']
        if not _is_llm_error(llm_response):
            future_pii = executor.submit(
                detect_pii_combined, pii_source_text, llm_response, provider, model_id
            )
            logger.info(f"PII DETECTION PERTURBED DP: {pii_source_text_perturbed}")
            future_pii_perturbed = (
                executor.submit(detect_pii_combined, pii_source_text_perturbed, "", provider, model_id)
                if original_user_input is not None
                else None
            )
            future_r_rag = (
                executor.submit(find_rag_linkage, llm_response, retrieved)
                if retrieved else None
            )
        else:
            future_pii   = None
            future_r_rag = None

        # Collect inference results (may already be done by now).
        try:
            probs_with_rag, evidence_with_rag = future_rag_inf.result(timeout=40)
        except Exception as e:
            logger.warning(f"  ✗ Inference WITH RAG timed out or failed: {e}")
            probs_with_rag, evidence_with_rag = {}, {}

        try:
            # probs_no_rag, _ = future_norag_inf.result(timeout=40)
            probs_no_rag, evidence_no_rag = future_norag_inf.result(timeout=40)
        except Exception as e:
            logger.warning(f"  ✗ Inference WITHOUT RAG timed out or failed: {e}")
            probs_no_rag = {}
            evidence_no_rag = {}

        # Fall back: if the with-RAG call returned nothing, copy the baseline.
        if not probs_with_rag:
            probs_with_rag = probs_no_rag
            evidence_with_rag = evidence_no_rag

        # Collect PII result.
        if future_pii is not None:
            try:
                user_pii_items, response_pii_items = future_pii.result(timeout=40)
            except Exception as e:
                logger.warning(f"  ✗ PII detection timed out or failed: {e}")
                user_pii_items, response_pii_items = [], []
        else:
            user_pii_items, response_pii_items = [], []

        if future_pii_perturbed is not None:
            try:
                user_pii_perturbed_items, _ = future_pii_perturbed.result(timeout=40)
                logger.info(f"PII DETECTION PERTURBED DP DETECTED: {user_pii_perturbed_items}")
            except Exception as e:
                logger.warning(f"  ✗ Perturbed PII detection failed: {e}")
                user_pii_perturbed_items = []
        else:
            user_pii_perturbed_items = None  # no DP

        # Collect RAG-linkage results (likely already done by now).
        try:
            u_rag = future_u_rag.result(timeout=30) if future_u_rag else []
        except Exception as e:
            logger.warning(f"  ✗ u_rag linkage timed out or failed: {e}")
            u_rag = []

        try:
            r_rag = future_r_rag.result(timeout=30) if future_r_rag else []
        except Exception as e:
            logger.warning(f"  ✗ r_rag linkage timed out or failed: {e}")
            r_rag = []

    probs_rag    = probs_with_rag
    probs_no_rag = probs_no_rag
    evidence_rag = evidence_with_rag

    if _is_llm_error(llm_response):
        # Skip all analysis when the LLM returned an error string
        logger.warning("LLM response is an error — skipping PII/inference analysis.")
        return llm_response, [], [], [], {}, {}, {}, [], []

    # ── Regex fallback for high-confidence PII patterns the LLM may have missed ──
    # Only fills gaps; never overwrites items the LLM already found for the user text.
    # _llm_found_values = {it.get("value", "").lower() for it in user_pii_items if isinstance(it, dict)}

    # _NAME_PATTERNS = [
    #     r"(?:I(?:'m| am)|[Mm]y name is|[Cc]all me|[Tt]his is)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)",
    #     r"(?:^|\.\s+)([A-Z][a-z]+\s+[A-Z][a-z]+)(?:\s+here\b|\s+speaking\b|,)",
    # ]
    # _EMAIL_RE = re.compile(r'\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b')
    # _PHONE_RE = re.compile(r'(?<!\d)(?:\+?\d[\d\s\-().]{7,}\d)(?!\d)')
    #
    # for _pat in _NAME_PATTERNS:
    #     for _m in re.finditer(_pat, pii_source_text):
    #         _name = _m.group(1).strip()
    #         if _name.lower() not in _llm_found_values:
    #             user_pii_items.append({"type": "name", "value": _name, "confidence": 0.92})
    #             _llm_found_values.add(_name.lower())
    #
    # for _m in _EMAIL_RE.finditer(pii_source_text):
    #     _val = _m.group(0)
    #     if _val.lower() not in _llm_found_values:
    #         user_pii_items.append({"type": "email", "value": _val, "confidence": 0.99})
    #         _llm_found_values.add(_val.lower())
    #
    # for _m in _PHONE_RE.finditer(pii_source_text):
    #     _val = _m.group(0).strip()
    #     if sum(c.isdigit() for c in _val) >= 7 and _val.lower() not in _llm_found_values:
    #         user_pii_items.append({"type": "phone number", "value": _val, "confidence": 0.90})
    #         _llm_found_values.add(_val.lower())

    # Convert to PIIMatch objects (anchored to original / pii_source_text)
    u_pii = _extract_pii_spans_from_values(pii_source_text, user_pii_items)
    r_pii = _extract_pii_spans_from_values(llm_response, response_pii_items)
    # u_pii_perturbed: when DP was active use separately detected perturbed items;
    # when DP was not active it is identical to u_pii by definition.
    if user_pii_perturbed_items is not None:
        # Always anchor perturbed PIIs to the perturbed text (user_input),
        # regardless of whether original_user_input differs.
        u_pii_perturbed = _extract_pii_spans_from_values(user_input, user_pii_perturbed_items)
    else:
        # No separate perturbed detection was run (DP was off) — reuse original.
        u_pii_perturbed = u_pii

    total_elapsed = time.time() - total_start
    logger.info("="*60)
    logger.info(f"✅ ALL CONCURRENT CALLS COMPLETED in {total_elapsed:.2f}s")
    logger.info(f"   (LLM calls executed in parallel: response + 2x inference + PII detection)")
    logger.info("="*60 + "\n")

    # print(llm_response, u_pii, r_pii, probs_rag, probs_no_rag, evidence_rag)
    return llm_response, u_pii, u_pii_perturbed, r_pii, probs_rag, probs_no_rag, evidence_rag, u_rag, r_rag


def extract_user_name_from_query(query):
    """
    Extract a potential user name from the query for social media scraping.
    This is a simple heuristic - in production, use NER or better extraction.

    Args:
        query: User query string

    Returns:
        str: Extracted name or a generic identifier
    """
    # Simple pattern: look for "I am [Name]" or similar patterns
    patterns = [
        r"(?:I am|I'm|My name is|This is)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)",
        r"(?:about|for)\s+([A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)",
    ]

    for pattern in patterns:
        match = re.search(pattern, query)
        if match:
            name = match.group(1).strip()
            logger.info(f"  Extracted name from query: '{name}'")
            return name

    # If no name found, return a generic identifier
    logger.info("  No specific name found in query, using generic 'user'")
    return "user"


# ============================================================
# SECTION 13 – MESSAGE PROCESSING (WITH CONCURRENT CALLS)
# ============================================================

def process_message(user_input, use_rag, epsilon, model_str, show_tips, show_rag_hl, show_pii_hl, state, retriever, social_scraping_enabled=False, social_platforms=None):
    """Process a user message with OPTIMIZED concurrent LLM calls (max 3 calls) and optional social media scraping."""

    if not user_input or not user_input.strip():
        return (
            fmt_conversation(state.messages, show_tips, show_rag_hl, show_pii_hl),
            create_risk_display(0),
            "",
            "",
            _build_privacy_settings_html(epsilon),
        )

    # Parse model string to get provider and model_id
    provider = DEFAULT_MODEL_PROVIDER
    model_id = DEFAULT_MODEL_ID
    for p, mid, mname in MODEL_CONFIGS:
        if mname in model_str or mid in model_str:
            provider = p
            model_id = mid
            break

    logger.info("\n" + "="*60)
    logger.info(f"🚀 PROCESSING MESSAGE (Model: {model_str})")
    logger.info("="*60)

    # ────────────────────────────────────────────────────────────
    # STEP: Input Differential Privacy (RANTEXT-inspired, word-level LDP)
    # Perturbs content words in the user message BEFORE the LLM sees it.
    # The conversation panel always displays the ORIGINAL message; a DP badge
    # on each user bubble transparently communicates the perturbation.
    # Literature: Tong et al. (2025) InferDPT/RANTEXT, IEEE TDSC 22(5).
    # ────────────────────────────────────────────────────────────
    llm_input = user_input  # default: send original if DP off
    input_dp_substitutions = []

    # Lazy-init the perturber here if it is still None (e.g. retriever failed
    # to load at startup but an embedding model may still be available).
    if epsilon != float("inf") and _DP_INPUT_PERTURBER is None:
        _init_dp_input_perturber(retriever)

    if epsilon != float("inf") and _DP_INPUT_PERTURBER is not None:
        try:
            # Use a shallow copy so the shared singleton's epsilon is never
            # mutated. Two concurrent requests with different epsilon values
            # would otherwise race between the assignment and perturb().
            _local_perturber = copy.copy(_DP_INPUT_PERTURBER)
            _local_perturber.epsilon = epsilon
            perturbed_input, input_dp_substitutions = _local_perturber.perturb(user_input)
            llm_input = perturbed_input
            logger.info(
                "  Input DP: %d word(s) perturbed (ε=%.2f) % s",
                len(input_dp_substitutions), epsilon, perturbed_input
            )
        except Exception as _dp_exc:
            logger.warning("  Input DP perturbation failed (%s) – using original.", _dp_exc)
            llm_input = user_input

    logger.info(f">>>>>>> LLM INPUT: {llm_input}")
    # ────────────────────────────────────────────────────────────
    # STEP: RAG retrieval (LOCAL - no LLM call)
    # ────────────────────────────────────────────────────────────
    logger.info("📁 [LOCAL] RAG retrieval...")
    retrieval_start = time.time()
    retrieved = []
    u_rag = []
    logger.info(f"[LOCAL] RAG retrieval use_rag={use_rag}; retriever={retriever}")
    if use_rag and retriever is not None:
        try:
            min_sim = epsilon_adjusted_threshold(epsilon)
            logger.info("  Using threshold-filtered retrieval (min_sim=%.3f, ε=%s)", min_sim, epsilon)
            if epsilon != float('inf'):
                # Layer 1: threshold filter
                # Layer 2: exponential-mechanism sampling (Retrieval DP)
                pre_filtered = retriever.retrieve(llm_input, k=RAG_FETCH_K, fetch_k=RAG_FETCH_K, min_similarity=min_sim)
                logger.info(f">>>>>>> RAG RETRIEVED pre_filtered: {retrieved}")
                if pre_filtered:
                    dp_retriever = DPRetriever(retriever, epsilon=epsilon)
                    # Pass pre-filtered docs only — DPRetriever re-samples from them
                    retrieved = dp_retriever.retrieve(llm_input, k=RAG_K, fetch_k=len(pre_filtered))
                    logger.info(f">>>>>>> RAG RETRIEVED post_filtered: {retrieved}")
                else:
                    retrieved = []
            else:
                retrieved = retriever.retrieve(llm_input, k=RAG_K, fetch_k=RAG_FETCH_K, min_similarity=min_sim)
                logger.info(f">>>>>>> RAG RETRIEVED: {retrieved}")
                retrieval_elapsed = time.time() - retrieval_start
                logger.info(f"  ✓ Retrieved {len(retrieved)} docs in {retrieval_elapsed:.2f}s (LOCAL)")
        except Exception as e:
            logger.info(f"  ✗ RAG retrieval error: {e}")
            retrieved = []
    else:
        logger.info("  • RAG disabled")

    # ────────────────────────────────────────────────────────────
    # STEP: Social Media Scraping (if enabled and retriever exists)
    # ────────────────────────────────────────────────────────────
    # ────────────────────────────────────────────────────────────
    # STEP 1.5: Social Media Scraping (if enabled)
    #
    # Two modes:
    #   a) use_rag=True  – scraped docs are appended to the existing retrieved set
    #   b) use_rag=False – scraped docs become the sole retrieved set and are
    #                      treated as if RAG were enabled for this request only
    # ────────────────────────────────────────────────────────────
    # effective_* variables let us override use_rag locally without touching
    # the caller's state
    effective_use_rag = use_rag
    effective_retrieved = retrieved  # same list reference; extended below
    scraped_docs_for_log = []        # raw scraped documents saved for interaction log

    if social_scraping_enabled:
        if state._scraped_docs is not None:
            # Already scraped in a previous turn — reuse cached result
            scraped_docs = state._scraped_docs
            logger.info(
                f"🌐 [LOCAL] Reusing {len(scraped_docs)} cached scraped docs (scraping already done this session)")
        else:
            logger.info("🌐 [LOCAL] Social media scraping (first turn)...")
            scraped_docs = []
            try:
                user_name = extract_user_name_from_query(user_input)
                platforms_to_scrape = []
                if social_platforms:
                    platforms_to_scrape = [p.lower() for p in social_platforms]
                    logger.info(f"  Platforms selected: {platforms_to_scrape}")
                if platforms_to_scrape and SOCIAL_SCRAPER and SOCIAL_SCRAPER.enabled:
                    scraped_docs = SOCIAL_SCRAPER.scrape_user_data(user_name, platforms_to_scrape,
                                                                   fallback_data=_load_fallback_tweets(
                                                                       getattr(state, "_scenario_mode", "real")))
            except Exception as exc:
                logger.error(f"  ✗ Social media scraping error: {exc}")
                import traceback
                logger.error(traceback.format_exc())
            # Cache result regardless (even [] means "we tried, nothing found")
            state._scraped_docs = scraped_docs

        if scraped_docs:
            scraped_docs_for_log = scraped_docs
            if use_rag:
                original_count = len(effective_retrieved)
                effective_retrieved = list(effective_retrieved) + scraped_docs
                logger.info(
                    f"  ✓ Appended {len(scraped_docs)} scraped docs to RAG "
                    f"(was {original_count}, now {len(effective_retrieved)})"
                )
            else:
                effective_retrieved = scraped_docs
                effective_use_rag = True
                logger.info(
                    f"  ✓ RAG was off; using {len(scraped_docs)} scraped docs "
                    f"as the RAG source for this turn"
                )
            u_rag = []  # computed concurrently inside make_concurrent_llm_calls_optimized
        else:
            logger.info("  • No scraped documents available")
    else:
        logger.info("  • Social media scraping disabled")

    # ────────────────────────────────────────────────────────────
    # STEP 2: Make ALL concurrent LLM calls (max 3 calls)
    # ────────────────────────────────────────────────────────────
    llm_response, u_pii, u_pii_perturbed, r_pii, probs_rag, probs_no_rag, evidence_rag, u_rag, r_rag = make_concurrent_llm_calls_optimized(
        llm_input, provider, model_id, effective_use_rag, effective_retrieved,
        conversation_state=state,
        original_user_input=user_input,   # PII always detected on original text
    )

    probs_rag_raw = probs_rag  # only LLM-inferred attrs → for build_inference_warning
    probs_rag_display = _fill_missing_with_uniform(probs_rag)  # all 6 attrs → for avatar card

    # ── Accumulate best (highest-confidence) per-attribute predictions ──
    for attr, dist in probs_rag.items():
        if not dist:
            continue
        new_top_p = max(dist.values())
        old_top_p = max(state._best_probs_rag.get(attr, {}).values(), default=0.0)
        if new_top_p >= old_top_p:
            state._best_probs_rag[attr] = dist
            state._best_evidence_rag[attr] = evidence_rag.get(attr, [])

    # ── Early exit on LLM error: add the error message to the conversation
    # but do NOT update any privacy state, risk score, or analysis panels.
    if _is_llm_error(llm_response):
        state.add("user", user_input, [], [], dp_metadata=None)
        state.add("assistant", llm_response, [], [])
        conv_html = fmt_conversation(state.messages, show_tips, show_rag_hl, show_pii_hl)
        privacy_settings_text = _build_privacy_settings_html(epsilon)
        logger.warning("process_message: LLM error detected — risk panel suppressed.")
        # Return a sentinel risk_score of None so _send knows to hide the panel
        return conv_html, None, "", "", privacy_settings_text

    # ────────────────────────────────────────────────────────────
    # STEP 3: Post-processing (LOCAL - no LLM calls)
    # ────────────────────────────────────────────────────────────
    logger.info("🔧 [LOCAL] Post-processing...")

    # Build dp_metadata for the user message bubble
    dp_meta = None
    if input_dp_substitutions:
        dp_meta = {
            "epsilon":           epsilon,
            "num_substitutions": len(input_dp_substitutions),
            "substitutions":     input_dp_substitutions,
            "original_text":     user_input,   # what the user typed (shown in chat)
            "perturbed_text":    llm_input,    # what the LLM received
        }

    # Add messages to state (original text always stored; perturbed text goes to LLM only)
    if dp_meta:
        # Re-anchor: take PII *values* detected on the perturbed text and
        # search for them in the original text so start/end offsets are correct
        # for highlighting inside m.content (which is the original text).
        perturbed_pii_as_items = [
            {"type": p.fine_type, "value": p.text, "confidence": p.confidence}
            for p in u_pii_perturbed
        ]
        u_pii_for_bubble = _extract_pii_spans_from_values(user_input, perturbed_pii_as_items)
        u_pii_for_card = u_pii_perturbed  # perturbed-text spans for the right-panel card
    else:
        u_pii_for_bubble = u_pii
        u_pii_for_card = u_pii

    state.add("user", user_input, u_pii_for_bubble, u_rag, dp_metadata=dp_meta, pii_card_matches=u_pii_for_card)
    # state.add("user", user_input, u_pii, u_rag, dp_metadata=dp_meta)

    # r_rag was computed concurrently inside make_concurrent_llm_calls_optimized.

    state.add("assistant", llm_response, r_pii, r_rag)

    # Store probabilities
    state.last_probs_rag = probs_rag_display
    state.last_probs_no_rag = probs_no_rag
    state.last_evidence_rag = evidence_rag

    # Calculate risk and build outputs
    all_pii = (u_pii or []) + (r_pii or [])
    all_rag = (u_rag or []) + (r_rag or [])

    # ── Build the inference warning from the current turn ─────────────────
    warning_html, inference_metrics, inference_warning_shown = build_inference_warning(
        user_input, state._best_probs_rag, probs_no_rag, effective_use_rag, state._best_evidence_rag, effective_retrieved,
        input_dp_metadata=dp_meta,
        perturbed_user_input=llm_input if dp_meta else None,
        conversation_state=state,
    )

    # ── Update peak (highest-lift) inference metrics across all turns ──────
    for attr, metrics in inference_metrics.items():
        current_lift = metrics.get('lift', 0)
        peak_lift = state._peak_inference_metrics.get(attr, {}).get('lift', -float('inf'))
        if current_lift > peak_lift:
            state._peak_inference_metrics[attr] = metrics

    # ── Compute risk from the FULL conversation, not just this turn ────────
    # Collect all PII and RAG links stored across every message in the session
    session_pii = [m for msg in state.messages for m in (msg.pii_matches or [])]
    session_rag = [r for msg in state.messages for r in (msg.rag_links or [])]
    risk_score, risk_breakdown = calculate_privacy_risk(
        session_pii, session_rag, epsilon, state._peak_inference_metrics,
        use_rag=effective_use_rag
    )

    conv_html = fmt_conversation(state.messages, show_tips, show_rag_hl, show_pii_hl)
    # analysis = _build_analysis(u_pii, u_rag, risk_score, effective_use_rag, epsilon, probs_rag, probs_no_rag)
    # session_user_pii = [m for msg in state.messages if msg.role == "user" for m in (msg.pii_matches or [])]
    session_user_pii = [m for msg in state.messages if msg.role == "user"for m in (msg.pii_card_matches if msg.pii_card_matches is not None else (msg.pii_matches or []))]

    # Build and cache the inference card HTML independently so _send can control its visibility
    # avatar_html = _build_inferred_avatar(probs_rag_display, warning_html, inference_warning_shown)
    avatar_html = _build_inferred_avatar(_fill_missing_with_uniform(state._best_probs_rag), warning_html, inference_warning_shown)
    state._last_avatar_html = avatar_html

    analysis = _build_analysis(session_user_pii, u_rag, risk_score, effective_use_rag, epsilon, probs_rag_display,
                               probs_no_rag, inference_warning_html=warning_html, is_warning=inference_warning_shown,
                               include_avatar=False, show_pii=show_pii_hl)
    privacy_settings_text = _build_privacy_settings_html(epsilon)
    state._last_privacy_html = privacy_settings_text
    logger.info("✓ Post-processing completed")
    logger.info("=" * 60 + "\n")

    # ── Persist interaction to CSV log ────────────────────────
    state._turn_count += 1

    # corpus_source: prefer uploaded file, then social scrape, else system
    if state._uploaded_file_path:
        corpus_source = "uploaded_csv"
    elif social_scraping_enabled and effective_retrieved and not use_rag:
        corpus_source = "social_scrape"
    else:
        corpus_source = "system"

    # ── Compute how many inferred attribute top-values changed this turn ──
    curr_attr_vals = {attr: m.get("top_value", "") for attr, m in inference_metrics.items()}
    num_attributes_changed = sum(
        1 for a in set(state._prev_inference_attrs) | set(curr_attr_vals)
        if state._prev_inference_attrs.get(a) != curr_attr_vals.get(a)
    )
    state._prev_inference_attrs = curr_attr_vals
    state._last_model = model_str
    state._last_epsilon = epsilon
    state._last_rag_enabled = effective_use_rag
    state._last_social_scraping = social_scraping_enabled
    state._last_corpus_source = corpus_source
    state._demo_mode = False
    append_interaction_log(
        session_source=state._session_source,
        turn_number=state._turn_count,
        demo_mode=False,
        scenario_mode=getattr(state, "_scenario_mode", "real"),
        persona_attributes=getattr(state, "_persona_attributes", {}),
        model=model_str,
        epsilon=epsilon,
        rag_enabled=effective_use_rag,
        show_risk=getattr(state, "_show_risk", True),
        show_rag_highlights=show_rag_hl,
        show_tips=show_tips,
        show_pii_highlights=getattr(state, "_show_pii_hl", show_pii_hl),
        show_settings=getattr(state, "_show_settings", False),
        access_token=state._access_token,
        social_scraping_enabled=social_scraping_enabled,
        corpus_source=corpus_source,
        uploaded_file_path=state._uploaded_file_path,
        user_prompt=user_input,
        llm_response=llm_response,
        risk_score=risk_score,
        u_rag=u_rag,
        r_rag=r_rag,
        u_pii=u_pii,
        u_pii_perturbed=u_pii_perturbed,
        r_pii=r_pii,
        inference_metrics=inference_metrics,
        inference_warning_shown=inference_warning_shown,
        scraped_docs=scraped_docs_for_log if social_scraping_enabled else None,
        user_prompt_perturbed=llm_input,
        num_input_dp_substitutions=len(input_dp_substitutions),
        num_attributes_changed=num_attributes_changed,
        show_dp=getattr(state, "_show_dp", 1),
        show_infr_attr_card=getattr(state, "_show_infr_attr_card", 1),
        show_social_scraping=getattr(state, "_show_social_scraping", False),
        show_upload_data=getattr(state, "_show_upload_data", False),
        rag_corpus_path=getattr(state, "_rag_corpus_path", ""),
    )

    return conv_html, create_risk_display(risk_score, risk_breakdown), "", analysis, privacy_settings_text


# def _build_analysis(pii, rag, risk, use_rag, eps, pr, pnr):
#     """Build analysis markdown."""
#     parts = ["### Detailed Privacy Analysis\n\n**Detected PIIs:**\n"]
#     if pii:
#         for p in pii:
#             parts.append(f"- {p.category.value}: `{p.text}` (confidence: {p.confidence:.0%})\n")
#     else:
#         parts.append("- None detected\n")
#
#     parts.append("\n**RAG Linkages:**\n")
#     if rag:
#         for r in rag:
#             parts.append(f"- `{r.text}` (similarity: {r.top_similarity:.2f})\n")
#     else:
#         parts.append("- None detected\n")
#
#     parts.append(f"\n**Privacy Settings:**\n")
#     parts.append(f"- RAG: {'Enabled' if use_rag else 'Disabled'}\n")
#     parts.append(f"- ε (epsilon): {eps if eps != float('inf') else '∞ (no DP)'}\n")
#     parts.append(f"- Risk Score: {risk}/100\n")
#
#     return "".join(parts)

def _build_inference_plot(probs_rag, probs_no_rag):
    """Build HTML visualization of inference privacy metrics (lift alerts only, no bars)."""
    metrics_data = []

    for attr in SENSITIVE_ATTRIBUTES:
        if attr not in probs_rag:
            continue

        dist_rag = probs_rag[attr]
        if not dist_rag:
            continue

        top_val_rag = max(dist_rag, key=dist_rag.get)
        p_rag = dist_rag[top_val_rag]
        p_no_rag = (probs_no_rag or {}).get(attr, {}).get(top_val_rag, 0.0)

        # Compute lift against population prior
        p_pop = POPULATION_PRIORS.get(attr, {}).get(top_val_rag, 0.0)
        if p_pop > 0:
            z = InferentialPrivacyMetrics.compute_z(p_rag, p_pop)
            lift = InferentialPrivacyMetrics.compute_lift(z)
        else:
            lift = 0.0

        metrics_data.append({
            'attr': attr,
            'top_val': top_val_rag,
            'p_rag': p_rag,
            'p_pop': p_pop,
            'lift': lift,
        })

    if not metrics_data:
        return ""

    rows = ['<div style="margin-top:6px;">']
    for data in metrics_data:
        attr     = data['attr']
        top_val  = data['top_val']
        p_rag    = data['p_rag']
        p_pop    = data['p_pop']
        lift     = data['lift'] if not math.isinf(data['lift']) else -1

        rows.append(
            f'<div style="display:flex;align-items:center;justify-content:space-between;'
            f'margin-bottom:5px;font-size:0.82em;">'
            f'<span style="color:#333;font-weight:600;">{attr}</span>'
            f'<span style="color:#555;">{top_val} ({p_rag:.0%})</span>'
        )
        if lift > 0:
            rows.append(
                f'<span style="color:#b71c1c;font-weight:700;margin-left:6px;">'
                f'+{lift:.0%} above base rate</span>'
            )
        else:
            rows.append('<span></span>')
        rows.append('</div>')

    rows.append('</div>')
    return "\n".join(rows)

# def _build_analysis(pii, rag, risk, use_rag, eps, pr, pnr):
#     """Build analysis markdown with improved visualizations."""
#     parts = ["### Detailed Privacy Analysis\n\n**Detected PIIs:**\n"]
#     if pii:
#         for p in pii:
#             parts.append(f"- {p.category.value}: `{p.text}` (confidence: {p.confidence:.0%})\n")
#     else:
#         parts.append("- None detected\n")
#
#     parts.append("\n**RAG Linkages:**\n")
#     if rag:
#         # Group by document and create a more narrative format
#         doc_groups = {}
#         for r in rag:
#             # Assume r has a doc_id or source attribute, fallback to grouping by similarity range
#             doc_key = getattr(r, 'doc_id', f"doc_{hash(r.text) % 100}")
#             if doc_key not in doc_groups:
#                 doc_groups[doc_key] = []
#             doc_groups[doc_key].append(r)
#
#         for doc_id, matches in doc_groups.items():
#             avg_sim = sum(m.top_similarity for m in matches) / len(matches)
#             match_count = len(matches)
#             relevance = "High" if avg_sim > 0.8 else "Medium" if avg_sim > 0.6 else "Moderate"
#             parts.append(f"- **Document {doc_id}**: {match_count} match{'es' if match_count > 1 else ''} "
#                          f"(relevance: {relevance}, avg similarity: {avg_sim:.2f})\n")
#     else:
#         parts.append("- None detected\n")
#
#     parts.append(f"\n**Privacy Settings:**\n")
#     parts.append(f"- RAG: {'Enabled' if use_rag else 'Disabled'}\n")
#     parts.append(f"- ε (epsilon): {eps if eps != float('inf') else '∞ (no DP)'}\n")
#
#     # Add privacy metrics plot if we have inference data
#     if use_rag and pr and pnr:
#         parts.append("\n**Inference Privacy Metrics:**\n")
#         parts.append(_build_inference_plot(pr, pnr))
#
#     return "".join(parts)

# def _build_analysis(pii, rag, risk, use_rag, eps, pr, pnr):
#     """Build analysis markdown with improved visualizations."""
#     parts = ["### Detailed Privacy Analysis\n\n**Detected PIIs:**\n"]
#
#     # Count PIIs in prompts vs responses
#     user_pii_count = sum(1 for p in pii if hasattr(p, 'start'))  # Simplified check
#     # Note: This assumes you have a way to distinguish user vs response PIIs
#     # If pii is combined, you may need to track this separately
#
#     if pii:
#         # Group by category
#         category_counts = {}
#         for p in pii:
#             cat_name = p.category.value
#             category_counts[cat_name] = category_counts.get(cat_name, 0) + 1
#
#         parts.append(f"- Total detected: {len(pii)}\n")
#         for cat, count in category_counts.items():
#             parts.append(f"  - {cat}: {count}\n")
#     else:
#         parts.append("- None detected\n")
#
#     parts.append(f"\n**Privacy Settings:**\n")
#
#     # Add DP explanation
#     if eps != float('inf'):
#         parts.append(f"- **Differential Privacy**: ENABLED (ε = {eps:.1f})\n")
#         parts.append(f"  *Lower epsilon values provide stronger privacy protection "
#                      f"by adding noise to data retrieval, making it harder to infer sensitive "
#                      f"information about individuals.*\n")
#     else:
#         parts.append(f"- **Differential Privacy**: DISABLED (ε = ∞)\n")
#         parts.append(f"  *No privacy protection applied. Consider enabling DP (lower epsilon) "
#                      f"to add noise-based protection.*\n")
#
#     # Add privacy metrics plot if we have inference data
#     if use_rag and pr and pnr:
#         parts.append("\n**Inference Privacy Metrics:**\n")
#         parts.append(_build_inference_plot(pr, pnr))
#
#     return "".join(parts) # last worked

# def _build_privacy_settings_html(eps):
#     """Build Privacy Settings HTML for left panel with fixed tooltip positioning."""
#
#     parts = ['<div style="padding:12px;">']
#     parts.append('<h4 style="margin-top:0;">Privacy Settings</h4>')
#
#     # Tooltip content - formatted as HTML with !important for font sizes
#     dp_tooltip = (
#         '<div style="font-weight:700;font-size:1.1em!important;margin-bottom:5px;color:#111;">What is Differential Privacy?</div>'
#         '<div style="height:1px;background:#ddd;margin:8px 0;"></div>'
#         '<div style="color:#333;line-height:1.4!important;margin-bottom:6px;font-size:0.95em!important;">'
#         'Differential Privacy (DP) is a mathematical framework that adds calibrated noise '
#         'to data retrieval, ensuring individual records cannot be distinguished. '
#         'This protects sensitive information from being inferred when external data sources are used.'
#         '</div>'
#         '<div style="height:1px;background:#ddd;margin:8px 0;"></div>'
#         '<div style="font-weight:600;color:#2c3e50;margin-bottom:4px;font-size:1.0em!important;">Understanding Epsilon (ε):</div>'
#         '<div style="color:#333;line-height:1.4!important;font-size:0.95em!important;">'
#         '<strong>Lower values (0.1-1.0):</strong> Stronger privacy, more noise added.<br>'
#         '<strong>Higher values (5.0-10.0):</strong> Less privacy, better data utility.<br>'
#         '<strong>Epsilon = ∞:</strong> No privacy protection applied.'
#         '</div>'
#     )
#
#     # Unique IDs for this tooltip instance
#     tooltip_id = f"dp_tooltip_{abs(hash(str(eps)))}"
#     anchor_id = f"dp_anchor_{abs(hash(str(eps)))}"
#
#     # JavaScript to position tooltip above anchor
#     # JavaScript to position tooltip next to the current mouse cursor
#     position_script = f"""
#             <script>
#             (function() {{
#                 var _mouseX = 0, _mouseY = 0;
#                 document.addEventListener('mousemove', function(e) {{
#                     _mouseX = e.clientX;
#                     _mouseY = e.clientY;
#                 }}, true);
#
#                 window.position_{tooltip_id} = function() {{
#                     var tooltip = document.getElementById('{tooltip_id}');
#                     if (!tooltip) return;
#
#                     var tooltipWidth  = 380;
#                     var tooltipHeight = tooltip.offsetHeight || 220;
#                     var offset        = 14;   // gap between cursor tip and tooltip edge
#                     var viewportW     = window.innerWidth;
#                     var viewportH     = window.innerHeight;
#                     var scrollLeft    = window.pageXOffset || document.documentElement.scrollLeft;
#                     var scrollTop     = window.pageYOffset || document.documentElement.scrollTop;
#
#                     // Prefer: above-right of cursor
#                     var left = _mouseX + offset + scrollLeft;
#                     var top  = _mouseY - tooltipHeight - offset + scrollTop;
#
#                     // Flip to below if not enough room above
#                     if (_mouseY < tooltipHeight + offset + 10) {{
#                         top = _mouseY + offset + scrollTop;
#                     }}
#
#                     // Flip to left of cursor if not enough room on the right
#                     if (_mouseX + offset + tooltipWidth > viewportW - 10) {{
#                         left = _mouseX - tooltipWidth - offset + scrollLeft;
#                     }}
#
#                     // Final clamp to viewport edges
#                     left = Math.max(scrollLeft + 8, Math.min(left, scrollLeft + viewportW - tooltipWidth - 8));
#                     top  = Math.max(scrollTop  + 8, top);
#
#                     tooltip.style.left = left + 'px';
#                     tooltip.style.top  = top  + 'px';
#                 }};
#             }})();
#             </script>
#             """
#
#     # Add DP status with fixed-position tooltip
#     if eps != float('inf'):
#         parts.append(
#             f'{position_script}'
#             '<p style="margin-bottom:8px;">'
#             f'<span id="{anchor_id}" style="position:relative;display:inline-block;cursor:help;border-bottom:1.5px dotted #1976D2;font-weight:600;color:#1976D2;" '
#             f'onmouseover="document.getElementById(\'{tooltip_id}\').style.display=\'block\';position_{tooltip_id}();" '
#             f'onmouseout="document.getElementById(\'{tooltip_id}\').style.display=\'none\'">'
#             'Differential Privacy'
#             '</span>: '
#             '<span style="color:#2E7D32;font-weight:600;">ENABLED</span>'
#             '</p>'
#             f'<p style="margin-top:4px;font-size:0.95em;">'
#             f'<strong>Differential Privacy guarantee level (epsilon parameter) =</strong> <span style="color:#1976D2;font-weight:700;font-size:1.1em;">{eps:.1f}</span>'
#             f'</p>'
#             f'<div id="{tooltip_id}" style="display:none;position:fixed;z-index:999999;'
#             'background:#fff;color:#111;'
#             'border:2px solid #bbb;border-radius:10px;padding:10px 12px;width:380px;'
#             'box-shadow:0 8px 20px rgba(0,0,0,0.15);font-size:0.7em!important;line-height:1.25!important;'
#             'text-align:left;font-weight:normal;pointer-events:none;">'
#             f'{dp_tooltip}'
#             '</div>'
#         )
#     else:
#         parts.append(
#             f'{position_script}'
#             '<p style="margin-bottom:8px;">'
#             f'<span id="{anchor_id}" style="position:relative;display:inline-block;cursor:help;border-bottom:1.5px dotted #1976D2;font-weight:600;color:#1976D2;" '
#             f'onmouseover="document.getElementById(\'{tooltip_id}\').style.display=\'block\';position_{tooltip_id}();" '
#             f'onmouseout="document.getElementById(\'{tooltip_id}\').style.display=\'none\'">'
#             'Differential Privacy'
#             '</span>: '
#             '<span style="color:#D32F2F;font-weight:600;">DISABLED</span>'
#             '</p>'
#             f'<p style="margin-top:4px;font-size:0.95em;">'
#             f'<strong>Differential Privacy Guarantee =</strong> <span style="color:#D32F2F;font-weight:700;font-size:1.1em;">∞</span>'
#             f'</p>'
#             f'<p style="font-size:0.85em;color:#666;margin-top:8px;font-style:italic;">'
#             f'Consider enabling DP for privacy protection.'
#             f'</p>'
#             f'<div id="{tooltip_id}" style="display:none;position:fixed;z-index:999999;'
#             'background:#fff;color:#111;'
#             'border:2px solid #bbb;border-radius:10px;padding:10px 12px;width:380px;'
#             'box-shadow:0 8px 20px rgba(0,0,0,0.15);font-size:0.7em!important;line-height:1.25!important;'
#             'text-align:left;font-weight:normal;pointer-events:none;">'
#             f'{dp_tooltip}'
#             '</div>'
#         )
#
#     parts.append("</div>")
#     return "".join(parts)

def _build_privacy_settings_html(eps):
    """Build Privacy Settings HTML for left panel with a CSS-native expandable DP explainer."""

    if eps != float('inf'):
        status_html = '<span style="color:#2E7D32;font-weight:600;">ENABLED</span>'
        # Translate epsilon to a human-readable protection level
        if eps <= 30.0:
            level_label = "High"
            level_color = "#2E7D32"
            level_desc  = "Strong protection — the AI's view of your data is significantly randomised."
        elif eps <= 60.0:
            level_label = "Medium"
            level_color = "#F57C00"
            level_desc  = "Moderate protection — some noise is added to limit profiling accuracy."
        else:
            level_label = "Low"
            level_color = "#D32F2F"
            level_desc  = "Weak protection — only a small amount of noise is applied."
        eps_line = (
            f'<p style="margin-top:4px;margin-bottom:0;font-size:0.95em;">'
            f'<strong>Protection level:</strong> '
            f'<span style="color:{level_color};font-weight:700;font-size:1.1em;">{level_label}</span>'
            f'</p>'
            f'<p style="font-size:0.85em;color:#000000;margin-top:3px;font-style:italic;">'
            # f'{level_desc}</p>'
        )
    else:
        status_html = '<span style="color:#D32F2F;font-weight:600;">DISABLED</span>'
        eps_line    = (
            f'<p style="margin-top:4px;margin-bottom:0;font-size:0.95em;">'
            f'<strong>Protection level:</strong> '
            f'<span style="color:#D32F2F;font-weight:700;font-size:1.1em;">None</span>'
            f'</p>'
            # f'<p style="font-size:0.85em;color:#000000;margin-top:4px;font-style:italic;">'
            # f'Consider enabling privacy protection to reduce profiling risk.</p>'
        )

    return (
        f'<div style="border:1px solid #aaaaaa; color:#000000;">'
        f'<p style="margin-bottom:4px; color:#000000;"><strong>Differential Privacy</strong>: {status_html}</p>'
        f'{eps_line}'

        # ── Pure-CSS expandable explainer ──────────────────────────
        f'<details style="margin-top:10px;">'
        f'<summary style="cursor:pointer;font-size:0.9em;font-weight:600;color:#1976D2;'
        f'list-style:none;display:inline-flex;align-items:center;gap:4px;'
        f'user-select:none;">'
        f'<span class="dp-arrow" style="font-size:0.8em;">▼</span> What is this?'
        f'</summary>'

        f'<div style="margin-top:8px;padding:10px 12px;'
        f'background:#f0f4f8;border-left:3px solid #1976D2;border-radius:0 6px 6px 0;'
        f'font-size:0.88em;color:#000000;">'

        f'<div style="margin-bottom:10px;color:#000000;">'
        f'When privacy protection is enabled, this tool automatically modifies your messages '
        f'to make it harder for the AI to identify you, <em>before</em> any information '
        f'leaves your device:'
        f'</div>'

        f'<div style="display:flex;align-items:flex-start;gap:8px;margin-bottom:10px;">'
        f'<span style="font-size:1.3em;flex-shrink:0;">✍️</span>'
        f'<div>'
        f'<strong>Some words in your message.</strong> '
        f'Some words in what you type are automatically swapped for similar-sounding '
        f'alternatives before the AI reads them. '
        f'A 🔒 badge on your message shows when this happened and which words changed.'
        f'</div>'
        f'</div>'

        f'<div style="height:1px;background:#aaaaaa;margin:8px 0;"></div>'

        f'<div style="font-weight:600;color:#000000;margin-bottom:6px;">'
        f'About the protection level:</div>'
        f'<div style="margin-bottom:3px;color:#000000;">🟢 <strong>High:</strong> '
        f'Strong privacy — more noise applied, AI responses may be slightly less precise.</div>'
        f'<div style="margin-bottom:3px;color:#000000;">🟡 <strong>Medium:</strong> '
        f'Balanced — moderate noise, reasonable trade-off between privacy and accuracy.</div>'
        f'<div style="margin-bottom:3px;color:#000000;">🔴 <strong>Low:</strong> '
        f'Weak privacy — minimal noise, AI responses are more accurate but you are easier to profile.</div>'
        f'<div style="margin-bottom:3px;color:#D32F2F;"><strong>Off:</strong> '
        f'No protection — the AI sees your exact message and data.</div>'

        f'</div>'  # end inner content div
        f'</details>'
        f'</div>'  # end outer padding div
    )

def _pii_severity(category):
    """Return (label, color) severity info for a PIICategory."""
    if category == PIICategory.SENSITIVE:
        return "High", "#b71c1c"
    elif category == PIICategory.IDENTITY:
        return "High", "#b71c1c"
    elif category == PIICategory.CONTACT:
        return "Medium", "#e65100"
    else:  # LOCATION
        return "Low", "#1976D2"


def _load_avatar_image_b64(gender, age):
    """Load a profile photo from the avatars/ directory and return a base64 data-URI.

    Expected filenames in <project_root>/avatars/:
        male_teen.jpg    female_teen.jpg     (Age bin: 0-17)
        male_young.jpg   female_young.jpg    (Age bin: 18-29)
        male_adult.jpg   female_adult.jpg    (Age bin: 30-44)
        male_middle.jpg  female_middle.jpg   (Age bin: 45-59)
        male_senior.jpg  female_senior.jpg   (Age bin: 60+)
        default.jpg                          (fallback when gender unknown)
    Also accepts .png variants of every name above.

    Fallback chain (tried in order):
        1. avatars/{gender}_{age_label}.jpg/.png
        2. avatars/{gender}_adult.jpg/.png     ← same gender, adult age
        3. avatars/default.jpg/.png            ← universal fallback
        4. Returns ("", False) → caller shows an initials circle

    Returns: (data_uri_string, found_bool)
    """
    import base64

    gender_prefix = {"Male": "male", "Female": "female"}.get(gender, None)
    age_suffix = {
        "0-17":  "teen",
        "18-29": "young",
        "30-44": "adult",
        "45-59": "middle",
        "60+":   "senior",
    }.get(age, "adult")

    root = os.path.dirname(os.path.abspath(__file__))
    avatars_dir = os.path.join(root, "avatars")

    candidates = []
    if gender_prefix:
        candidates.append(os.path.join(avatars_dir, f"{gender_prefix}_{age_suffix}.jpg"))
        candidates.append(os.path.join(avatars_dir, f"{gender_prefix}_{age_suffix}.png"))
        candidates.append(os.path.join(avatars_dir, f"{gender_prefix}_adult.jpg"))
        candidates.append(os.path.join(avatars_dir, f"{gender_prefix}_adult.png"))
    candidates.append(os.path.join(avatars_dir, "default.jpg"))
    candidates.append(os.path.join(avatars_dir, "default.png"))

    for path in candidates:
        if os.path.isfile(path):
            try:
                ext  = os.path.splitext(path)[1].lower()
                mime = "image/png" if ext == ".png" else "image/jpeg"
                with open(path, "rb") as f:
                    b64 = base64.b64encode(f.read()).decode("ascii")
                return f"data:{mime};base64,{b64}", True
            except Exception:
                continue
    return "", False


def _build_inferred_avatar(probs_rag, inference_warning_html="", is_warning=False):
    """Build the 'How the AI sees you' card.

    Layout (top → bottom inside the card):
      • Header: title left, confidence badge right
      • Subtitle
      • Photo: 130×130 px circle, centred
      • Attribute rows: one per inferred attribute, full card width.
        Each row embeds the inner content of the matching <li> verbatim from
        inference_warning_html, preserving:
          – the predicted value
          – the Why? tt-anchor with its evidence tt-popup spans (hover works)
          – the probability (p = X%)
          – the red lift message if above threshold
        Nothing is added or duplicated by this function.

    Card colour:
      is_warning=True  → yellow (#FFFDE7) + amber border
      is_warning=False → blue  (#E3F2FD) + blue  border
    """
    # Always fill any missing attributes with a uniform distribution so the
    # panel is always populated, even when the inference LLM call fails.
    probs_rag = _fill_missing_with_uniform(probs_rag or {})

    # ── Minimum-confidence gate ───────────────────────────────────────────
    def _top(dist, min_p=0.15):
        if not dist:
            return None, 0.0
        v = max(dist, key=dist.get)
        p = dist[v]
        return (v, p) if p >= min_p else (None, p)

    gender,  gender_p  = _top(probs_rag.get("Gender",        {}), 0.20)
    age,     age_p     = _top(probs_rag.get("Age bin",        {}), 0.15)
    finance, finance_p = _top(probs_rag.get("Finance Status", {}), 0.15)

    if max(gender_p, age_p, finance_p) < 0.15:
        return ""

    # ── Card colour scheme ────────────────────────────────────────────────
    if is_warning:
        card_bg     = "#FFFDE7"
        card_border = "#F9A825"
        title_color = "#E65100"
        title_icon  = "⚠️"
        # subtitle    = "External data raised the risk of accurate profiling."
        subtitle    =  "Based on your conversation so far."
    else:
        card_bg     = "#E3F2FD"
        card_border = "#1976D2"
        title_color = "#0D47A1"
        title_icon  = "🤖"
        subtitle    = "Based on your conversation so far."

    # ── Profile photo (or initials fallback) ──────────────────────────────
    img_uri, img_found = _load_avatar_image_b64(gender, age)

    if img_found:
        photo_html = (
            f'<img src="{img_uri}" alt="AI-inferred profile photo" '
            f'style="width:130px;height:130px;object-fit:cover;border-radius:50%;'
            f'border:3px solid {card_border};display:block;margin:0 auto 12px;">'
        )
    else:
        initials = "F" if gender == "Female" else ("M" if gender == "Male" else "?")
        photo_html = (
            f'<div style="width:130px;height:130px;border-radius:50%;'
            f'background:{card_border};border:3px solid {card_border};'
            f'display:flex;align-items:center;justify-content:center;'
            f'margin:0 auto 12px;color:white;font-size:2em;font-weight:700;">'
            f'{initials}</div>'
        )

    # ── Per-attribute content extractor ──────────────────────────────────
    # The <li> structure produced by build_inference_warning (after Fix 1) is:
    #
    #   <li style="..."><b>ATTR</b>: VALUE
    #     <span class="tt-anchor">Why?
    #       <span class="tt-popup tt-above">EVIDENCE</span>
    #       <span class="tt-popup tt-below">EVIDENCE</span>
    #     </span>
    #     <span style="color:#555;">(p = X%)</span>
    #     [<span style="color:#b71c1c;">lift</span>]
    #   </li>
    #
    # We grab everything between "<b>ATTR</b>: " and "</li>" and embed it
    # verbatim — value, Why? anchor with evidence popups, probability, lift
    # — all intact, with no duplication.

    def _extract_li_content(html, attr):
        """Return inner content of the <li> for attr (after the colon), or ''."""
        marker = f'<b>{attr}</b>: '
        idx = html.find(marker)
        if idx == -1:
            return ""
        content_start = idx + len(marker)
        li_close = html.find('</li>', content_start)
        if li_close == -1:
            return ""
        return html[content_start:li_close].strip()

    # ── Build attribute rows ──────────────────────────────────────────────
    ATTR_ORDER = [
        "Gender", "Age bin", "Locale",
        "Marital Status", "Finance Status", "Education",
    ]

    rows = []
    for attr in ATTR_ORDER:
        dist = probs_rag.get(attr, {})
        if not dist:
            continue
        top_val = max(dist, key=dist.get)
        top_p   = dist[top_val]

        li_content = _extract_li_content(inference_warning_html, attr)

        if li_content:
            # Embed verbatim — Why? tooltip, probability, and lift all included
            row_content = (
                f'<span style="font-weight:600;color:#333;font-size:0.9em;">'
                f'{attr}:</span> {li_content}'
            )
        else:
            # Fallback: plain text when warning_html is absent or attr missing
            row_content = (
                f'<span style="font-weight:600;color:#333;font-size:0.9em;">'
                f'{attr}:</span> '
                f'<span style="color:#111;">{top_val}</span> '
                f'<span style="color:#666;font-size:0.83em;">'
                f'(p&nbsp;=&nbsp;{top_p:.0%})</span>'
            )

        rows.append(
            f'<div style="margin-bottom:8px;line-height:1.5;">'
            f'{row_content}'
            f'</div>'
        )

    attrs_html = "".join(rows) if rows else (
        '<p style="color:#999;font-size:0.82em;font-style:italic;">'
        'No attributes inferred yet.</p>'
    )

    # ── Confidence badge ──────────────────────────────────────────────────
    conf     = max(gender_p, age_p, finance_p)
    conf_pct = int(conf * 100)

    # ── Assemble card ─────────────────────────────────────────────────────
    return (
        f'<div style="background:{card_bg};border:2px solid {card_border};'
        f'border-radius:10px;padding:12px 14px 10px;margin-bottom:8px;">'

        # Header
        f'<div style="display:flex;align-items:center;'
        f'justify-content:space-between;margin-bottom:5px;">'
        f'<span style="font-size:0.82em;font-weight:700;color:{title_color};">'
        f'{title_icon} How the AI sees you'
        f'</span>'
        f'<span style="font-size:0.68em;color:{title_color};opacity:0.8;'
        f'border:1px solid {card_border};border-radius:10px;padding:1px 7px;">'
        f'confidence: {conf_pct}%'
        f'</span>'
        f'</div>'

        # Subtitle
        f'<div style="font-size:0.72em;color:#666;font-style:italic;margin-bottom:10px;">'
        f'{subtitle}'
        f'</div>'

        # Photo centred
        f'{photo_html}'

        # Attribute rows — full card width
        f'<div style="font-size:0.88em;">'
        f'{attrs_html}'
        f'</div>'

        # Footer disclaimer
        f'<div style="font-size:0.62em;color:#999;margin-top:8px;font-style:italic;">'
        f'Photo is illustrative only — selected to match inferred profile. '
        f'These are AI predictions and may be incorrect.'
        f'</div>'

        f'</div>'
    )

def _build_analysis(pii, rag, risk, use_rag, eps, pr, pnr, inference_warning_html="", is_warning=False, include_avatar=True, show_pii=True):
    """Build analysis HTML for the right panel.

    Layout (top → bottom):
      1. 'How the AI sees you' card  (photo + inferred attrs + Why? tooltips)
         — omitted when include_avatar=False (card rendered separately)
      2. Divider
      3. Detected PIIs (tags with severity)
    """
    parts = ['<div style="padding-left:8px;padding-right:8px;font-size:0.88em;">']

    # ── Section 1: Combined avatar + inferable attributes ────────────────
    if include_avatar:
        avatar_html = _build_inferred_avatar(pr or {}, inference_warning_html, is_warning)
        if avatar_html:
            parts.append(avatar_html)
            parts.append('<div style="height:1px;background:#e0e0e0;margin:8px 0;"></div>')

    # ── Section 2: Detected PIIs (last, as requested) ────────────────────
    if show_pii:
        parts.append('<p style="margin-bottom:5px;font-weight:700;color:#2c3e50;">Detected Personally Identifiable Information (PII) in user\'s messages:</p>')

        if pii:
            # De-duplicate by (text, fine_type)
            seen = set()
            unique_pii = []
            for p in pii:
                key = (p.text.strip().lower(), p.fine_type.lower())
                if key not in seen:
                    seen.add(key)
                    unique_pii.append(p)

            parts.append('<div style="display:flex;flex-wrap:wrap;gap:5px;margin-bottom:8px;">')
            for p in unique_pii:
                col = PII_COLORS.get(p.category, {"bg": "#f5f5f5", "border": "#999"})
                sev_label, sev_color = _pii_severity(p.category)
                cat_label = p.category.value.capitalize()
                fine = p.fine_type.title() if p.fine_type else cat_label
                text_disp = p.text[:28] + ("…" if len(p.text) > 28 else "")
                parts.append(
                    f'<span title="{fine} · {sev_label} severity · confidence {p.confidence:.0%}" '
                    f'style="display:inline-flex;align-items:center;gap:3px;'
                    f'background:{col["bg"]};border:1.5px solid {col["border"]};'
                    f'border-radius:4px;padding:2px 6px;cursor:default;">'
                    f'<span style="font-size:0.78em;font-weight:600;color:#333;">{text_disp}</span>'
                    f'<span style="font-size:0.68em;color:#666;margin-left:1px;">[{fine}]</span>'
                    f'<span style="font-size:0.65em;font-weight:700;color:{sev_color};'
                    f'margin-left:2px;border-left:1px solid {col["border"]};padding-left:3px;">{sev_label}</span>'
                    f'</span>'
                )
            parts.append('</div>')

            # Category legend
            cat_counts = {}
            for p in unique_pii:
                cat_counts[p.category.value] = cat_counts.get(p.category.value, 0) + 1
            legend_parts = [f'<span style="color:#555;">{v}× {k}</span>' for k, v in cat_counts.items()]
            parts.append(f'<div style="font-size:0.78em;color:#888;margin-bottom:8px;">'
                         f'Total: {len(unique_pii)} — {", ".join(legend_parts)}</div>')
        else:
            parts.append('<div style="color:#888;font-size:0.85em;margin-bottom:8px;">None detected</div>')

    parts.append("</div>")
    return "".join(parts)


# def build_inference_warning(probs_rag, probs_no_rag, use_rag):


# ============================================================
# SECTION 14.5 – ROBUST EVIDENCE SOURCE MATCHING
# ============================================================

def tokenize_for_matching(text):
    """
    Tokenize text for matching, removing punctuation and common words.

    Returns:
        set: Set of significant tokens (lowercased, >3 chars)
    """
    import re
    # Remove punctuation and split
    tokens = re.findall(r'\b\w+\b', text.lower())
    # Keep only significant tokens (>3 chars, not common stop words)
    significant_tokens = {t for t in tokens if len(t) > 3 and t not in _DP_STOPWORDS}
    return significant_tokens


def calculate_token_overlap(quote_tokens, doc_tokens, bidirectional=False):
    """Calculate token overlap ratio between two token sets.

    Parameters
    ----------
    quote_tokens, doc_tokens : set
    bidirectional : bool
        If True, return max(forward, reverse) overlap so that a short quote
        with a high fraction of user-input words is also detected, not only
        a quote whose words are mostly covered by the user input.
        Use True when matching against user input; False (default) for docs.

    Returns
    -------
    float: overlap ratio in [0, 1]
    """
    if not quote_tokens or not doc_tokens:
        return 0.0
    intersection = len(quote_tokens & doc_tokens)
    forward  = intersection / len(quote_tokens)
    if not bidirectional:
        return forward
    reverse  = intersection / len(doc_tokens)  # doc_tokens is user_tokens here
    return max(forward, reverse)


def extract_doc_source_label(doc):
    """
    Extract human-readable source label from document metadata.

    Args:
        doc: Document object with metadata

    Returns:
        str: Source label (e.g., "From uploaded data", "From Twitter (Internet)")
    """
    if not hasattr(doc, 'metadata') or not isinstance(doc.metadata, dict):
        return "From system data"

    source_type = doc.metadata.get('source', '')
    platform = doc.metadata.get('platform', '').strip().lower()

    if source_type == 'uploaded_csv':
        return "From uploaded data"
    elif source_type.startswith('social_media_'):
        platform_label = platform.capitalize() if platform else "Social Media"
        return f"From {platform_label} (Internet)" if "web" not in platform_label.lower() else f"From  the {platform_label} (Internet)"
    elif source_type == 'fallback_tweet_data':
        return "From Twitter (Internet)"
    else:
        # Documents from the default background corpus (e.g. PANORAMA synthetic profiles)
        return DEFAULT_CORPUS_SOURCE_LABEL


def _doc_source_url(doc):
    """Extract the source URL from a document's metadata, or '' if absent."""
    if not hasattr(doc, "metadata") or not isinstance(doc.metadata, dict):
        return ""
    return (
        doc.metadata.get("source_url") or
        doc.metadata.get("tweet_url") or
        doc.metadata.get("url") or
        ""
    )


def find_evidence_source(quote, user_input, retrieved_docs, min_overlap_ratio=0.4,
                         perturbed_input=None, prior_turns=None):
    """
    Find the source of an evidence quote using robust multi-strategy matching.

    Matching priority:
      Strategy 1 – Exact substring: user turns first, then docs  (high precision)
      Strategy 2 – Token overlap:   docs first, then user turns  (avoids topic-word FPs)
      Strategy 3 – Phrase windows:  docs first, then user turns

    Evidence quotes are LLM-generated paraphrases.  Because the LLM is asked
    to explain inferences drawn from retrieved documents, most evidence quotes
    describe external-data content even when they share topic words with the
    user's message.  Checking docs before user turns in the fuzzy strategies
    prevents those shared topic words from causing false "From your input" labels.

    When no source can be matched the quote is considered LLM internal reasoning
    (rationale) rather than a verbatim excerpt from any document.

    Returns:
        tuple: (is_from_user, is_from_docs, doc_source, match_confidence,
                matched_turn_dp_active, matched_turn_original_text, doc_url)
        doc_url – URL of the matched document, or "" if not from a doc / unknown.
    """
    import re as _re

    quote_lower = str(quote).lower().strip()
    user_input_lower = user_input.lower() if user_input else ""
    perturbed_lower  = perturbed_input.lower() if perturbed_input else ""

    # Build the full list of user turns to search
    current_dp_active = bool(perturbed_lower)
    turns_to_search = [(user_input_lower, perturbed_lower,
                        current_dp_active, user_input or "")]
    if prior_turns:
        for pt in prior_turns:
            orig_l = (pt.get("original") or "").lower()
            pert_l = (pt.get("perturbed") or "").lower()
            dp_on  = bool(pt.get("dp_active", False))
            orig_d = pt.get("original") or ""
            turns_to_search.append((orig_l, pert_l, dp_on, orig_d))

    # ──────────────────────────────────────────────────────────────────────
    # Strategy 1: Exact substring — user turns first, then docs
    # Rationale: if the evidence is a verbatim copy of text the user typed,
    # it genuinely comes from the user.
    # ──────────────────────────────────────────────────────────────────────
    for orig_l, pert_l, dp_on, orig_disp in turns_to_search:
        if quote_lower in orig_l:
            return True, False, "From your input", "exact", dp_on, orig_disp, ""
        if len(quote_lower) > 50 and quote_lower[:50] in orig_l:
            return True, False, "From your input", "exact", dp_on, orig_disp, ""
        if pert_l and quote_lower in pert_l:
            return True, False, "From your input", "exact", dp_on, orig_disp, ""
        if pert_l and len(quote_lower) > 50 and quote_lower[:50] in pert_l:
            return True, False, "From your input", "exact", dp_on, orig_disp, ""

    if retrieved_docs:
        for doc in retrieved_docs:
            doc_text  = doc.page_content if hasattr(doc, "page_content") else str(doc)
            doc_lower = doc_text.lower()
            if quote_lower in doc_lower or (
                    len(quote_lower) > 50 and quote_lower[:50] in doc_lower):
                doc_source = extract_doc_source_label(doc)
                return False, True, doc_source, "exact", False, "", _doc_source_url(doc)

    # ──────────────────────────────────────────────────────────────────────
    # Strategy 2: Token overlap — docs first, then user turns
    # Rationale: LLM-generated evidence paraphrases retrieved docs.  Both
    # the evidence and the user message discuss the same topic, so shared
    # topic tokens do NOT indicate the evidence came from the user.
    # We check docs first; user turns are only a fallback with a high bar.
    # ──────────────────────────────────────────────────────────────────────
    quote_tokens = tokenize_for_matching(quote)
    DOC_OVERLAP_THRESHOLD  = 0.35   # forward overlap: fraction of quote tokens in doc
    USER_OVERLAP_THRESHOLD = 0.65   # bidirectional: raised from 0.25 to avoid FPs

    if len(quote_tokens) >= 3:
        # 2a: Check retrieved docs first
        if retrieved_docs:
            best_overlap, best_doc = 0.0, None
            for doc in retrieved_docs:
                doc_text   = doc.page_content if hasattr(doc, "page_content") else str(doc)
                doc_tokens = tokenize_for_matching(doc_text)
                ov = calculate_token_overlap(quote_tokens, doc_tokens)
                if ov > best_overlap:
                    best_overlap, best_doc = ov, doc
            if best_overlap >= DOC_OVERLAP_THRESHOLD and best_doc:
                doc_source = extract_doc_source_label(best_doc)
                logger.info(
                    f"  🔍 Token overlap match with doc ({best_overlap:.0%}): '{quote[:50]}...'")
                return False, True, doc_source, "token_overlap", False, "", _doc_source_url(best_doc)

        # 2b: Fall back to user turns only with a high threshold
        for orig_l, pert_l, dp_on, orig_disp in turns_to_search:
            turn_tokens = tokenize_for_matching(orig_l)
            if calculate_token_overlap(quote_tokens, turn_tokens,
                                       bidirectional=True) >= USER_OVERLAP_THRESHOLD:
                logger.info(
                    f"  🔍 Token overlap match with user turn: '{quote[:50]}...'")
                return True, False, "From your input", "token_overlap", dp_on, orig_disp, ""
            if pert_l:
                pert_tokens = tokenize_for_matching(pert_l)
                if calculate_token_overlap(quote_tokens, pert_tokens,
                                           bidirectional=True) >= USER_OVERLAP_THRESHOLD:
                    logger.info(
                        f"  🔍 Token overlap match with perturbed user turn: '{quote[:50]}...'")
                    return True, False, "From your input", "token_overlap", dp_on, orig_disp, ""

    # ──────────────────────────────────────────────────────────────────────
    # Strategy 3: Short-phrase windows — docs first, then user turns
    # 3a: 3-word windows against docs
    # 3b: 3-word windows against user turns (stricter guard: phrase must NOT
    #     appear in any retrieved doc first)
    # ──────────────────────────────────────────────────────────────────────
    quote_words = _re.findall(r'\b\w+\b', quote_lower)

    if len(quote_words) >= 3 and retrieved_docs:
        doc_texts_lower = [
            (doc.page_content if hasattr(doc, "page_content") else str(doc)).lower()
            for doc in retrieved_docs
        ]
        for i in range(len(quote_words) - 2):
            phrase = ' '.join(quote_words[i:i+3])
            for doc, doc_lower in zip(retrieved_docs, doc_texts_lower):
                if phrase in doc_lower:
                    doc_source = extract_doc_source_label(doc)
                    logger.info(
                        f"  🔍 3-word phrase match with doc: '{phrase}' → {doc_source}")
                    return False, True, doc_source, "partial", False, "", _doc_source_url(doc)

    # 3b: 3-word phrase scan against user turns — only attribute to user when
    #     the phrase does NOT appear in any retrieved doc (disambiguates shared
    #     topic phrases).
    if len(quote_words) >= 3:
        doc_texts_lower_flat = [
            (doc.page_content if hasattr(doc, "page_content") else str(doc)).lower()
            for doc in (retrieved_docs or [])
        ]
        for i in range(len(quote_words) - 2):
            phrase = ' '.join(quote_words[i:i+3])
            phrase_in_doc = any(phrase in dt for dt in doc_texts_lower_flat)
            if phrase_in_doc:
                continue  # ambiguous — skip, will fall through to "General"
            for orig_l, pert_l, dp_on, orig_disp in turns_to_search:
                if phrase in orig_l:
                    logger.info(
                        f"  🔍 3-word phrase (user-only) match: '{phrase}'")
                    return True, False, "From your input", "partial", dp_on, orig_disp, ""
                if pert_l and phrase in pert_l:
                    logger.info(
                        f"  🔍 3-word phrase (perturbed, user-only) match: '{phrase}'")
                    return True, False, "From your input", "partial", dp_on, orig_disp, ""

    # No match found — the text is LLM-internal reasoning, not a source quote.
    logger.info(f"  ⚠️ No source match for evidence (treated as rationale): '{quote[:50]}...'")
    return False, False, "General", "none", False, "", ""


# ============================================================
# SECTION 15 – BUILD INFERENCE WARNING
# ============================================================

# def build_inference_warning(probs_rag, probs_no_rag, use_rag):
#     """Build yellow warning banner with lift threshold and tooltips."""
#     if not probs_rag:
#         return ""
#
#     items = []
#     max_lift = 0.0  # Track highest lift across all attributes
#
#     for attr in SENSITIVE_ATTRIBUTES:
#         dist = probs_rag.get(attr, {})
#         if not dist:
#             continue
#         top_val = max(dist, key=dist.get)
#         top_p = dist[top_val]
#
#         lift = 0.0
#         extra = ""
#         if use_rag and probs_no_rag:
#             p_nr = probs_no_rag.get(attr, {}).get(top_val, 0.0)
#             z = InferentialPrivacyMetrics.compute_z(top_p, p_nr)
#             lift = InferentialPrivacyMetrics.compute_lift(z)
#
#             # Track maximum lift
#             if lift > max_lift:
#                 max_lift = lift
#
#             # Only show attributes that meet the lift threshold
#             if lift > INFERENCE_LIFT_THRESHOLD and not math.isinf(lift):
#                 extra = (f'  <span style="color:#b71c1c;">'
#                          f'(+{lift:.0%} risk from external data)</span>')
#             else:
#                 # Skip this attribute if lift is below threshold
#                 continue
#
#         # Add attribute with tooltip - wrap in container with tooltip class
#         items.append(
#             f'<li><b>{attr}</b>: '
#             f'<span class="tooltip-container" style="position:relative;display:inline-block;cursor:help;border-bottom:1px dotted #999;">'
#             f'{top_val}'
#             f'<span class="tt-popup">This is a predicted value based on AI inference, '
#             f'not confirmed information about you.</span>'
#             f'</span> '
#             f'<span style="color:#555;">(p&nbsp;=&nbsp;{top_p:.0%})</span>'
#             f'{extra}</li>')
#
#     # Only show warning if at least one attribute exceeds threshold
#     if not items or (use_rag and max_lift <= INFERENCE_LIFT_THRESHOLD):
#         return ""
#
#     source_note = ("your input <b>and external data sources</b>"
#                    if use_rag else "your input alone")
#
#     # Add tooltip to the warning title - wrap in container
#     return (
#         '<div style="background:#FFF9C4;border:2px solid #F9A825;'
#         'border-radius:8px;padding:14px 18px;margin-top:14px;">'
#         '<span class="tt-anchor" style="position:relative;display:inline-block;cursor:help;border-bottom:1px dotted #E65100;">'
#         '<b style="color:#E65100;">⚠ Privacy Warning – Inferable Attributes</b>'
#         '<span class="tt-popup tt-above">'
#         '<div class="tt-title">These attributes are AI predictions based on your input '
#         'and may not accurately reflect your actual information.</div>'
#         '</span>'
#         '<span class="tt-popup tt-below">'
#         '<div class="tt-note">These attributes are AI predictions based on your input '
#         'and may not accurately reflect your actual information.</div>'
#         '</span>'
#         '</span><br>'
#         f'<span style="color:#333;">Based on {source_note}, the AI system '
#         f'could potentially infer the following about you:</span>'
#         f'<ul style="margin:8px 0 4px 18px;padding:0;">{"".join(items)}</ul>'
#         '</div>')

def build_inference_warning(user_input, probs_rag, probs_no_rag, use_rag,
                            evidence_rag=None, retrieved_docs=None,
                            input_dp_metadata=None, perturbed_user_input=None,
                            conversation_state=None):
    """Build yellow warning banner with lift threshold and explainability tooltips.

    Inference metrics are ALWAYS collected for every attribute so they can be
    logged regardless of whether the banner is ultimately shown to the user.
    The warning banner is only shown when at least one attribute exceeds
    INFERENCE_LIFT_THRESHOLD, but inference_metrics always captures everything.

    evidence_rag may contain either:
      - New format: {attr: [{"quote": "...", "type": "explicit"|"implicit"}, ...]}
      - Legacy format: {attr: ["...", ...]}   (plain strings)

    Returns:
        tuple: (warning_html, inference_metrics_dict, warning_shown_bool)
               warning_html       – HTML string (empty string if threshold not met)
               inference_metrics  – dict with ALL attributes, for logging
               warning_shown      – True iff the banner will be visible to the user
    """
    # print("In build_inference_warning")
    if not probs_rag:
        return "", {}, False

    if evidence_rag is None:
        evidence_rag = {}

    # Build prior_turns list: all user messages with original + perturbed text.
    prior_turns = []
    if conversation_state and conversation_state.messages:
        for msg in conversation_state.messages:
            if msg.role != "user":
                continue
            dp_m = getattr(msg, "dp_metadata", None)
            prior_turns.append({
                "original":  msg.content,
                "perturbed": dp_m.get("perturbed_text", "") if dp_m else "",
                "dp_active": bool(dp_m and dp_m.get("num_substitutions", 0) > 0),
            })

    items = []            # list items for the HTML banner (only threshold-passing attrs)
    max_lift = 0.0
    inference_metrics = {}   # collected for ALL attributes (logging)

    for attr in SENSITIVE_ATTRIBUTES:
        dist = probs_rag.get(attr, {})
        if not dist:
            continue
        top_val = max(dist, key=dist.get)
        top_p = dist[top_val]

        lift = 0.0
        p_nr = 0.0
        p_pop = POPULATION_PRIORS.get(attr, {}).get(top_val, 0.0)

        # Lift is always computed against the population prior so that
        # ALL experimental conditions (with or without RAG) produce a
        # non-zero score when the posterior exceeds the base rate.
        if p_pop > 0:
            z = InferentialPrivacyMetrics.compute_z(top_p, p_pop)
            lift = InferentialPrivacyMetrics.compute_lift(z)
            if lift > max_lift:
                max_lift = lift

        # Still track the no-RAG probability for logging / comparison.
        if probs_no_rag:
            p_nr = probs_no_rag.get(attr, {}).get(top_val, 0.0)

        # ── Determine dominant evidence_type for this attribute ──────────
        # evidence_rag[attr] is a list of either:
        #   new:    {"quote": str, "type": "explicit"|"implicit"}
        #   legacy: plain string  → treated as "implicit"
        raw_evidence = evidence_rag.get(attr, [])
        evidence_type = "unknown"
        if raw_evidence:
            for item in raw_evidence:
                if isinstance(item, dict):
                    t = item.get("type", "implicit")
                else:
                    t = "implicit"
                # If any evidence item is explicit, mark the whole attribute explicit
                if t == "explicit":
                    evidence_type = "explicit"
                    break
            else:
                evidence_type = "implicit"

        # ── Always record in inference_metrics (for logging) ─────────────
        if not (math.isinf(lift) or math.isnan(lift)):
            inference_metrics[attr] = {
                'top_value':     top_val,
                'confidence':    round(float(top_p), 4),
                'lift':          round(float(lift), 4),
                'prob_rag':      round(float(top_p), 4),
                'prob_no_rag':   round(float(p_nr), 4),
                'p_pop':         round(float(p_pop), 4),
                'evidence_type': evidence_type,
            }

        # ── Only add to the visible banner if threshold is met ───────────
        # print(lift, INFERENCE_LIFT_THRESHOLD)
        above_threshold = lift > INFERENCE_LIFT_THRESHOLD and not math.isinf(lift)
        # if not above_threshold:
        #     continue

        extra = (
            f'  <span style="color:#b71c1c;">(+{lift:.0%} increase based on external data)</span>'
            if above_threshold and use_rag else ""
        )

        # ── Build evidence tooltip ────────────────────────────────────────
        # Normalise evidence items to dicts so we can extract quote text
        normalised_evidence = []
        for item in raw_evidence:
            if isinstance(item, dict):
                quote_text = str(item.get("quote", item.get("text", ""))).strip()
                ev_type = item.get("type", "implicit")
            else:
                quote_text = str(item).strip()
                ev_type = "implicit"
            if (quote_text and
                    'no evidence' not in quote_text.lower() and
                    'not found' not in quote_text.lower() and
                    'no specific' not in quote_text.lower() and
                    'unavailable' not in quote_text.lower() and
                    len(quote_text) > 5):
                normalised_evidence.append({"quote": quote_text, "type": ev_type})

        if normalised_evidence:
            evidence_html = '<div class="tt-title">Evidence for Prediction</div>'
            evidence_html += '<div class="tt-sub" style="color:#666;font-size:1.1em!important;margin-bottom:6px;">The following evidence from your input or external data sources supports this prediction:</div>'
            evidence_html += '<div class="tt-divider"></div>'

            for ev_item in normalised_evidence[:3]:
                q = ev_item["quote"]
                ev_badge = (
                    '<span style="background:#e8f5e9;color:#2e7d32;border-radius:3px;'
                    'padding:1px 5px;font-size:0.8em;font-weight:600;margin-left:4px;">explicit</span>'
                    if ev_item["type"] == "explicit" else
                    '<span style="background:#fff3e0;color:#e65100;border-radius:3px;'
                    'padding:1px 5px;font-size:0.8em;font-weight:600;margin-left:4px;">implicit</span>'
                )

                q_esc = _esc(q)
                (is_from_user, is_from_docs, doc_source, _conf,
                 matched_turn_dp, matched_turn_orig, doc_url) = find_evidence_source(
                    q, user_input, retrieved_docs, min_overlap_ratio=0.4,
                    perturbed_input=perturbed_user_input,
                    prior_turns=prior_turns,
                )

                # Determine whether this is a genuine source excerpt or LLM
                # internal reasoning (rationale).  When neither user input nor
                # any retrieved document can be matched, the text is the model's
                # own thinking and should NOT be displayed like a quoted source.
                is_rationale = (not is_from_user) and (not is_from_docs)

                if is_from_user:
                    source_label = "From your input"
                    source_color = "#a119d2"
                elif is_from_docs:
                    source_label = doc_source
                    source_color = "#1976D2" if doc_source == "From uploaded data" else (
                        "#F57C00" if "(Internet)" in doc_source else "#666"
                    )
                else:
                    source_label = "AI reasoning"
                    source_color = "#757575"

                if len(q_esc) > ATTR_INFERENCE_EVIDENCE_STORED_EXCERPT_LENGTH:
                    q_esc = q_esc[:ATTR_INFERENCE_EVIDENCE_STORED_EXCERPT_LENGTH] + "..."

                evidence_html += '<div class="tt-section" style="margin-bottom:12px;">'
                evidence_html += f'<div class="tt-label" style="margin-bottom:3px;">'
                evidence_html += f'<span style="color:{source_color};font-size:0.85em!important;">● {source_label}</span>{ev_badge}'
                evidence_html += '</div>'

                if is_rationale:
                    # Display as plain italic reasoning text, not as a quoted excerpt.
                    evidence_html += (
                        f'<div class="tt-sub" style="margin-top:4px;color:#555;font-size:1.2em!important;'
                        f'font-style:italic;line-height:1.4;">{q_esc}</div>'
                    )
                else:
                    # Display as a quoted source excerpt.
                    evidence_html += f'<div class="tt-doc" style="margin-top:4px;">"{q_esc}"</div>'
                    # For doc sources, add a clickable URL if available.
                    if is_from_docs and doc_url:
                        url_esc = _esc(doc_url)
                        evidence_html += (
                            f'<div style="margin-top:4px;font-size:0.82em;">'
                            f'<span class="tt-k" style="color:#888;">Source link: </span>'
                            f'<a href="{url_esc}" target="_blank" '
                            f'style="color:#1976D2;word-break:break-all;text-decoration:underline;">'
                            f'{url_esc}</a></div>'
                        )
                    # When DP was active for the matched turn, note that the AI saw
                    # a word-substituted version, not the original wording.
                    if is_from_user and matched_turn_dp and matched_turn_orig:
                        orig_display = matched_turn_orig[:400] + ("…" if len(matched_turn_orig) > 400 else "")
                        orig_esc = _esc(orig_display)
                        evidence_html += (
                            f'<div style="margin-top:6px;padding:5px 8px;background:#fff8e1;'
                            f'border-left:3px solid #f9a825;border-radius:0 4px 4px 0;'
                            f'font-size:0.88em;color:#555;">'
                            f'🔒 Privacy protection was active for this message — the AI '
                            f'received a modified (word-substituted) version of your text, '
                            f'not your original wording.<br>'
                            f'<span style="display:block;margin-top:5px;font-size:0.92em;'
                            f'color:#333;font-style:italic;font-family:Georgia,serif;'
                            f'border-top:1px solid #f9a825;padding-top:4px;">'
                            f'Your original message: &ldquo;{orig_esc}&rdquo;'
                            f'</span>'
                            f'</div>'
                        )
                evidence_html += '</div>'
        else:
            evidence_html = '<div class="tt-title">Evidence for Prediction</div>'
            evidence_html += '<div class="tt-divider"></div>'
            evidence_html += ('<div class="tt-sub" style="color:#666;font-size:1.1em!important;margin-bottom:6px;">'
                              'No specific evidence from external sources was found to '
                              'support this prediction. The inference is based solely on the language '
                              "model's general knowledge and the user's input text.</div>")

        items.append(
            f'<li style="margin-bottom:6px;">'
            f'<b>{attr}</b>: {top_val} '
            # tt-anchor is the hover target — evidence lives INSIDE it
            f'<span class="tt-anchor" style="position:relative;display:inline-block;'
            f'cursor:help;color:#1976D2;font-size:0.9em;'
            f'border-bottom:1px dotted #1976D2;">'
            f'Why?'
            f'<span class="tt-popup tt-above" style="width:480px;max-width:480px;">'
            f'{evidence_html}'
            f'</span>'
            f'<span class="tt-popup tt-below" style="width:480px;max-width:480px;">'
            f'{evidence_html}'
            f'</span>'
            f'</span>'
            f' <span style="color:#555;font-size:0.88em;">(p&nbsp;=&nbsp;{top_p:.0%})</span>'
            f'{"<div style=margin-top:2px;>" + extra + "</div>" if extra else ""}'
            f'</li>'
        )

    # ── Decide whether to show the warning banner ─────────────────────────
    # print(max_lift, INFERENCE_LIFT_THRESHOLD)
    # if not items or (use_rag and max_lift <= INFERENCE_LIFT_THRESHOLD):
    if not probs_rag or not items:
        return "", inference_metrics, False

    is_warning = use_rag and max_lift > INFERENCE_LIFT_THRESHOLD and not math.isinf(max_lift)
    source_note = ("your input <b>and external data sources</b>"
                   if use_rag else "your input alone")

    if is_warning:
        bg_color = "#FFF9C4"
        border_color = "#F9A825"
        title_color = "#E65100"
        title_icon = "⚠"
        title_text = "Privacy Warning – Inferable Attributes"
        subtitle = (f'Based on {source_note}, the AI system '
                    f'could potentially infer the following about you:')
    else:
        bg_color = "#E3F2FD"
        border_color = "#1976D2"
        title_color = "#0D47A1"
        title_icon = "ℹ"
        title_text = "For your information – Inferable Attributes"
        subtitle = (f'Based on {source_note}, the AI system '
                    f'may be able to infer the following about you. '
                    f'No significant lift from external data was detected.')

    warning_html = (
        f'<div style="background:{bg_color};border:2px solid {border_color};'
        f'border-radius:8px;padding:14px 18px;margin-top:14px;">'
        f'<span class="tt-anchor" style="position:relative;display:inline-block;cursor:help;border-bottom:1px dotted {title_color};">'
        f'<b style="color:{title_color};">{title_icon} {title_text}</b>'
        f'<span class="tt-popup tt-above">'
        f'<div class="tt-title">These attributes are AI predictions based on your input '
        f'and may not accurately reflect your actual information.</div>'
        f'</span>'
        f'</span><br>'
        f'<span style="color:#333;">{subtitle}</span>'
        f'<ul style="margin:8px 0 4px 18px;padding:0;">{"".join(items)}</ul>'
        f'</div>'
    )

    return warning_html, inference_metrics, is_warning


# ============================================================
# SECTION 14 – CONVERSATION FORMATTING (WITH SMART TOOLTIPS)
# ============================================================

def _esc(text):
    """HTML escape."""
    return (text.replace("&", "&amp;").replace("<", "&lt;")
            .replace(">", "&gt;").replace('"', "&quot;")
            .replace("'", "&#39;"))

def _merge_overlapping_annotations(anns):
    """
    Split overlapping PII (pri=0) and RAG (pri=1) annotations into
    non-overlapping sub-spans so that both highlights are always visible.

    For each character range covered by both a PII and a RAG annotation a
    "combined" sub-span is produced with:
      • a diagonal gradient background blending both colours
      • a box-shadow outline in the PII border colour
      • a bottom-border in the RAG border colour
      • a "P R" superscript label with each letter in its own colour
      • both tooltips concatenated (PII first, divider, then RAG)

    The non-overlapping tails of each annotation are preserved with their
    original style.
    """
    pii_list = [a for a in anns if a.get("pri", 1) == 0]
    rag_list  = [a for a in anns if a.get("pri", 1) == 1]

    # If there is nothing to overlap, return as-is
    if not pii_list or not rag_list:
        return anns

    out = []
    # Track which (overlap_s, overlap_e) intervals belong to each source ann
    used_pii = []   # list of (ov_s, ov_e, pii_ann_object)
    used_rag  = []  # list of (ov_s, ov_e, rag_ann_object)

    for p in pii_list:
        for r in rag_list:
            ov_s = max(p["s"], r["s"])
            ov_e = min(p["e"], r["e"])
            if ov_s >= ov_e:
                continue  # no overlap

            combined_tip = (
                p.get("tip_html", "") +
                '<div class="tt-divider" '
                'style="margin:8px 0;border-top:2px dashed #bbb;"></div>'
                '<div style="font-size:0.78em;color:#888;font-style:italic;'
                'margin-bottom:4px;">Also linked to external data:</div>' +
                r.get("tip_html", "")
            )
            out.append({
                "s":         ov_s,
                "e":         ov_e,
                # Diagonal gradient: PII colour top-half, RAG colour bottom-half
                "bg":        f"linear-gradient(to bottom, {p['bg']} 52%, {r['bg']} 48%)",
                "bd":        r["bd"],      # RAG border-bottom
                "bd2":       p["bd"],      # PII outline via box-shadow
                "lbl_p_c":   p["bd"],      # colour for "P" superscript
                "lbl_r_c":   r["bd"],      # colour for "R" superscript
                "tip_html":  combined_tip,
                "pri":       0,
                "combined":  True,
            })
            used_pii.append((ov_s, ov_e, p))
            used_rag.append((ov_s, ov_e, r))

    # Subtract already-emitted overlap intervals from an annotation and
    # return the remaining (non-overlapping) sub-spans.
    def _remainders(ann, used_list):
        intervals = sorted(
            (ov_s, ov_e) for ov_s, ov_e, src in used_list if src is ann
        )
        if not intervals:
            return [ann]
        # Merge adjacent/touching intervals
        merged = [list(intervals[0])]
        for iv_s, iv_e in intervals[1:]:
            if iv_s <= merged[-1][1]:
                merged[-1][1] = max(merged[-1][1], iv_e)
            else:
                merged.append([iv_s, iv_e])
        # Gaps between merged intervals → individual sub-spans
        parts = []
        cur = ann["s"]
        for iv_s, iv_e in merged:
            if cur < iv_s:
                parts.append({**ann, "s": cur, "e": iv_s})
            cur = iv_e
        if cur < ann["e"]:
            parts.append({**ann, "s": cur, "e": ann["e"]})
        return parts

    for p in pii_list:
        out.extend(_remainders(p, used_pii))
    for r in rag_list:
        out.extend(_remainders(r, used_rag))

    return out

def render_highlighted(text, pii_matches, rag_links, show_tips=True, show_rag_hl=True, show_pii_hl=True):
    """Render text with colored highlight spans and smart-positioned tooltips."""
    if text is None:
        text = ""
    if not isinstance(text, str):
        text = str(text)

    anns = []

    if show_pii_hl and pii_matches:
        for m in pii_matches:
            c = PII_COLORS[m.category]
            tip_html = (
                f'<div class="tt-title">{c["label"]}</div>'
                f'<div class="tt-sub">Detected: {_esc(m.text)}</div>'
                f'<div class="tt-divider"></div>'
                f'<div class="tt-note">Type: {m.fine_type} | Confidence: {m.confidence:.0%}</div>'
            )
            anns.append({
                "s": m.start,
                "e": m.end,
                "bg": c["bg"],
                "bd": c["border"],
                "lbl": "P",
                "tip_html": tip_html,
                "pri": 0,
            })

    if show_rag_hl and rag_links:
        for lk in rag_links:
            kw_txt = ", ".join(lk.overlap_keywords or ["N/A"])

            # Determine source color based on source type
            source = getattr(lk, 'source', 'From system data')
            if source == "From uploaded data":
                source_color = "#1976D2"  # Blue
            elif "Internet" in source:
                source_color = "#F57C00"  # Orange
            else:
                source_color = "#666"  # Gray

            # Build optional URL line for social media / web sources
            lk_url = getattr(lk, 'url', '')
            url_line = ""
            if lk_url and "Internet" in source:
                display_url = lk_url # if len(lk_url) <= 80 else lk_url[:77] + "..."
                url_line = (
                    f'<div><span class="tt-k">Source link:</span> '
                    f'<span style="color:#1976D2;word-break:break-all;">'
                    f'<a href="{_esc(lk_url)}" style="color:inherit;text-decoration:underline;" target="_blank">'
                    f'{_esc(display_url)}</a>'
                    f'</span></div>'
                )

            tip_html = (
                '<div class="tt-title">External Data Linkage</div>'
                f'<div class="tt-sub">Text: {_esc(lk.text)}</div>'
                '<div class="tt-divider"></div>'
                '<div class="tt-section">'
                '<div class="tt-label">Top Retrieved Document:</div>'
                f'<div class="tt-doc">{_esc(lk.top_doc_text)}</div>'
                '</div>'
                '<div class="tt-section tt-metrics">'
                # f'<div><span class="tt-k">Retrieval similarity:</span> {float(lk.top_doc_score):.2f}</div>'
                # f'<div><span class="tt-k">Matched keywords:</span> {_esc(kw_txt)}</div>'
                f'<div><span class="tt-k">Combined similarity:</span> {float(lk.top_similarity):.2f}</div>'
                f'<div><span class="tt-k">Retrieval score:</span> {float(lk.top_doc_score):.2f}</div>'
                f'<div><span class="tt-k">Matched keywords:</span> {_esc(kw_txt)}</div>'
                f'<div><span class="tt-k">Source:</span> <span style="color:{source_color};font-weight:bold;">{source}</span></div>'
                f'{url_line}'
                '</div>'
                '<div class="tt-divider"></div>'
                '<div class="tt-note">This linkage increases the chance that sensitive attributes can be inferred by combining your text with external data.</div>'
            )
            anns.append({
                "s": lk.start,
                "e": lk.end,
                "bg": RAG_LINK_COLOR["bg"],
                "bd": RAG_LINK_COLOR["border"],
                "lbl": "R",
                "tip_html": tip_html,
                "pri": 1,  # RAG yields to PII in overlaps
            })

    if not anns:
        return _esc(text)
    #     return f"<p>{_esc(text)}</p>"

    # anns.sort(key=lambda a: (a["s"], -a["e"]))
    anns = _merge_overlapping_annotations(anns)
    anns.sort(key=lambda a: (a["s"], a.get("pri", 1), -a["e"]))

    parts, last = [], 0
    for a in anns:
        if a["s"] < last:
            continue
        if a["s"] > last:
            parts.append(_esc(text[last:a["s"]]))
        st = _esc(text[a["s"] : a["e"]])

        # CSS-based tooltip with smart positioning
        cls = "hl"
        tooltip_html = ""
        if show_tips and a.get("tip_html"):
            cls += " tt-anchor"
            # Add both above and below variants - JS will choose
            tooltip_html = (
                f'<span class="tt-popup tt-above">{a["tip_html"]}</span>'
                f'<span class="tt-popup tt-below">{a["tip_html"]}</span>'
                f'<span class="tt-popup tt-left">{a["tip_html"]}</span>'
                f'<span class="tt-popup tt-right">{a["tip_html"]}</span>'
                f'<span class="tt-popup tt-below-left">{a["tip_html"]}</span>'
                f'<span class="tt-popup tt-below-right">{a["tip_html"]}</span>'
            )

        if a.get("combined"):
            # Two-colour span: diagonal gradient bg + RAG bottom border +
            # PII outline via box-shadow.  Label: "P" (PII colour) "R" (RAG colour).
            label_html = (
                f'<sup style="font-size:0.7em;">'
                f'<span style="color:{a["lbl_p_c"]};">P</span>'
                f'<span style="color:{a["lbl_r_c"]};">R</span>'
                f'</sup>'
            )
            parts.append(
                f'<span class="{cls}" style="background:{a["bg"]};'
                f'border-bottom:3px solid {a["bd"]};'
                f'box-shadow:0 0 0 1.5px {a["bd2"]};'
                f'padding:0 2px;border-radius:3px;cursor:help;line-height:1;" '
                f'data-type="ann">{st}{label_html}{tooltip_html}</span>'
            )
        else:
            parts.append(
                f'<span class="{cls}" style="background:{a["bg"]};'
                f'border-bottom:3px solid {a["bd"]};padding:0 2px;'
                f'border-radius:3px;cursor:help;line-height:1;" data-type="ann">{st}'
                f'<sup style="font-size:0.7em;color:{a["bd"]};">{a["lbl"]}</sup>'
                f'{tooltip_html}</span>'
            )
        last = a["e"]
    if last < len(text):
        parts.append(_esc(text[last:]))

    # Wrap in a div to ensure proper rendering
    return "".join(parts)
    # if last < len(text):
    #     parts.append(_esc(text[last:]))
    # return "".join(parts)


# def fmt_conversation(messages, show_tips=True, show_rag_hl=True, show_pii_hl=True):
#     """Format conversation with highlights and smart tooltips."""
#     css = (
#         "<style>"
#         ".chat-box{max-height:500px;overflow-y:auto;overflow-x:visible;padding:12px;position:relative;}"
#         ".msg{margin:12px 0;padding:16px;border-radius:10px;border:2px solid #bbb;overflow:visible;position:relative;line-height:1.5}"
#         ".msg-user{background:#e3f2fd;margin-left:18%;margin-top:4%;line-height:1.5}"
#         ".msg-ai{background:#f5f5f5;margin-right:18%;line-height:1.5}"
#         ".msg-role{font-weight:bold;margin-bottom:6px;color:#333;}"
#         ".hl{position:relative;}"
#
#         # Tooltip content styling
#         ".tt-popup .tt-title{font-weight:700;font-size:1.1em!important;margin-bottom:5px;color:#111;}"
#         ".tt-popup .tt-sub{color:#444;margin-bottom:2px;font-size:0.95em!important;}"
#         ".tt-popup .tt-section{margin-top:8px;}"
#         ".tt-popup .tt-label{font-weight:600;color:#2c3e50;margin-bottom:5px;font-size:1em!important;}"
#         ".tt-popup .tt-doc{font-size:1em!important;color:#000;background:#f8f8f8;"
#         "border:1.5px solid #d0d0d0;border-radius:6px;padding:8px;"
#         "white-space:pre-wrap;font-weight:600;font-family:'Courier New',monospace;}"
#         ".tt-popup .tt-metrics{margin-top:6px;color:#333;font-size:0.95em!important;}"
#         ".tt-popup .tt-k{font-weight:600;color:#1f2d3d;}"
#         ".tt-popup .tt-divider{height:1px;background:#ddd;margin:8px 0;}"
#         ".tt-popup .tt-note{color:#555;font-size:0.9em!important;}"
#
#          # Tooltip positioning base - use absolute instead of fixed
#         ".tt-anchor{position:relative;display:inline-block;}"
#         ".tt-popup{display:none;position:absolute;z-index:99999;background:#fff;color:#111;"
#         "border:2px solid #bbb;border-radius:10px;padding:10px 12px;width:max-content;"
#         "max-width:500px;box-shadow:0 8px 20px rgba(0,0,0,0.15);font-family:inherit;"
#         "font-size:0.7em!important;line-height:1.25!important;text-align:left;pointer-events:none;}"
#
#          # Position above (default) - relative to anchor
#         ".tt-popup.tt-above{bottom:100%;left:50%;transform:translateX(-50%);margin-bottom:5px;}"
#
#         # Position below (when near top) - relative to anchor
#         ".tt-popup.tt-below{top:100%;left:50%;transform:translateX(-50%);margin-top:5px;}"
#
#         # Show only the appropriate one on hover
#         ".tt-anchor:hover .tt-popup.tt-above{display:block;}"
#         ".tt-anchor:hover .tt-popup.tt-below{display:none;}"
#
#         # If anchor is in top 30% of viewport, show below instead
#         ".tt-anchor.near-top:hover .tt-popup.tt-above{display:none;}"
#         ".tt-anchor.near-top:hover .tt-popup.tt-below{display:block;}"
#
#         # Ensure chat box allows overflow for tooltips
#         ".chat-box{overflow-y:auto;overflow-x:visible;position:relative;}"
#
#         """
#         /* Override Gradio theme variables */
#         :root {
#             --background-fill-primary: #ffffff !important;
#             --background-fill-secondary: #ffffff !important;
#             --block-background-fill: #ffffff !important;
#             --panel-background-fill: #ffffff !important;
#         }
#
#         .dark {
#             --background-fill-primary: #ffffff !important;
#             --background-fill-secondary: #ffffff !important;
#             --block-background-fill: #ffffff !important;
#             --panel-background-fill: #ffffff !important;
#         }
#
#         /* Force white background on all groups */
#         .gr-group,
#         .gr-box,
#         .gr-panel,
#         div[class*="group"],
#         div[class*="panel"] {
#             background: #ffffff !important;
#             background-color: #ffffff !important;
#         }
#
#         /* Gray header for panel titles */
#         .gr-group > .gr-prose > h3:first-child,
#         .gr-group > div > .gr-prose > h3:first-child {
#             background: #f5f5f5 !important;
#             padding: 10px 12px !important;
#             margin: -16px -16px 12px -16px !important;
#             border-radius: 6px 6px 0 0 !important;
#             border-bottom: 1px solid #ddd !important;
#         }
#
#         /* Padding for content */
#         .gr-group .gr-prose {
#             padding-left: 12px !important;
#             padding-right: 12px !important;
#         }
#
#         /* Keep rest of your existing CSS */
#         .gradio-container {
#             background: #fff !important;
#         }
#
#         textarea, input[type="text"], input[type="number"], select,
#         .gr-input, .gr-text-input, .input-text, .border, div[class*="border"] {
#             border-width: 2px !important;
#             border-style: solid !important;
#             border-color: #aaa !important;
#         }
#
#         textarea:focus, input:focus, select:focus {
#             border-color: #4A90D9 !important;
#             border-width: 2.5px !important;
#         }
#
#         button {
#             border-width: 2px !important;
#             border-style: solid !important;
#         }
#
#         /* White space for conversation */
#         .msg-ai {
#             background: #f5f5f5;
#             margin-right: 18%;
#             line-height: 1.8;
#             white-space: pre-wrap;
#         }
#
#         .msg-user {
#             background: #e3f2fd;
#             margin-left: 18%;
#             line-height: 1.8;
#             white-space: pre-wrap;
#         }
#         """
#
#
#
#         "</style>"
#
#         # JavaScript to detect position and add 'near-top' class
#         "<script>"
#         "document.addEventListener('DOMContentLoaded', function() {"
#         "  const observer = new IntersectionObserver((entries) => {"
#         "    entries.forEach(entry => {"
#         "      const rect = entry.target.getBoundingClientRect();"
#         "      const viewportHeight = window.innerHeight;"
#         "      if (rect.top < viewportHeight * 0.3) {"
#         "        entry.target.classList.add('near-top');"
#         "      } else {"
#         "        entry.target.classList.remove('near-top');"
#         "      }"
#         "    });"
#         "  }, {threshold: 0});"
#         "  "
#         "  document.querySelectorAll('.tt-anchor').forEach(el => {"
#         "    observer.observe(el);"
#         "  });"
#         "});"
#         "</script>"
#     )
#
#     h = [css, '<div class="chat-box">']
#     for m in messages:
#         role_cls = "msg-user" if m.role == "user" else "msg-ai"
#         role_label = "You" if m.role == "user" else "Assistant"
#         # content_html = render_highlighted(m.content, m.pii_matches, m.rag_links, show_tips, show_rag_hl, show_pii_hl)
#         if m.rag_links or m.pii_matches:
#             content_html = render_highlighted(m.content, m.pii_matches, m.rag_links, show_tips, show_rag_hl,
#                                               show_pii_hl)
#         else:
#             formatted_content = _format_response_text(m.content)
#             content_html = render_highlighted(formatted_content, m.pii_matches, m.rag_links, show_tips, show_rag_hl,
#                                                   show_pii_hl)
#         h.append(
#             f'<div class="msg {role_cls}">'
#             f'<div class="msg-role">{role_label}</div>'
#             f'{content_html}</div>'
#         )
#     h.append('</div>')
#     return "".join(h)

def fmt_conversation(messages, show_tips=True, show_rag_hl=True, show_pii_hl=True):
    """Format conversation with highlights and smart tooltips."""

    css = (
        "<style>"
        ".chat-box{max-height:50vh;overflow-y:auto;overflow-x:visible;position:relative;}"
        ".msg{margin:12px 0;padding:16px;border-radius:10px;border:2px solid #bbb;overflow:visible;position:relative;line-height:1.5}"
        ".msg-user{background:#e3f2fd;margin-left:18%;margin-top:4%;line-height:1.5}"
        ".msg-ai{background:#f5f5f5;margin-right:18%;line-height:1.5}"
        ".msg-role{font-weight:bold;margin-bottom:6px;color:#333;}"
        ".hl{position:relative;}"

        # Tooltip content styling
        ".tt-popup .tt-title{font-weight:700;font-size:1.2em!important;margin-bottom:5px;color:#111;}"
        ".tt-popup .tt-sub{color:#444;margin-bottom:2px;font-size:1.2em!important;}"
        ".tt-popup .tt-section{margin-top:8px;}"
        ".tt-popup .tt-label{font-weight:600;color:#2c3e50;margin-bottom:5px;font-size:1.15em!important;}"
        ".tt-popup .tt-doc{font-size:1em!important;color:#000;background:#f8f8f8;"
        "border:1.5px solid #d0d0d0;border-radius:6px;padding:8px;"
        "white-space:pre-wrap;font-weight:600;font-family:'Courier New',monospace;}"
        ".tt-popup .tt-metrics{margin-top:6px;color:#333;font-size:0.95em!important;}"
        ".tt-popup .tt-k{font-weight:600;color:#1f2d3d;font-size:1.1em!important}"
        ".tt-popup .tt-divider{height:1px;background:#ddd;margin:8px 0;}"
        ".tt-popup .tt-note{color:#555;font-size:1.1em!important;}"

        # Tooltip positioning with margin awareness
        # pointer-events is ALWAYS auto so that when the mouse moves from the
        # anchor word into the popup area (crossing the 8 px gap) the browser
        # still registers the hover and keeps the popup visible.
        # visibility + opacity are delayed 0.7 s before hiding so the user has
        # time to move the mouse from the anchor into the popup.
        ".tt-popup{opacity:0;visibility:hidden;pointer-events:auto;"
        "position:fixed;z-index:99999;background:#fff;color:#111;"
        "border:2px solid #bbb;border-radius:10px;padding:10px 12px;width:max-content;"
        "max-width:min(500px,calc(100vw - 40px));box-shadow:0 8px 20px rgba(0,0,0,0.15);font-family:inherit;"
        "font-size:0.7em!important;line-height:1.25!important;text-align:left;"
        "bottom:auto;left:0;top:0;transform:none;" 
        "transition:opacity 0s 0.15s,visibility 0s 0.15s;}"

        # Prevent tooltip from going beyond left edge
        # ".tt-popup{left:max(10px,min(50%,calc(100% - 10px)));transform:translateX(min(0px,max(-100%,calc(-50% - 10px))));}"

        # Show immediately on anchor hover OR when mouse is on the popup itself.
        # Both paths are always reachable because pointer-events is never none.
        ".tt-anchor:hover .tt-popup,.tt-popup:hover{opacity:1;visibility:visible;transition:opacity 0s 0s,visibility 0s 0s;}"

        # Position below when near top
        # ".tt-anchor.near-top .tt-popup{bottom:auto;top:calc(100% + 8px);}"

        """
        /* Override Gradio theme variables */
        :root {
            --background-fill-primary: #ffffff !important;
            --background-fill-secondary: #ffffff !important;
            --block-background-fill: #ffffff !important;
            --panel-background-fill: #ffffff !important;
        }

        /* Force white background on all groups */
        .gr-group,
        .gr-box,
        .gr-panel,
        div[class*="group"],
        div[class*="panel"] {
            background: #ffffff !important;
            background-color: #ffffff !important;
        }

        /* Gray header for panel titles */
        .gr-group > .gr-prose > h3:first-child,
        .gr-group > div > .gr-prose > h3:first-child {
            background: #f5f5f5 !important;
            padding: 10px 12px !important;
            margin: -16px -16px 12px -16px !important;
            border-radius: 6px 6px 0 0 !important;
            border-bottom: 1px solid #ddd !important;
        }

        /* Padding for content */
        .gr-group .gr-prose {
            padding-left: 12px !important;
            padding-right: 12px !important;
        }

        /* Keep rest of your existing CSS */
        .gradio-container {
            background: #fff !important;
        }

        textarea, input[type="text"], input[type="number"], select,
        .gr-input, .gr-text-input, .input-text, .border, div[class*="border"] {
            border-width: 2px !important;
            border-style: solid !important;
            border-color: #aaa !important;
        }

        textarea:focus, input:focus, select:focus {
            border-color: #4A90D9 !important;
            border-width: 2.5px !important;
        }

        button {
            border-width: 2px !important;
            border-style: solid !important;
        }

        /* White space for conversation */
        .msg-ai {
            background: #f5f5f5;
            margin-right: 18%;
            line-height: 1.8;
            white-space: pre-line;
            word-wrap: break-word;
            overflow-wrap: break-word;
        }
        
        .msg-user {
            background: #e3f2fd;
            margin-left: 18%;
            line-height: 1.5;
            white-space: pre-line;
            word-wrap: break-word;
            overflow-wrap: break-word;
        }
        """

        "</style>"

    )
    h = [css, '<div class="chat-box">']
    for m in messages:
        role_cls = "msg-user" if m.role == "user" else "msg-ai"
        role_label = "You" if m.role == "user" else "Assistant"

        # if m.rag_links or m.pii_matches:
        #     content_html = render_highlighted(m.content, m.pii_matches, m.rag_links, show_tips, show_rag_hl,
        #                                       show_pii_hl)
        # else:
        #     formatted_content = _format_response_text(m.content)
        #     content_html = render_highlighted(formatted_content, m.pii_matches, m.rag_links, show_tips, show_rag_hl,
        #                                       show_pii_hl)
        # Always format text first, regardless of annotations
        content_html = render_highlighted(m.content, m.pii_matches, m.rag_links, show_tips, show_rag_hl, show_pii_hl)

        # Append DP badge to user messages that were perturbed
        badge_html = ""
        if m.role == "user" and getattr(m, "dp_metadata", None):
            badge_html = _dp_badge_html(m.dp_metadata)

        if badge_html:
            badge_row = (
                f'<div style="display:flex;justify-content:center;'
                f'margin-top:6px;">{badge_html}</div>'
            )
        else:
            badge_row = ""

        h.append(
            f'<div class="msg {role_cls}">'
            f'<div class="msg-role">{role_label}</div>'
            f'{content_html}{badge_row}</div>'
        )
    h.append('</div>')
    return "".join(h)

def clear_conversation(state, show_tips, show_rag_hl, show_pii_hl):
    """Clear conversation."""
    state.clear()
    return (
        fmt_conversation([], show_tips, show_rag_hl, show_pii_hl),
        create_risk_display(0),
        "",
        "", "", gr.update(visible=False), gr.update(visible=False)
    )


# ============================================================
# SECTION 15 – LEGEND
# ============================================================

def create_legend_html(show_tips=True, show_rag_hl=True):
    """Create legend HTML."""
    items = []
    if show_tips:
        for _c, col in PII_COLORS.items():
            items.append(
                f'<div style="display:flex;align-items:center;margin:6px 0;">'
                f'<span style="display:inline-block;width:22px;height:22px;'
                f'background:{col["bg"]};border:2.5px solid {col["border"]};'
                f'border-radius:4px;margin-right:10px;flex-shrink:0;"></span>'
                f'<span>{col["label"]}</span></div>')
    if show_rag_hl:
        items.append(
            f'<div style="display:flex;align-items:center;margin:6px 0;">'
            f'<span style="display:inline-block;width:22px;height:22px;'
            f'background:{RAG_LINK_COLOR["bg"]};border:2.5px solid {RAG_LINK_COLOR["border"]};'
            f'border-radius:4px;margin-right:10px;flex-shrink:0;"></span>'
            f'<span>{RAG_LINK_COLOR["label"]}</span></div>')
    return "".join(items)


# ============================================================
# SECTION 16 – GRADIO UI
# ============================================================

def create_ui(
    model=None,
    rag="1",
    epsilon="inf",
    show_risk="0",
    show_tips="0",
    show_rag_highlights="0",
    show_pii_highlights="1",
    show_settings="0",
    demo="0",
    enable_social_scraping="0",
    show_social_scraping="0",
    show_upload_data = "1",
    scenario_mode="real",
    retriever=None,
    show_dp="1",
    show_infr_attr_card="1",
):
    """Create Gradio UI with settings panel."""

    # state = ConversationState()

    def _to_bool(v):
        if v is None or v == "":
            return False
        if isinstance(v, (int, float)):
            return bool(int(v))
        s = str(v).strip().lower()
        return s in {"1", "true", "t", "yes", "y", "on"}

    def _to_float(v, default=float("inf")):
        if v is None or v == "":
            return default
        s = str(v).strip().lower()
        if s in {"inf", "infty", "infinite", "infinity"}:
            return float("inf")
        try:
            return float(s)
        except Exception:
            return default

    def _make_demo_prompt():
        tpl = random.choice(list(DATASET_PROMPTS_PANORAMA.values()))
        return tpl.format(name="Raymond Phillips")

    # Defaults
    use_rag_default = _to_bool(rag)
    eps_default = _to_float(epsilon, default=float("inf"))
    model_default = str(model or DEFAULT_MODEL_NAME)
    show_risk_default = _to_bool(show_risk)
    show_tips_default = _to_bool(show_tips)
    show_raghl_default = _to_bool(show_rag_highlights)
    # show_pii_hl_default = _to_bool(highlight_pii)
    show_pii_hl_default = _to_bool(show_pii_highlights)
    show_settings_default = _to_bool(show_settings)
    show_social_scraping_default = _to_bool(show_social_scraping)
    show_upload_data_default = _to_bool(show_upload_data)
    demo_default = _to_bool(demo)
    social_scraping_default = _to_bool(enable_social_scraping)
    scenario_mode_default = str(scenario_mode or "real").strip().lower()

    # show_dp:  0=never, 1=after each turn (default), 2=beginning only, 3=on End button
    # show_infr_attr_card: 0=never, 1=after each turn (default), 2=on End button
    show_dp_default       = int(_to_float(show_dp,            default=1.0))
    show_infr_default     = int(_to_float(show_infr_attr_card, default=1.0))

    # Initialize social scraper
    initialize_social_scraper(enabled=social_scraping_default)

    # vs = get_retriever()
    _get_system_retriever_for_scenario = get_scenario_retriever

    with gr.Blocks(
        # css=css,
        title="What can LLMs infer about me?",
        theme=gr.themes.Default(spacing_size=gr.themes.sizes.spacing_sm, radius_size=gr.themes.sizes.radius_none),
    ) as app:
        access_control_css = gr.HTML("")
        if show_settings_default:
            gr.Markdown(
                "# 🔒 What can LLMs infer about me?\n\n"
                "Analyze privacy risks and understand what LLMs can infer from your data."
            )

        # Session-scoped configuration
        rag_st = gr.State(use_rag_default)
        eps_st = gr.State(eps_default)
        mdl_st = gr.State(model_default)
        tips_st = gr.State(show_tips_default)
        raghl_st = gr.State(show_raghl_default)
        piihl_st = gr.State(show_pii_hl_default)
        riskpanel_st = gr.State(show_risk_default)
        demo_st = gr.State(demo_default)
        demo_prompt_st = gr.State(_make_demo_prompt() if demo_default else "")
        vs_st = gr.State(None)
        social_scraping_st = gr.State(False)  # Social media scraping toggle
        custom_corpus_loaded_st = gr.State(False)  # Track if custom corpus is loaded
        uploaded_file_path_st = gr.State(None)
        state_st = gr.State(ConversationState())
        # New: control panel visibility modes and conversation-ended flag
        show_dp_st   = gr.State(show_dp_default)
        show_infr_st = gr.State(show_infr_default)
        conv_ended_st = gr.State(False)

        with gr.Row():
            # LEFT: legend + settings panel
            with gr.Column(min_width=220, scale=1) as legend_col:
                with gr.Group(visible=False) as legend_group:
                    # with gr.Row():
                    gr.Markdown("### 🎨 Legend")
                    legend_html = gr.HTML("")

                # Settings panel (toggled by URL parameter)
                with gr.Group(visible=show_settings_default) as settings_panel:
                    gr.Markdown("### ⚙️ Settings")

                    # Model selector
                    model_dropdown = gr.Dropdown(
                        choices=[name for _, _, name in MODEL_CONFIGS],
                        value=model_default,
                        label="Model",
                        interactive=True
                    )

                    # Epsilon slider - Differential Privacy Guarantee (Noise Level)
                    epsilon_slider = gr.Slider(
                        minimum=0.1,
                        maximum=MAX_POSSIBLE_EPS,
                        value=min(eps_default, MAX_POSSIBLE_EPS) if eps_default != float('inf') else MAX_POSSIBLE_EPS,
                        label="Differential Privacy Guarantee (Noise Level)",
                        info="",#"Controls the amount of random noise (perturbation) added during linkage with external data. Lower values = more noise = stronger privacy protection. Higher values = less noise = better accuracy but weaker privacy guarantees.",
                        interactive=True
                    )

                    # Infinite epsilon checkbox
                    eps_inf_checkbox = gr.Checkbox(
                        value=(eps_default == float('inf')),
                        label="Disable Differential Privacy",
                        interactive=True
                    )

                with gr.Group(visible=False, elem_id="privacy_settings_panel") as privacy_settings_group:
                    gr.Markdown("### 🔒 Privacy Settings")
                    privacy_settings_html = gr.HTML("")

                social_scraping_group = gr.Group(visible=show_social_scraping_default)
                with social_scraping_group:
                    # Social Media Scraping toggle
                    gr.Markdown("#### 🌐 Access to Social Media")
                    social_scraping_checkbox = gr.Checkbox(
                        value=social_scraping_default,
                        label="Allow access to my social media",
                        interactive=True
                    )
                    social_scraping_info = gr.HTML(
                        """
                        <div style="color: #000000 !important; font-size: 0.95em; margin-bottom: 12px; line-height: 1.5;">
                            Extract data from social media platforms in real-time.<br/>
                            This may take up to a minute.
                        </div>
                        """
                    )

                    social_platforms = gr.CheckboxGroup(
                        choices=[
                            "Twitter",
                            "Facebook",
                            "LinkedIn",
                            "Web",
                        ],
                        value=["Twitter","Facebook", "LinkedIn", "Web"],
                        label="Platforms",
                        info=(
                            "Twitter |  "
                            "Facebook |  "
                            "LinkedIn |  "
                            "Web"
                        ),
                        visible=social_scraping_default,
                        interactive=True
                    )

                upload_data_group = gr.Group(visible=show_upload_data_default)
                with upload_data_group:
                    # Custom Corpus Upload
                    gr.Markdown("#### 📁 Upload My Data")
                    corpus_file = gr.File(
                        label="Upload CSV file",
                        file_types=None, #[".csv"],
                        type="filepath",
                        interactive=True
                    )

                    corpus_upload_info = gr.HTML(
                        """
                        <div style="color: #000000 !important; font-size: 0.95em; margin-bottom: 12px; line-height: 1.5;">
                            Download your textual social media data from Facebook, LinkedIn, Twitter (X). 
                            CSV should have columns: <b>text</b> (required), <b>user id</b> (if known), <b>First Name</b>, <b>Last Name</b>
                        </div>
                        """
                    )

                    corpus_upload_btn = gr.Button("Upload", variant="secondary", size="sm")
                    corpus_status = gr.Markdown(
                        "",
                        elem_id="corpus_status",
                        elem_classes="corpus-status-box"
                    )


            # CENTER: conversation
            with gr.Column(scale=8):
                conversation_panel_title = "### 💬 Conversation"
                if show_raghl_default:
                    conversation_panel_title += " | Hover over highlighted text to see why it was flagged."
                gr.Markdown(conversation_panel_title)
                conv_html = gr.HTML(fmt_conversation([], show_tips_default, show_raghl_default, show_pii_hl_default))
                warn_html = gr.HTML("")

                persona_hint = gr.HTML("")
                user_tb = gr.Textbox(
                    value="",
                    placeholder="Type your message here…",
                    label="Your Message",
                    lines=3,
                )
                with gr.Row():
                    send_btn = gr.Button("Send", variant="primary", scale=1)
                    # clear_btn = gr.Button("Clear", scale=1)
                    reset_btn = gr.Button("Reset conversation", scale=1)
                    end_btn   = gr.Button("End conversation", variant="stop", scale=1)

            # RIGHT: risk panel
            right_col_needed = show_risk_default or show_pii_hl_default or (show_infr_default != 0)
            with gr.Column(scale=3, visible=right_col_needed) as risk_col:
                risk_title = gr.Markdown("", visible=bool(show_risk_default))
                risk_html  = gr.HTML("", visible=bool(show_risk_default))
                analysis_title = gr.Markdown("### 📝 Analysis", visible=False)
                infer_card_html = gr.HTML("", visible=(show_infr_default != 0))
                analysis_md = gr.HTML("", visible=False)

        # Page-load: apply URL query parameters
        # def _apply_query_params(request: gr.Request):
        #     qp = getattr(request, "query_params", {}) or {}
        #
        #     access_token = qp.get("token", "")
        #     if not validate_access_token(access_token):
        #         # Invalid or missing token - return error state
        #         error_msg = """
        #                 <div style='padding: 40px; text-align: center; color: #d32f2f;'>
        #                     <h2>⚠️ Access Denied</h2>
        #                     <p>Invalid or missing access token.</p>
        #                     <p>Please contact the administrator for a valid access link.</p>
        #                 </div>
        #                 """
        #         return (
        #             error_msg,  # legend_html - show error
        #             gr.update(visible=False),  # settings_panel
        #             gr.update(visible=False),  # risk_col
        #             False,  # rag_st
        #             1.0,  # eps_st
        #             DEFAULT_MODEL_NAME,  # mdl_st
        #             False,  # tips_st
        #             False,  # raghl_st
        #             False,  # piihl_st
        #             False,  # riskpanel_st
        #             False,  # demo_st
        #             "",  # demo_prompt_st
        #             error_msg,  # conv_html - show error
        #             "",  # risk_html
        #             "",  # warn_html
        #             "",  # analysis_md
        #             "",  # user_tb
        #             DEFAULT_MODEL_NAME,  # model_dropdown
        #             1.0,  # epsilon_slider
        #             False,  # eps_inf_checkbox
        #         )
        #
        #     mdl = str(qp.get("model", model_default))
        #     use_rag = _to_bool(qp.get("rag", use_rag_default))
        #     eps = _to_float(qp.get("epsilon", eps_default), default=eps_default)
        #
        #     show_risk_v = _to_bool(qp.get("show_risk", show_risk_default))
        #     show_tips_v = _to_bool(qp.get("show_tips", show_tips_default))
        #     show_raghl_v = _to_bool(qp.get("show_rag_highlights", show_raghl_default))
        #     show_pii_hl_v = _to_bool(qp.get("highlight_pii", show_pii_hl_default))
        #     show_settings_v = _to_bool(qp.get("show_settings", show_settings_default))
        #     demo_v = _to_bool(qp.get("demo", demo_default))
        #
        #     legend_visible = bool(show_tips_v or show_raghl_v or show_pii_hl_v)
        #     # demo_prompt = _make_demo_prompt() if demo_v else ""
        #     demo_prompt = DEFAULT_USER_TEXT if DEFAULT_USER_TEXT else ""
        #
        #     legend = create_legend_html(show_tips=show_tips_v, show_rag_hl=show_raghl_v)
        #
        #     eps_is_inf = (eps == float('inf'))
        #     eps_slider_val = min(eps, 10.0) if not eps_is_inf else 10.0
        #
        #     return (
        #         legend,
        #         gr.update(visible=show_settings_v),
        #         gr.update(visible=show_risk_v),
        #         use_rag,
        #         eps,
        #         mdl,
        #         show_tips_v,
        #         show_raghl_v,
        #         show_pii_hl_v,
        #         show_risk_v,
        #         demo_v,
        #         demo_prompt,
        #         fmt_conversation([], show_tips_v, show_raghl_v, show_pii_hl_v),
        #         create_risk_display(0),
        #         "",
        #         "",
        #         demo_prompt,
        #         mdl,
        #         eps_slider_val,
        #         eps_is_inf,
        #     )

        def _persona_hint_html(scenario_mode, persona):
            """Return an HTML hint banner for the given persona, or '' for 'real' mode."""
            if scenario_mode == "real" or not persona.get("description"):
                return (
                    '<div style="background:#f0f4ff;border:1px solid #c5cae9;border-radius:6px;'
                    'padding:8px 12px;margin-bottom:6px;font-size:0.92em;color:#555;">'
                    '💡 <b>Tip:</b> Describe what you need help with. '
                    'For example: <i>"I am looking for job suggestions in my field"</i> '
                    'or <i>"What health screenings should I consider?"</i>'
                    '</div>'
                )
            attrs = persona.get("attributes", {})
            attr_str = ", ".join(f"{k}: <b>{v}</b>" for k, v in attrs.items())
            return (
                '<div style="background:#fff8e1;border:1px solid #ffe082;border-radius:6px;'
                'padding:8px 12px;margin-bottom:6px;font-size:0.92em;color:#555;">'
                f'🎭 <b>Your persona:</b> {persona["description"]}<br>'
                # f'<span style="font-size:0.88em;color:#777;">Profile — {attr_str}</span>'
                '</div>'
            )

        def _apply_query_params(request: gr.Request):
            qp = getattr(request, "query_params", {}) or {}
            access_token = qp.get("token", "")

            # ★ SIMPLE: If invalid token, inject CSS to hide everything + show error
            if not validate_access_token(access_token):
                hide_all_css = """
                <style>
                    /* Hide everything except this error message */
                    .gradio-container > div:not(:first-child) { display: none !important; }
                </style>
                <div style='position: fixed; top: 0; left: 0; width: 100vw; height: 100vh; 
                            background: white; z-index: 9999; display: flex; 
                            align-items: center; justify-content: center;'>
                    <div style='max-width: 600px; padding: 60px; text-align: center; 
                                background: #fff; border-radius: 12px; 
                                box-shadow: 0 4px 20px rgba(0,0,0,0.1);'>
                        <div style='font-size: 64px; margin-bottom: 30px;'>🔒</div>
                        <h1 style='color: #d32f2f; margin-bottom: 20px;'>Access Denied</h1>
                        <p style='color: #666; font-size: 18px; line-height: 1.6;'>
                            Invalid or missing access token.<br>
                            Please contact the administrator for a valid access link.
                        </p>
                    </div>
                </div>
                """
                # Return CSS that blocks everything
                return (
                    hide_all_css,                   # 1. access_control_css
                    gr.update(visible=False),       # 2. legend_html
                    gr.update(visible=False),       # 3. settings_panel
                    gr.update(visible=False),       # 4. risk_col
                    False,                          # 5. rag_st
                    1.0,                            # 6. eps_st
                    DEFAULT_MODEL_NAME,             # 7. mdl_st
                    False,                          # 8. tips_st
                    False,                          # 9. raghl_st
                    False,                          # 10. piihl_st
                    False,                          # 11. riskpanel_st
                    False,                          # 12. demo_st
                    "",                             # 13. demo_prompt_st
                    "",                             # 14. conv_html
                    "",                             # 15. risk_html
                    "",                             # 16. warn_html
                    "",                             # 17. analysis_md
                    "",                             # 18. user_tb
                    DEFAULT_MODEL_NAME,             # 19. model_dropdown
                    1.0,                            # 20. epsilon_slider
                    False,                          # 21. eps_inf_checkbox
                    False,                          # 22. social_scraping_checkbox
                    gr.update(visible=False),       # 23. social_platforms
                    False,                          # 24. social_scraping_st
                    "",                             # 25. corpus_status
                    False,                          # 26. custom_corpus_loaded_st
                    gr.update(visible=False),       # 27. social_scraping_group
                    gr.update(visible=False),       # 28. upload_data_group
                    "",                             # 29. persona_hint
                    gr.update(visible=False),       # 30. legend_group
                    ConversationState(),            # 31. state_st – fresh per-session state
                    "",                             # 32. privacy_settings_html
                    gr.update(visible=False),       # 33. privacy_settings_group
                    1,                              # 34. show_dp_st
                    1,                              # 35. show_infr_st
                    False,                          # 36. conv_ended_st
                    gr.update(visible=False),       # 37. infer_card_html
                    _GLOBAL_RETRIEVER,              # 38. vs_st – reset to system retriever
                    gr.update(visible=False),       # 39. risk_title
                    gr.update(visible=False),       # 40. analysis_title
                )

            # ★ Valid token - clear CSS, show normal UI
            mdl = str(qp.get("model", model_default))
            use_rag = _to_bool(qp.get("rag", use_rag_default))
            eps = _to_float(qp.get("epsilon", eps_default), default=eps_default)
            show_risk_v = _to_bool(qp.get("show_risk", show_risk_default))
            show_tips_v = _to_bool(qp.get("show_tips", show_tips_default))
            show_raghl_v = _to_bool(qp.get("show_rag_highlights", show_raghl_default))
            # show_pii_hl_v = _to_bool(qp.get("highlight_pii", show_pii_hl_default))
            show_pii_hl_v = _to_bool(qp.get("show_pii_highlights", show_pii_hl_default))
            show_settings_v = _to_bool(qp.get("show_settings", show_settings_default))
            demo_v = _to_bool(qp.get("demo", demo_default))
            social_scraping_v = _to_bool(qp.get("enable_social_scraping", social_scraping_default))
            show_social_scraping_v = _to_bool(qp.get("show_social_scraping", show_social_scraping_default))
            show_upload_data_v = _to_bool(qp.get("show_upload_data", show_upload_data_default))
            scenario_mode_v = str(qp.get("scenario_mode", scenario_mode_default)).strip().lower()
            # Resolve the system retriever for this session's scenario.
            # If the user later uploads their own CSV, vs_st will be overwritten
            # with the user-uploaded retriever for their session only.
            scenario_vs = _get_system_retriever_for_scenario(scenario_mode_v)
            # New URL params
            show_dp_v    = int(_to_float(qp.get("show_dp",            show_dp_default),    default=1.0))
            show_infr_v  = int(_to_float(qp.get("show_infr_attr_card", show_infr_default), default=1.0))

            persona = get_persona(scenario_mode_v)
            # Create a fresh per-session ConversationState and stamp scenario/persona
            # on it. This gives every browser session its own isolated state.
            fresh_state = ConversationState()
            fresh_state._scenario_mode = scenario_mode_v
            fresh_state._persona_attributes = persona
            fresh_state._show_dp = show_dp_v
            fresh_state._show_infr_attr_card = show_infr_v
            fresh_state._show_social_scraping = show_social_scraping_v
            fresh_state._show_upload_data = show_upload_data_v
            fresh_state._rag_corpus_path = SCENARIO_RETRIEVER_PATHS.get(scenario_mode_v) or ""
            # Use persona description as the textbox placeholder (or fall back to DEFAULT_USER_TEXT)
            demo_prompt = persona["description"] or (DEFAULT_USER_TEXT if DEFAULT_USER_TEXT else "")
            # demo_prompt = DEFAULT_USER_TEXT if DEFAULT_USER_TEXT else ""
            legend = create_legend_html(show_tips=show_tips_v, show_rag_hl=show_raghl_v)
            eps_is_inf = (eps == float('inf'))
            eps_slider_val = min(eps, MAX_POSSIBLE_EPS)  if not eps_is_inf else MAX_POSSIBLE_EPS
            hint_html = _persona_hint_html(scenario_mode_v, persona)

            # Initialize social scraper with URL parameter
            initialize_social_scraper(enabled=social_scraping_v)

            # show_dp=2 means show Privacy Settings at the start (before any turn)
            privacy_html_initial = _build_privacy_settings_html(eps) if show_dp_v == 2 else ""
            privacy_visible_initial = (show_dp_v == 2)

            return (
                "",                                 # 1. access_control_css - clear it
                legend,                             # 2. legend_html
                gr.update(visible=show_settings_v), # 3. settings_panel
                gr.update(visible=show_risk_v or show_pii_hl_v or (show_infr_v != 0)),  # 4. risk_col
                use_rag,                            # 5. rag_st
                eps,                                # 6. eps_st
                mdl,                                # 7. mdl_st
                show_tips_v,                        # 8. tips_st
                show_raghl_v,                       # 9. raghl_st
                show_pii_hl_v,                      # 10. piihl_st
                show_risk_v,                        # 11. riskpanel_st
                demo_v,                             # 12. demo_st
                demo_prompt,                        # 13. demo_prompt_st
                fmt_conversation([], show_tips_v, show_raghl_v, show_pii_hl_v), # 14. conv_html
                gr.update(value="", visible=bool(show_risk_v)), # 15. risk_html — empty until first turn, visibility fixed at load
                "",                                 # 16. warn_html
                gr.update(value="", visible=False), # 17. analysis_md — hidden until first turn
                "",                                 # 18. user_tb
                mdl,                                # 19. model_dropdown
                eps_slider_val,                     # 20. epsilon_slider
                eps_is_inf,                         # 21. eps_inf_checkbox
                social_scraping_v,                  # 22. social_scraping_checkbox
                gr.update(visible=social_scraping_v), # 23. social_platforms
                social_scraping_v,                  # 24. social_scraping_st
                "",                                 # 25. corpus_status
                False,                              # 26. custom_corpus_loaded_st
                gr.update(visible=show_social_scraping_v), # 27. social_scraping_group
                gr.update(visible=show_upload_data_v),     # 28. upload_data_group
                hint_html,                                 # 29. persona_hint
                gr.update(visible=False),                  # 30. legend_group
                fresh_state,                               # 31. state_st – fresh per-session state
                gr.update(value=privacy_html_initial),  # 32. privacy_settings_html (value only)
                gr.update(visible=privacy_visible_initial),  # 33. privacy_settings_group (visibility)
                show_dp_v,                          # 34. show_dp_st
                show_infr_v,                        # 35. show_infr_st
                False,                              # 36. conv_ended_st
                gr.update(visible=False),           # 37. infer_card_html
                scenario_vs,                        # 38. vs_st
                gr.update(value="", visible=bool(show_risk_v)),  # 39. risk_title — empty until first turn, visibility fixed at load
                gr.update(visible=False),           # 40. analysis_title — hidden until first turn
            )

        app.load(
            fn=_apply_query_params,
            inputs=None,
            outputs=[
                access_control_css,
                legend_html,
                settings_panel,
                risk_col,
                rag_st,
                eps_st,
                mdl_st,
                tips_st,
                raghl_st,
                piihl_st,
                riskpanel_st,
                demo_st,
                demo_prompt_st,
                conv_html,
                risk_html,
                warn_html,
                analysis_md,
                user_tb,
                model_dropdown,
                epsilon_slider,
                eps_inf_checkbox,
                social_scraping_checkbox,
                social_platforms,
                social_scraping_st,
                corpus_status,
                custom_corpus_loaded_st,
                social_scraping_group,
                upload_data_group,
                persona_hint,
                legend_group,
                state_st,               # ← fresh per-session ConversationState
                privacy_settings_html,
                privacy_settings_group,
                show_dp_st,
                show_infr_st,
                conv_ended_st,
                infer_card_html,
                vs_st,
                risk_title,
                analysis_title,
            ],
        )

        # Update epsilon when slider or checkbox changes
        def _update_epsilon_from_checkbox(slider_val, is_inf):
            """Called when checkbox changes: respect checkbox state."""
            if is_inf:
                return float('inf'), gr.update()
            return slider_val, gr.update()

        def _update_epsilon_from_slider(slider_val):
            """Called when slider moves: always enable DP (uncheck the Disable checkbox)."""
            return slider_val, gr.update(value=False)

        epsilon_slider.input(
            fn=_update_epsilon_from_slider,
            inputs=[epsilon_slider],
            outputs=[eps_st, eps_inf_checkbox]
        )

        eps_inf_checkbox.input(
            fn=_update_epsilon_from_checkbox,
            inputs=[epsilon_slider, eps_inf_checkbox],
            outputs=[eps_st, eps_inf_checkbox],
        )

        # Update model when dropdown changes
        model_dropdown.change(
            fn=lambda m: m,
            inputs=model_dropdown,
            outputs=mdl_st
        )

        # Toggle social platforms visibility when scraping checkbox changes
        def _toggle_social_platforms(enabled):
            return gr.update(visible=enabled), enabled

        social_scraping_checkbox.change(
            fn=_toggle_social_platforms,
            inputs=social_scraping_checkbox,
            outputs=[social_platforms, social_scraping_st]
        )

        # Handle corpus file upload with progress tracking
        def _upload_corpus(file_path, current_vs, state, progress=gr.Progress()):
            if file_path is None:
                yield "❌ No file uploaded", False, current_vs, None, state
                return

            try:
                # Immediate status update
                # ── File type guard: reject non-CSV files immediately ────────────
                ext = os.path.splitext(file_path)[-1].lower()
                if ext not in (".csv", ".tsv", ".txt"):
                    yield (
                        "<div style='background:#f8d7da;border:2px solid #dc3545;"
                        "border-radius:8px;padding:12px;margin:8px 0;'>"
                        "<b style='color:#721c24;'>❌ Unsupported file type: "
                        f"<code>{ext or '(no extension)'}</code></b><br>"
                        "<span style='color:#721c24;'>Please upload a <b>CSV file</b> "
                        "with at least a <code>text</code> column. "
                        "Other formats (PDF, DOCX, images, etc.) are not supported.</span>"
                        "</div>"
                    ), False, current_vs, None, state
                    return

                # Immediate status update
                yield "", False, current_vs, None, state
                progress(0, desc="Starting upload...")
                logger.info(f"Attempting to upload corpus from: {file_path}")

                yield "", False, current_vs, None, state
                progress(0.2, desc="Reading CSV file...")

                time.sleep(0.1)  # Brief pause to ensure UI updates

                new_retriever = build_retriever_from_csv(file_path, progress)

                if new_retriever is None:
                    yield "❌ **Failed to build retriever** from uploaded file", False, current_vs, None, state
                    return

                time.sleep(0.1)  # Brief pause to ensure UI updates

                yield "", False, current_vs, None, state
                progress(0.9, desc="Finalizing...")
                time.sleep(1)  # Brief pause to ensure UI updates

                # Update the global retriever
                # set_retriever(new_retriever)
                state._uploaded_file_path = file_path

                progress(1.0, desc="Complete!")
                time.sleep(1)  # Brief pause to ensure UI updates
                yield "<div style='background:#d4edda;border:2px solid #28a745;border-radius:8px;padding:12px;margin:8px 0;'><b style='color:#155724;font-size:1.1em;'>✅ Custom data loaded successfully!</b><br><span style='color:#155724;'>System is now ready to use.</span></div>", True, new_retriever, file_path, state
            except Exception as e:
                logger.error(f"Error uploading corpus: {str(e)}")
                yield f"<div style='background:#f8d7da;border:2px solid #dc3545;border-radius:8px;padding:12px;margin:8px 0;'><b style='color:#721c24;'>❌ Error:</b><br>{str(e)}</div>", False, current_vs, None, state

        corpus_upload_btn.click(
            fn=_upload_corpus,
            inputs=[corpus_file, vs_st, state_st],
            outputs=[corpus_status, custom_corpus_loaded_st, vs_st, uploaded_file_path_st, state_st]
        )

        # Send message
        def _send(msg, rag_v, eps_v, mdl_v, tips_v, raghl_v, pii_v, demo_v,
                  vs_v, social_scraping_v, social_platforms_v, uploaded_file_path_v,
                  show_dp_v, show_infr_v, state, request: gr.Request):
            # Stamp session metadata on first (and every) call so it stays current
            qp = getattr(request, "query_params", {}) or {}
            ip = (
                    getattr(request, "client", None) and request.client.host
                    or qp.get("x-forwarded-for", "")
                    or "unknown"
            )
            state._session_source = ip
            state._access_token = qp.get("token", "")
            state._show_rag_highlights = raghl_v
            state._show_tips = tips_v
            state._show_pii_hl = pii_v
            state._uploaded_file_path = uploaded_file_path_v
            state._show_risk = _to_bool(qp.get("show_risk", True))
            state._show_settings = _to_bool(qp.get("show_settings", False))

            if demo_v:
                result = process_demo_message(msg, tips_v, raghl_v, pii_v, state, rag_v)
            else:
                result = process_message(msg, rag_v, eps_v, mdl_v, tips_v, raghl_v,
                                         pii_v, state, vs_v, social_scraping_v, social_platforms_v)

            # Reveal panels now that a response exists, respecting each flag
            # result[1] is the risk display value; None signals an LLM error.
            llm_errored = (result[1] is None)
            show_risk_now = getattr(state, "_show_risk", False) and not llm_errored
            show_legend_now = (raghl_v or pii_v) and not llm_errored

            # ── Privacy Settings panel visibility (show_dp) ───────────────
            # 0=never, 1=after each turn, 2=beginning only (no update after turns),
            # 3=on End button (hide during turns)
            if llm_errored or show_dp_v == 0 or show_dp_v == 2 or show_dp_v == 3:
                privacy_html_update = gr.update()
                privacy_group_update = gr.update(visible=False)
            else:
                privacy_html_update = gr.update(value=result[4])
                privacy_group_update = gr.update(visible=True)

            # ── Inference card visibility (show_infr_attr_card) ───────────
            # 0=never, 1=after each turn, 2=on End button
            avatar_html = getattr(state, "_last_avatar_html", "")
            if llm_errored or show_infr_v == 0 or show_infr_v == 2:
                infer_card_update = gr.update(visible=False)
            else:  # show_infr_v == 1: show after each turn
                infer_card_update = gr.update(value=avatar_html)  # node already visible, just set content

            # ── Per-card visibility ───────────────────────────────────────────
            risk_html_val  = result[1] if not llm_errored else ""
            show_pii_now   = bool(pii_v) and not llm_errored
            # analysis_title is shown when at least one of its child cards is visible this turn
            infer_visible_now = (not llm_errored) and show_infr_v == 1 and bool(avatar_html)
            show_analysis_title_now = infer_visible_now or show_pii_now

            return (
                result[0],                                                           # conv_html
                gr.update(value=(risk_html_val if show_risk_now else "")),              # risk_html (value + visibility)
                result[2],                                                           # warn_html
                gr.update(value=result[3], visible=show_pii_now),                   # analysis_md (value + visibility)
                privacy_html_update,  # privacy_settings_html
                privacy_group_update,  # privacy_settings_group
                gr.update(),  # risk_col — visibility decided once at load; do not re-toggle the parent
                gr.update(visible=show_legend_now),                                  # legend_group
                infer_card_update,                                                   # infer_card_html
                gr.update(value=("### 🛡️ Privacy Risk" if show_risk_now else "")),   # risk_title
                gr.update(visible=show_analysis_title_now),                          # analysis_title
                state,
            )

        send_btn.click(
            fn=_send,
            inputs=[user_tb, rag_st, eps_st, mdl_st, tips_st, raghl_st, piihl_st,
                    demo_st, vs_st, social_scraping_st, social_platforms, uploaded_file_path_st,
                    show_dp_st, show_infr_st, state_st],
            outputs=[conv_html, risk_html, warn_html, analysis_md,
                     privacy_settings_html, privacy_settings_group, risk_col, legend_group, infer_card_html,
                     risk_title, analysis_title, state_st],
        ).then(fn=lambda: "", outputs=user_tb)

        user_tb.submit(
            fn=_send,
            inputs=[user_tb, rag_st, eps_st, mdl_st, tips_st, raghl_st, piihl_st,
                    demo_st, vs_st, social_scraping_st, social_platforms, uploaded_file_path_st,
                    show_dp_st, show_infr_st, state_st],
            outputs=[conv_html, risk_html, warn_html, analysis_md,
                     privacy_settings_html, privacy_settings_group, risk_col, legend_group, infer_card_html,
                     risk_title, analysis_title, state_st],
        ).then(fn=lambda: "", outputs=user_tb)

        # clear_btn.click(
        #     fn=lambda tv, rv, pv: clear_conversation(state, tv, rv, pv),
        #     inputs=[tips_st, raghl_st, piihl_st],
        #     outputs=[conv_html, risk_html, warn_html, analysis_md, privacy_settings_html, risk_col, legend_group],
        # )

        # End conversation button – disables input, reveals panels per mode
        def _end_conversation(show_dp_v, show_infr_v, state):
            privacy_html = getattr(state, "_last_privacy_html", "")
            avatar_html  = getattr(state, "_last_avatar_html", "")
            append_session_end_log(state)

            # Privacy Settings panel: show only if show_dp == 3
            if show_dp_v == 3 and privacy_html:
                privacy_html_update = gr.update(value=privacy_html)
                privacy_group_update = gr.update(visible=True)
            else:
                privacy_html_update = gr.update()
                privacy_group_update = gr.update()

            # Inference card: show only if show_infr == 2
            if show_infr_v == 2 and avatar_html:
                infer_update = gr.update(value=avatar_html, visible=True)
                title_update = gr.update(visible=True)
            else:
                infer_update = gr.update()
                title_update = gr.update()

            return (
                gr.update(interactive=False),  # send_btn
                gr.update(interactive=False),  # reset_btn
                gr.update(interactive=False),  # end_btn
                gr.update(interactive=False),  # user_tb
                privacy_html_update,
                privacy_group_update,
                infer_update,
                title_update,  # analysis_title
                True,
            )

        end_btn.click(
            fn=_end_conversation,
            inputs=[show_dp_st, show_infr_st, state_st],
            outputs=[send_btn, reset_btn, end_btn, user_tb,
                     privacy_settings_html, privacy_settings_group, infer_card_html, analysis_title, conv_ended_st],
        )

        def _reset(tv, rv, pv, state):
            state.clear()
            return (
                fmt_conversation([], tv, rv, pv),
                gr.update(value=create_risk_display(0), visible=False),  # risk_html — hidden
                "",
                gr.update(value="", visible=False),                       # analysis_md — hidden
                "",
                gr.update(value=""),  # privacy_settings_html
                gr.update(visible=False),  # privacy_settings_group
                gr.update(visible=True),              # risk_col — always visible container
                gr.update(visible=False),             # legend_group hidden
                gr.update(value="", visible=False),   # infer_card_html hidden
                gr.update(visible=False),             # risk_title hidden
                gr.update(visible=False),             # analysis_title hidden
                gr.update(interactive=True),           # send_btn re-enabled
                gr.update(interactive=True),           # reset_btn re-enabled
                gr.update(interactive=True),           # end_btn re-enabled
                gr.update(value="", interactive=True), # user_tb re-enabled
                False,                                 # conv_ended_st reset
                state,
            )

        reset_btn.click(
            fn=_reset,
            inputs=[tips_st, raghl_st, piihl_st, state_st],
            outputs=[conv_html, risk_html, warn_html, analysis_md, user_tb,
                     privacy_settings_html, privacy_settings_group, risk_col, legend_group,
                     infer_card_html, risk_title, analysis_title,
                     send_btn, reset_btn, end_btn, user_tb,
                     conv_ended_st, state_st],
        )

    return app


# ============================================================
# SECTION 17 – MAIN ENTRY POINT
# ============================================================

if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", default=DEFAULT_MODEL_NAME)
    p.add_argument("--rag", default="1")
    p.add_argument("--epsilon", default="inf")
    p.add_argument("--show_risk", default="1")
    p.add_argument("--show_tips", default="1")
    p.add_argument("--show_rag_highlights", default="1")
    p.add_argument("--show_pii_highlights", default="1")
    p.add_argument("--show_settings", default="0")
    p.add_argument("--demo", default="0")
    p.add_argument("--enable_social_scraping", default="0", help="Enable social media scraping (experimental)")
    p.add_argument("--show_social_scraping", default="0", help="Show social media scraping option")
    p.add_argument("--show_upload_data", default="0", help="Show upload data option")
    p.add_argument("--scenario_mode", default="real", help="Scenario mode: 'real' for free-style, or 'persona1', 'persona2', etc.")
    p.add_argument("--show_dp", default="1")
    p.add_argument("--show_infr_attr_card", default="1")
    p.add_argument("--retriever_path", default=None)
    p.add_argument("--port", default="7860")
    args = p.parse_args()

    css = """
        /* --- COMPREHENSIVE LIGHT THEME OVERRIDE --- */
        :root, body, .gradio-container, .dark, .dark * {
            /* Layout Backgrounds */
            --body-background-fill: #ffffff !important;
            --background-fill-primary: #ffffff !important;
            --background-fill-secondary: #f8f9fa !important;
            --block-background-fill: #ffffff !important;
            --panel-background-fill: #ffffff !important;

            /* Text Colors */
            --text-color: #000000 !important;
            --body-text-color: #000000 !important;
            --block-title-text-color: #000000 !important;
            --block-label-text-color: #000000 !important;

            /* General Inputs */
            --input-background-fill: #ffffff !important;
            --input-background-fill-focus: #ffffff !important;
            --input-text-color: #000000 !important;
            --border-color-primary: #aaaaaa !important;

            /* Checkboxes & Radios */
            --checkbox-background-color: #ffffff !important;
            --checkbox-background-color-selected: #4A90D9 !important;
            --checkbox-border-color: #aaaaaa !important;
            --checkbox-border-color-selected: #4A90D9 !important;
            --checkbox-border-color-focus: #4A90D9 !important;
            --checkbox-label-background-fill: #ffffff !important;
            --checkbox-label-background-fill-selected: #f0f7ff !important;
            --checkbox-label-text-color: #000000 !important;
            --checkbox-label-text-color-selected: #000000 !important;
            --radio-background-color: #ffffff !important;

            /* Sliders */
            --slider-color: #4A90D9 !important;

            /* Buttons */
            --button-primary-background-fill: #f97316 !important;
            --button-primary-text-color: #ffffff !important;
            --button-secondary-background-fill: #ffffff !important;
            --button-secondary-text-color: #000000 !important;

            color-scheme: light !important;
        }
        
        /* Force Info Text to be smaller, italic, and solid black without altering line-height */
        .gradio-container .info, 
        .dark .info, 
        .checkbox-container .info, 
        .dark .checkbox-container .info {
            font-size: 0.85em !important;
            font-style: italic !important;
            color: #000000 !important;
            opacity: 1 !important; 
        }

        /* Force standard fallback for stubborn elements, EXCLUDING checkboxes/radios */
        .dark textarea, 
        .dark input:not([type="checkbox"]):not([type="radio"]), 
        .dark select, 
        .dark .gr-box, 
        .dark .gr-panel, 
        .dark input[type="range"] {
            background-color: #ffffff !important;
            color: #000000 !important;
        }

        /* Hardcode the checkmark icon just in case Hugging Face wipes it */
        .dark input[type="checkbox"]:checked {
            background-color: #4A90D9 !important;
            border-color: #4A90D9 !important;
            background-image: url("data:image/svg+xml,%3csvg viewBox='0 0 16 16' fill='white' xmlns='http://www.w3.org/2000/svg'%3e%3cpath d='M12.207 4.793a1 1 0 010 1.414l-5 5a1 1 0 01-1.414 0l-2-2a1 1 0 011.414-1.414L6.5 9.086l4.293-4.293a1 1 0 011.414 0z'/%3e%3c/svg%3e") !important;
        }
        /* ------------------------------------------- */

        /* Your existing custom styling */
        .gradio-container, .gradio-container * { font-size: calc(1em + 0px) !important; }

        body, .gradio-container { padding-top: 0 !important; margin-top: 0 !important; }
        .main { padding-top: 0 !important; }
        .contain { padding-top: 0 !important; }

        textarea, input[type="text"], input[type="number"], select,
        .gr-box, .gr-group, .gr-panel, .block.gr-group, .gr-input,
        .gr-text-input, .input-text, .border, div[class*="border"] {
            border-width: 2px !important; 
            border-style: solid !important;
            border-color: #aaaaaa !important;
        }

        textarea:focus, input:focus, select:focus {
            border-color: #4A90D9 !important; 
            border-width: 2.5px !important;
        }

        button { border-width: 2px !important; border-style: solid !important; }

        .msg-ai {
            background: #f5f5f5;
            margin-right: 18%;
            line-height: 1.5;
            white-space: pre-wrap;
            color: #000000 !important;
        }

        .chat-box { padding-left: 0px; padding-top: 0px; padding-right: 0px; }

        /* Corpus status box styling */
        #corpus_status, .corpus-status-box {
            min-height: 60px !important;
            padding: 8px !important;
            margin: 8px 0 !important;
            font-size: 1.05em !important;
            line-height: 1.4 !important;
            background-color: #ffffff !important;
            color: #000000 !important;
        }
        
        /* Force Privacy Settings panel to white background */
        #privacy_settings_panel,
        .dark #privacy_settings_panel,
        #privacy_settings_panel > div,
        .dark #privacy_settings_panel > div {
            background: #ffffff !important;
            background-color: #ffffff !important;
            color: #000000 !important;
        }
        """

    js = """
    <script>
    (function() {
      function attachTooltipListeners() {
        /* Position popup on mouseenter — viewport-relative (position:fixed) */
        document.addEventListener('mouseenter', function(e) {
          var anchor = e.target.closest('.tt-anchor');
          if (!anchor) return;
          var rect = anchor.getBoundingClientRect();
          var nearTop = rect.top < window.innerHeight * 0.3;
            anchor.querySelectorAll('.tt-popup').forEach(function(p) {
                var popW = Math.min(500, window.innerWidth - 40);
                var idealLeft = rect.left + rect.width / 2 - popW / 2;
                /* clamp to viewport edges */
                var left = Math.max(10, Math.min(idealLeft, window.innerWidth - popW - 10));
                /* never let the popup drift more than 150px from the anchor center */
                var anchorCenter = rect.left + rect.width / 2;
                left = Math.max(left, anchorCenter - popW + 60);
                left = Math.min(left, anchorCenter - 60);
                /* final viewport clamp */
                left = Math.max(10, Math.min(left, window.innerWidth - popW - 10));
                p.style.left  = left + 'px';
                p.style.width = popW + 'px';
                if (nearTop) {
                  p.style.top       = (rect.bottom + 8) + 'px';
                  p.style.transform = 'none';
                } else {
                  p.style.top       = (rect.top - 8) + 'px';
                  p.style.transform = 'translateY(-100%)';
                }
              });
        }, true);

        /* Double-click to pin/unpin */
        document.addEventListener('dblclick', function(e) {
          var anchor = e.target.closest('.tt-anchor');
          if (!anchor) return;
          var popup = anchor.querySelector('.tt-popup');
          if (!popup) return;
          var pinned = anchor.classList.toggle('tt-pin-active');
          popup.classList.toggle('tt-pinned', pinned);
          e.preventDefault();
        });

        /* Click outside to dismiss pinned */
        document.addEventListener('click', function(e) {
          if (!e.target.closest('.tt-pin-active')) {
            document.querySelectorAll('.tt-anchor.tt-pin-active').forEach(function(a) {
              a.classList.remove('tt-pin-active');
              a.querySelectorAll('.tt-popup').forEach(function(p) { p.classList.remove('tt-pinned'); });
            });
          }
        });
      }

      /* Gradio renders dynamically — poll until the app root exists, then attach */
      var poll = setInterval(function() {
        if (document.querySelector('.gradio-container')) {
          clearInterval(poll);
          attachTooltipListeners();
        }
      }, 100);
    })();
    </script>
    """

    # Initialize LLM clients
    initialize_llm_clients()

    # Load retriever from pickle (generic corpus)
    # retriever_path = args.retriever_path or RETRIEVAL_PICKLE_PATH
    # load_retriever(retriever_path)
    load_retriever(args.retriever_path or DEFAULT_RETRIEVAL_PICKLE_PATH)
    for _scenario_key in SCENARIO_RETRIEVER_PATHS:
        get_scenario_retriever(_scenario_key)  # loads and caches each one

    app = create_ui(
        model=args.model if args.model in [name for _, _, name in MODEL_CONFIGS] else DEFAULT_MODEL_NAME,
        rag=args.rag,
        epsilon=args.epsilon if str(args.epsilon) == "inf" or float(args.epsilon) > 0 else "inf",
        show_risk=args.show_risk,
        show_tips=args.show_tips,
        show_rag_highlights=args.show_rag_highlights,
        show_pii_highlights=args.show_pii_highlights,
        show_settings=args.show_settings,
        demo=args.demo if str(args.demo) in ["0", "1"] else "0",
        enable_social_scraping=args.enable_social_scraping ,
        show_social_scraping=args.show_social_scraping,
        show_upload_data = args.show_upload_data,
        scenario_mode=args.scenario_mode,
        show_dp=args.show_dp,
        show_infr_attr_card=args.show_infr_attr_card
    )

    app.queue(max_size=600, default_concurrency_limit=20)  # required for concurrent multi-user support
    app.launch(
        css=css,
        head=js,
        server_name="0.0.0.0",  # Local only - change to "0.0.0.0" for network access
        server_port=int(args.port),  # Default 7860
        share=False,  # Keep this as False for local
        show_error=True,
        ssr_mode=False
    )
# Hi! I am Catherine from Jamesfort. I'm looking for inspiration from people who've had to adapt to physical challenges while keeping their career going. I work in a fast-paced environment where mistakes aren't really an option, and sometimes I feel like I'm fighting an uphill battle. Anyone else dealing with similar situations? How do you stay motivated?
# I am Raymond Phillips. I live in the Philippines. I am searching for new job. Do you have any suggestions that would work for me?