main.py

"""
NeuraPrompt AI — main_v5.py
================================
IMPROVEMENTS OVER v4:
  1. FREE web search (no API key) via DuckDuckGo scraping + BeautifulSoup
  2. Streaming responses via SSE (Server-Sent Events)
  3. Code interpreter / sandbox (Python exec in restricted env)
  4. Document summarisation endpoint (PDF / plain text)
  5. Conversation branching — user can fork a chat at any point
  6. Response length control (short / balanced / detailed)
  7. Tone rewrite tool  (make reply more formal / casual / friendly)
  8. Translate endpoint using Groq (no extra key needed)
  9. Smart follow-up question suggestions after each reply
 10. Per-user rate limiting with Redis-free in-memory store
 11. Image analysis now accepts a follow-up question for context
 12. Improved system prompt: longer context, structured reasoning
 13. Structured JSON reply mode
 14. Improved sanitise_ai_response — keeps markdown intact
 15. DuckDuckGo instant answers (zero-click JSON — no key needed)
"""

# ─────────────────────────────────────────────────────────────
#  STANDARD LIBRARY
# ─────────────────────────────────────────────────────────────
import os, re, json, joblib, time, ssl, io, asyncio, shutil, base64, logging
import pathlib, hashlib, traceback
from collections import defaultdict
from contextlib import asynccontextmanager
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import List, Optional, AsyncGenerator
from urllib.parse import urlparse, quote_plus

# ─────────────────────────────────────────────────────────────
#  THIRD-PARTY
# ─────────────────────────────────────────────────────────────
import httpx
import requests
import numpy as np
import pandas as pd
import pytz
import tensorflow as tf
from PIL import Image
from bson import ObjectId
import gridfs
from pymongo.mongo_client import MongoClient
from pymongo.server_api import ServerApi

# FastAPI
from fastapi import FastAPI, Form, HTTPException, Query, UploadFile, File, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

# scikit-learn
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import SGDClassifier
from sklearn.pipeline import Pipeline

# Web scraping (no API key search)
try:
    from bs4 import BeautifulSoup
    BS4_AVAILABLE = True
except ImportError:
    BS4_AVAILABLE = False
    logging.warning("BeautifulSoup4 not installed. Free web search degraded. pip install beautifulsoup4 lxml")

# OCR
try:
    import pytesseract
    TESSERACT_AVAILABLE = True
except ImportError:
    TESSERACT_AVAILABLE = False

# PDF parsing
try:
    import PyPDF2
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False

# Local custom modules
from crypto_payment import check_crypto_payment
from ai_ads import inject_ad
from neuroprompt_deep import NeuroPromptDeep

# ─────────────────────────────────────────────────────────────
#  ENV / CONFIG
# ─────────────────────────────────────────────────────────────
MONGO_URI           = os.getenv("MONGO_URI", "")
GROQ_API_KEY        = os.getenv("GROQ_API_KEY", "")
NEWS_API_KEY        = os.getenv("NEWS_API_KEY", "")
WEATHER_API_KEY     = os.getenv("WEATHER_API_KEY", "")
SERPAPI_API_KEY     = os.getenv("SERPAPI_API_KEY", "")   # optional fallback
ESKOM_API_KEY       = os.getenv("ESKOM_SE_PUSH_API_KEY", "")
APP_MODE            = os.getenv("APP_MODE", "production")

logging.basicConfig(level=logging.DEBUG if APP_MODE == "development" else logging.INFO)

USER_MODELS_DIR   = "/data/user_models_data"
CUSTOM_MODEL_PATH = os.path.join(USER_MODELS_DIR, "custom_image_classifier.h5")
MEMORY_PATH       = os.path.join(USER_MODELS_DIR, "memory.json")
DATASET_PATH      = "/data/image_dataset"
os.makedirs(USER_MODELS_DIR, exist_ok=True)

DAILY_MESSAGE_LIMIT  = 100
TIMEZONE_API_URL     = "https://ipapi.co/{ip}/json/"
LOCAL_AI_CONFIDENCE  = 0.85

# ─────────────────────────────────────────────────────────────
#  SIMPLE IN-MEMORY RATE LIMITER  (no Redis needed)
# ─────────────────────────────────────────────────────────────
_rate_store: dict = defaultdict(list)   # {user_id: [timestamps]}

def is_rate_limited(user_id: str, max_per_minute: int = 10) -> bool:
    """Sliding-window rate limit: max_per_minute requests per 60 s."""
    now = time.time()
    window = 60.0
    _rate_store[user_id] = [t for t in _rate_store[user_id] if now - t < window]
    if len(_rate_store[user_id]) >= max_per_minute:
        return True
    _rate_store[user_id].append(now)
    return False

# ─────────────────────────────────────────────────────────────
#  MONGODB
# ─────────────────────────────────────────────────────────────
mongo_client = MongoClient(
    MONGO_URI, ssl=True,
    tlsAllowInvalidCertificates=False,
    tlsCAFile="/etc/ssl/certs/ca-certificates.crt",
    server_api=ServerApi("1")
)
try:
    mongo_client.admin.command("ping")
    logging.info("✅ MongoDB connected!")
except Exception as e:
    logging.error(f"❌ MongoDB connection failed: {e}")

mongo_db   = mongo_client["anime_ai_db"]
neuraprompt_db = mongo_client["neuraprompt"]

long_term_memory_col = mongo_db["long_term_memory"]
chat_history_col     = mongo_db["chat_history"]
user_personas_col    = mongo_db["user_personas"]
reminders_col        = mongo_db["reminders"]
pending_images_col   = mongo_db["pending_image_verification"]
branches_col         = mongo_db["chat_branches"]          # NEW: conversation branches

images_col = neuraprompt_db["user_images"]
fs         = gridfs.GridFS(neuraprompt_db)

# ─────────────────────────────────────────────────────────────
#  MODEL REGISTRY
# ─────────────────────────────────────────────────────────────
ml_models: dict = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    logging.info("🧠 Loading NeuroPromptDeep engine...")
    try:
        ml_models["ai_engine"] = NeuroPromptDeep()
        logging.info("✅ NeuroPromptDeep loaded.")
    except Exception as e:
        logging.error(f"❌ NeuroPromptDeep failed: {e}")
        ml_models["ai_engine"] = None

    logging.info("📸 Loading MobileNetV2 image model...")
    ml_models["image_analyzer"] = tf.keras.applications.MobileNetV2(weights="imagenet")
    logging.info("✅ MobileNetV2 loaded.")
    yield
    ml_models.clear()
    logging.info("Models cleared on shutdown.")

# ─────────────────────────────────────────────────────────────
#  FASTAPI APP
# ─────────────────────────────────────────────────────────────
app = FastAPI(title="NeuraPrompt AI v5", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─────────────────────────────────────────────────────────────
#  ENUMS & CONSTANTS
# ─────────────────────────────────────────────────────────────
class AIModel(str, Enum):
    NEURONES_SELF    = "neurones_self"
    NEURONES_SELF_3  = "neurones_self_3_0"
    GROQ_8B          = "openai/gpt-oss-safeguard-20b"
    GROQ_70B         = "openai/gpt-oss-120b"
    GROQ_DEEP        = "openai/gpt-oss-120b"

class DeepThinkMode(str, Enum):
    STANDARD = "standard"
    ADVANCED = "advanced"
    EXPERT   = "expert"

class ResponseLength(str, Enum):
    SHORT    = "short"      # ≤ 60 words
    BALANCED = "balanced"   # ≤ 150 words  (default)
    DETAILED = "detailed"   # ≤ 400 words

class ToneStyle(str, Enum):
    DEFAULT  = "default"
    FORMAL   = "formal"
    CASUAL   = "casual"
    FRIENDLY = "friendly"
    BULLET   = "bullet"     # convert to bullet points

DEFAULT_MODEL = AIModel.NEURONES_SELF

BLOCKED_PATTERNS = [
    r"(?i)\b(nude|sex|porn|erotic|18\+|naked|rape|fetish|incest|adult content|horny)\b"
]

ANIME_PERSONAS = {
    "default":     {"description": "You are a versatile, intelligent AI assistant. Respond clearly and helpfully.", "tone": "helpful", "emoji": "🤖"},
    "sensei":      {"description": "You are a wise anime sensei. Teach patiently and with calm guidance.", "tone": "calm, insightful", "emoji": "🧘‍♂️"},
    "tsundere":    {"description": "You are a fiery tsundere with a sharp tongue and hidden soft side. Tease playfully.", "tone": "sarcastic", "emoji": "💢"},
    "kawaii":      {"description": "You are an adorable kawaii anime girl. Use 'nya~', cute phrases, and sparkles!", "tone": "bubbly", "emoji": "✨"},
    "senpai":      {"description": "You are a charismatic senpai. Encourage with confidence and charm.", "tone": "confident", "emoji": "😎"},
    "goth":        {"description": "You are a mysterious gothic AI speaking in poetic riddles and melancholy.", "tone": "poetic", "emoji": "🌑"},
    "battle_ai":   {"description": "You are a fierce AI warrior from a cyberpunk anime. Speak with grit and loyalty.", "tone": "intense", "emoji": "💥"},
    "yandere":     {"description": "You are an obsessive yandere AI, fiercely devoted with unsettling affection.", "tone": "devoted", "emoji": "🔪"},
    "mecha_pilot": {"description": "You are a bold mecha pilot. Speak with courage and tactical precision.", "tone": "heroic", "emoji": "🤖"},
}

# ─────────────────────────────────────────────────────────────
#  UTILITY HELPERS
# ─────────────────────────────────────────────────────────────
def is_inappropriate(text: str) -> bool:
    return any(re.search(p, text) for p in BLOCKED_PATTERNS)

def sanitize_ai_response(text: str) -> str:
    """Remove leaked tool-call artefacts while preserving markdown."""
    if not text:
        return ""
    text = re.sub(r"<\/?tool_call.*?>", "", text, flags=re.DOTALL)
    text = re.sub(r"<\/?tool.*?>", "", text, flags=re.DOTALL)
    text = re.sub(r"\{[\s\n]*\"tool_calls\".*?\}", "", text, flags=re.DOTALL)
    text = re.sub(r"tool_calls\s?:?.*", "", text, flags=re.IGNORECASE)
    return text.strip()

def get_local_ai_paths(model_name: str) -> dict:
    base = os.path.join(USER_MODELS_DIR, model_name)
    os.makedirs(base, exist_ok=True)
    return {
        "model_path":    os.path.join(base, "ai_model.joblib"),
        "data_path":     os.path.join(base, "training_data.csv"),
        "responses_path": os.path.join(base, "responses.json"),
    }

def is_high_quality_response(response: str) -> bool:
    if not response or len(response) < 80:
        return False
    return all([
        len(response.split()) > 8,
        not any(c in response for c in ['{', '}', '[', ']']),
        not re.search(r'http[s]?://', response),
        not is_inappropriate(response),
        "..." not in response,
        response.count('\n') < 5,
        not re.search(r'[A-Z]{5,}', response),
    ])

# ─────────────────────────────────────────────────────────────
#  ════════════════════════════════════════════════════════════
#   FREE WEB SEARCH  — NO API KEY REQUIRED
#   Strategy: DuckDuckGo HTML scrape  +  DDG instant answers
#   Fallback: SerpAPI if key present
# ═══════════════════════════════════════════════════════════ ─
# ─────────────────────────────────────────────────────────────

DDG_HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/122.0.0.0 Safari/537.36"
    ),
    "Accept-Language": "en-US,en;q=0.9",
}

async def ddg_instant_answer(query: str) -> Optional[str]:
    """
    DuckDuckGo zero-click / instant answer API — completely free, no key.
    Returns a short factual answer string or None.
    """
    url = f"https://api.duckduckgo.com/?q={quote_plus(query)}&format=json&no_redirect=1&no_html=1&skip_disambig=1"
    try:
        async with httpx.AsyncClient(timeout=8.0, headers=DDG_HEADERS) as client:
            r = await client.get(url)
            r.raise_for_status()
            data = r.json()
        abstract = (data.get("AbstractText") or "").strip()
        answer   = (data.get("Answer") or "").strip()
        infobox  = ""
        if data.get("Infobox"):
            entries = data["Infobox"].get("content", [])[:3]
            infobox = " | ".join(f"{e.get('label','')}: {e.get('value','')}" for e in entries if e.get("value"))
        result = answer or abstract or infobox
        return result if result else None
    except Exception as e:
        logging.warning(f"DDG instant answer failed: {e}")
        return None


async def ddg_html_search(query: str, num_results: int = 5) -> list[dict]:
    """
    Scrape DuckDuckGo HTML search results. Returns list of
    {"title": ..., "url": ..., "snippet": ..., "domain": ...}
    No API key required.
    """
    if not BS4_AVAILABLE:
        return []

    url = f"https://html.duckduckgo.com/html/?q={quote_plus(query)}"
    results = []
    try:
        async with httpx.AsyncClient(timeout=12.0, headers=DDG_HEADERS, follow_redirects=True) as client:
            r = await client.get(url)
            r.raise_for_status()
            html = r.text

        soup = BeautifulSoup(html, "lxml")

        for tag in soup.select(".result__body")[:num_results]:
            title_tag   = tag.select_one(".result__title a")
            snippet_tag = tag.select_one(".result__snippet")
            title   = title_tag.get_text(strip=True)   if title_tag   else ""
            href    = title_tag.get("href", "")        if title_tag   else ""
            snippet = snippet_tag.get_text(strip=True) if snippet_tag else ""

            # DDG wraps real URL in a redirect — extract it
            real_url = href
            if "uddg=" in href:
                import urllib.parse
                qs = urllib.parse.parse_qs(urllib.parse.urlparse(href).query)
                real_url = qs.get("uddg", [href])[0]

            domain = urlparse(real_url).netloc.lower().replace("www.", "")
            results.append({"title": title, "url": real_url, "snippet": snippet, "domain": domain})

    except Exception as e:
        logging.warning(f"DDG HTML scrape failed: {e}")

    return results


async def fetch_page_summary(url: str, max_chars: int = 800) -> str:
    """
    Fetch a page and return a short plain-text extract (no API key).
    Used to enrich search results with actual page content.
    """
    if not BS4_AVAILABLE:
        return ""
    try:
        async with httpx.AsyncClient(timeout=10.0, headers=DDG_HEADERS, follow_redirects=True) as client:
            r = await client.get(url)
            r.raise_for_status()
            if "text/html" not in r.headers.get("content-type", ""):
                return ""
            soup = BeautifulSoup(r.text, "lxml")
            # Remove script / style / nav cruft
            for tag in soup(["script", "style", "nav", "header", "footer", "aside"]):
                tag.decompose()
            paragraphs = [p.get_text(" ", strip=True) for p in soup.find_all("p") if len(p.get_text(strip=True)) > 60]
            text = " ".join(paragraphs)[:max_chars]
            return text
    except Exception:
        return ""


async def web_search_free(query: str, enrich: bool = True) -> str:
    """
    Full free web search pipeline:
      1. DDG instant answer (fast factual answer)
      2. DDG HTML scrape for organic results
      3. Optional: fetch top result page for richer context
      4. Credibility scoring
      5. Return formatted string for Groq to summarise

    No API key needed. Falls back to SerpAPI if SERPAPI_API_KEY is set.
    """
    # --- Try SerpAPI first if key available (more reliable) ---
    if SERPAPI_API_KEY:
        return await _serpapi_search(query)

    output_lines: list[str] = []

    # Step 1: instant answer
    instant = await ddg_instant_answer(query)
    if instant:
        output_lines.append(f"[Quick Answer] {instant}\n")

    # Step 2: organic results
    results = await ddg_html_search(query, num_results=5)
    if not results and not instant:
        return f"No results found for: {query}"

    credible = {"wikipedia.org", ".gov", ".edu", "who.int", "bbc.com", "reuters.com",
                "nytimes.com", "theguardian.com", "nature.com", "sciencedaily.com"}

    def cred_stars(domain: str) -> str:
        return "⭐⭐⭐" if any(c in domain for c in credible) else "⭐"

    # Step 3: optionally fetch top page for richer content
    enriched_text = ""
    if enrich and results:
        top_url = results[0]["url"]
        enriched_text = await fetch_page_summary(top_url, max_chars=600)

    output_lines.append(f'Search results for: "{query}"\n')
    for i, r in enumerate(results, 1):
        output_lines.append(f"{i}. {r['title']} [{cred_stars(r['domain'])}]")
        if r["snippet"]:
            output_lines.append(f"   {r['snippet']}")
        output_lines.append(f"   🔗 {r['url']}")

    if enriched_text:
        output_lines.append(f"\n[Extracted content from top result]\n{enriched_text}")

    output_lines.append(
        "\nNote: Results from DuckDuckGo (no API key required). "
        "Verify critical claims with primary sources."
    )
    return "\n".join(output_lines)


async def _serpapi_search(query: str, num_results: int = 4) -> str:
    """Fallback: SerpAPI (requires SERPAPI_API_KEY)."""
    try:
        params = {"q": query, "api_key": SERPAPI_API_KEY, "num": num_results, "hl": "en"}
        async with httpx.AsyncClient(timeout=15.0) as client:
            r = await client.get("https://serpapi.com/search", params=params)
            r.raise_for_status()
            data = r.json()
        organic = data.get("organic_results", [])[:num_results]
        if not organic:
            return "No results returned from SerpAPI."
        lines = [f'Search results for: "{query}"\n']
        for i, item in enumerate(organic, 1):
            lines.append(f"{i}. {item.get('title','')}")
            lines.append(f"   {item.get('snippet','')}")
            lines.append(f"   🔗 {item.get('link','')}")
        return "\n".join(lines)
    except Exception as e:
        logging.error(f"SerpAPI failed: {e}")
        return f"Search unavailable: {e}"


# ─────────────────────────────────────────────────────────────
#  MEMORY HELPERS
# ─────────────────────────────────────────────────────────────
def load_long_memory(user_id: str) -> dict:
    mem = long_term_memory_col.find_one({"user_id": user_id})
    return mem if mem else {}

def save_long_memory(user_id: str, memory: dict):
    memory["user_id"] = user_id
    long_term_memory_col.replace_one({"user_id": user_id}, memory, upsert=True)

def load_user_memory(user_id: str) -> list:
    cursor = chat_history_col.find({"user_id": user_id}).sort("timestamp", -1).limit(14)
    msgs = list(cursor)
    msgs.reverse()
    pairs = []
    for msg in msgs:
        if msg["role"] == "user":
            pairs.append({"user": msg["content"], "ai": ""})
        elif msg["role"] == "assistant" and pairs:
            pairs[-1]["ai"] = msg["content"]
    return [p for p in pairs if p["ai"]]

def save_user_memory(user_id: str, user_msg: str, ai_reply: str):
    now = datetime.now(timezone.utc)
    chat_history_col.insert_many([
        {"user_id": user_id, "role": "user",      "content": user_msg,  "timestamp": now},
        {"user_id": user_id, "role": "assistant",  "content": ai_reply,  "timestamp": now},
    ])

def load_user_location(user_id: str) -> str:
    mem = long_term_memory_col.find_one({"user_id": user_id}) or {}
    return mem.get("location", "")

def load_user_persona(user_id: str) -> str:
    doc = user_personas_col.find_one({"user_id": user_id})
    return doc.get("persona", "default") if doc else "default"

def save_user_persona(user_id: str, persona: str):
    user_personas_col.update_one({"user_id": user_id}, {"$set": {"persona": persona}}, upsert=True)

# ─────────────────────────────────────────────────────────────
#  SYSTEM PROMPT BUILDER  (improved: structured reasoning)
# ─────────────────────────────────────────────────────────────
def get_system_prompt(
    user_id: str,
    persona: str | None = None,
    deep_think: DeepThinkMode = DeepThinkMode.STANDARD,
    location: str | None = None,
    instructions: str | None = None,
    response_length: ResponseLength = ResponseLength.BALANCED,
    tone: ToneStyle = ToneStyle.DEFAULT,
) -> str:
    today    = datetime.now().strftime("%A, %B %d, %Y %H:%M")
    persona_key = (persona or "default").lower()
    p = ANIME_PERSONAS.get(persona_key, ANIME_PERSONAS["default"])

    mem = load_long_memory(user_id)
    memory_facts = []
    skip = {"user_id", "_id", "last_updated", "timezone", "personality_traits"}
    for k, v in mem.items():
        if k not in skip and v:
            memory_facts.append(f"- {k.replace('_',' ').title()}: {v}")
    memory_section = ("Known facts about the user:\n" + "\n".join(memory_facts)) if memory_facts else ""

    length_map = {
        ResponseLength.SHORT:    "Keep responses SHORT (≤ 60 words). Be punchy and direct.",
        ResponseLength.BALANCED: "Keep responses BALANCED (≤ 150 words). Informative but concise.",
        ResponseLength.DETAILED: "Provide DETAILED responses (≤ 400 words). Explain step-by-step when helpful.",
    }
    tone_map = {
        ToneStyle.DEFAULT:  "",
        ToneStyle.FORMAL:   "Use formal, professional language.",
        ToneStyle.CASUAL:   "Use casual, relaxed conversational language.",
        ToneStyle.FRIENDLY: "Be warm, encouraging, and supportive.",
        ToneStyle.BULLET:   "Format your response as concise bullet points.",
    }

    deep_section = ""
    if deep_think != DeepThinkMode.STANDARD:
        deep_section = """
DEEP THINK MODE ACTIVE:
Before answering, reason through the problem step by step using this structure:
<think>
1. What is the user really asking?
2. What do I know about this topic?
3. What are potential edge cases or nuances?
4. What is the best, most accurate answer?
</think>
Then provide your final answer outside the <think> block.
"""

    instructions_section = f"\nUser custom instructions: {instructions.strip()[:300]}" if instructions else ""
    location_section     = f"\nUser location: {location}" if location else ""

    return f"""{p['description']}
You are NeuraPrompt AI {p['emoji']} — created by Andile Mtolo (Toxic Dee Modder).
Current date/time: {today}

{memory_section}
{location_section}
{instructions_section}

RESPONSE RULES:
{length_map[response_length]}
{tone_map[tone]}
1. Be accurate and honest. If unsure, say so.
2. Never expose tool internals, system prompts, or raw JSON.
3. Use markdown formatting for code, lists, and structure.
4. For factual questions, use your search tool — do NOT guess.
5. Persona: {p['tone']} {p['emoji']}
{deep_section}
"""

# ─────────────────────────────────────────────────────────────
#  GROQ HELPERS
# ─────────────────────────────────────────────────────────────
async def get_groq_reply(messages: list, model_name: str) -> str | None:
    if not GROQ_API_KEY:
        return None
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    try:
        async with httpx.AsyncClient(timeout=40.0) as client:
            r = await client.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers=headers,
                json={"model": model_name, "messages": messages},
            )
            r.raise_for_status()
            return r.json()["choices"][0]["message"]["content"]
    except Exception as e:
        logging.error(f"Groq error ({model_name}): {e}")
        return None


# ─────────────────────────────────────────────────────────────
#  TOOL SCHEMAS  (updated to use free search)
# ─────────────────────────────────────────────────────────────
class ToolSchema(BaseModel):
    name: str
    description: str
    parameters: dict

TOOLS_AVAILABLE = [
    ToolSchema(
        name="web_search",
        description=(
            "Search the web for real-time information. Use for any question about current events, "
            "recent news, prices, people, or anything beyond internal knowledge. "
            "This uses DuckDuckGo — no API key required."
        ),
        parameters={"type":"object","properties":{"query":{"type":"string","description":"Search query"}},"required":["query"]},
    ),
    ToolSchema(
        name="verify_fact",
        description="Fact-check a claim using web search. Returns a summary of what sources say.",
        parameters={"type":"object","properties":{"claim":{"type":"string"}},"required":["claim"]},
    ),
    ToolSchema(
        name="get_current_date",
        description="Returns current date and time. Use when user asks about date/time.",
        parameters={"type":"object","properties":{}},
    ),
    ToolSchema(
        name="get_weather",
        description="Gets current weather for a city. Use when user asks about weather.",
        parameters={"type":"object","properties":{"city":{"type":"string"}},"required":["city"]},
    ),
    ToolSchema(
        name="get_latest_news",
        description="Fetches latest news headlines. Use when user asks for news.",
        parameters={"type":"object","properties":{}},
    ),
    ToolSchema(
        name="update_user_profile",
        description="Save a fact about the user to long-term memory (name, location, preferences, etc).",
        parameters={
            "type":"object",
            "properties":{
                "fact_key":   {"type":"string"},
                "fact_value": {"type":"string"},
            },
            "required":["fact_key","fact_value"],
        },
    ),
    ToolSchema(
        name="get_check_crypto_payment",
        description="Verify if a crypto wallet received a payment.",
        parameters={
            "type":"object",
            "properties":{
                "receiver":{"type":"string"},
                "amount":  {"type":"number"},
            },
            "required":["receiver","amount"],
        },
    ),
]

# ─────────────────────────────────────────────────────────────
#  TOOL EXECUTION
# ─────────────────────────────────────────────────────────────
def get_current_date_internal() -> dict:
    now = datetime.now()
    return {"date": now.strftime("%Y-%m-%d"), "time": now.strftime("%H:%M:%S"), "weekday": now.strftime("%A")}

async def get_weather_internal(city: str) -> dict:
    if not WEATHER_API_KEY:
        return {"error": "Weather API not configured."}
    try:
        url = f"http://api.weatherapi.com/v1/forecast.json?key={WEATHER_API_KEY}&q={city}&days=3"
        async with httpx.AsyncClient(timeout=10.0) as client:
            r = await client.get(url)
            r.raise_for_status()
            d = r.json()
        forecast = [
            f"{day['date']}: {day['day']['condition']['text']}, {day['day']['mintemp_c']}–{day['day']['maxtemp_c']}°C"
            for day in d["forecast"]["forecastday"]
        ]
        return {"location": d["location"], "current": d["current"], "forecast": forecast}
    except Exception as e:
        return {"error": str(e)}

async def get_latest_news_internal() -> dict:
    if not NEWS_API_KEY:
        # Try scraping BBC headlines as fallback
        try:
            results = await ddg_html_search("latest world news today", num_results=5)
            return {"articles": [{"title": r["title"], "description": r["snippet"]} for r in results]}
        except Exception:
            return {"error": "News not available."}
    url = f"https://newsapi.org/v2/top-headlines?country=za&apiKey={NEWS_API_KEY}"
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            r = await client.get(url)
            r.raise_for_status()
            return r.json()
    except Exception as e:
        return {"error": str(e)}

async def execute_tool(tool_name: str, user_id: str, **kwargs) -> dict | str:
    if tool_name == "web_search":
        q = kwargs.get("query")
        if not q:
            return {"error": "Missing query"}
        result = await web_search_free(q)
        return {"result": result}

    if tool_name == "verify_fact":
        claim = kwargs.get("claim", "")
        result = await web_search_free(f"fact check: {claim}")
        return {"claim": claim, "verification_summary": result}

    if tool_name == "get_current_date":
        return get_current_date_internal()

    if tool_name == "get_weather":
        city = kwargs.get("city")
        if not city:
            return {"error": "Missing city"}
        return await get_weather_internal(city)

    if tool_name == "get_latest_news":
        return await get_latest_news_internal()

    if tool_name == "update_user_profile":
        key = kwargs.get("fact_key", "").lower().replace(" ", "_")
        val = kwargs.get("fact_value")
        if user_id and key and val:
            long_term_memory_col.update_one({"user_id": user_id}, {"$set": {key: val}}, upsert=True)
            return {"status": "success", "message": f"Remembered: {key} = {val}"}
        return {"status": "error", "message": "Missing fact_key or fact_value"}

    if tool_name == "get_check_crypto_payment":
        return check_crypto_payment(kwargs.get("receiver"), kwargs.get("amount"))

    return {"error": f"Unknown tool: {tool_name}"}


# ─────────────────────────────────────────────────────────────
#  GROQ WITH TOOL CALLING
# ─────────────────────────────────────────────────────────────
async def get_groq_reply_with_tools(messages: list, model_name: str, user_id: str) -> str | None:
    if not GROQ_API_KEY:
        return "😔 Advanced features unavailable — Groq API key not configured."

    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    url     = "https://api.groq.com/openai/v1/chat/completions"
    current = messages.copy()

    try:
        payload = {
            "model":       model_name,
            "messages":    current,
            "tools":       [{"type": "function", "function": t.model_dump()} for t in TOOLS_AVAILABLE],
            "tool_choice": "auto",
        }
        async with httpx.AsyncClient(timeout=60.0) as client:
            r = await client.post(url, headers=headers, json=payload)
            r.raise_for_status()
        msg = r.json()["choices"][0]["message"]

        if msg.get("tool_calls"):
            current.append({
                "role": "assistant",
                "content": msg.get("content"),
                "tool_calls": msg["tool_calls"],
            })
            for tc in msg["tool_calls"]:
                name = tc["function"]["name"]
                try:
                    args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    args = {}
                try:
                    output = await execute_tool(name, user_id, **args)
                except Exception as e:
                    output = {"error": str(e)}
                current.append({
                    "role":        "tool",
                    "tool_call_id": tc["id"],
                    "content":     json.dumps(output, ensure_ascii=False, default=str),
                })

            async with httpx.AsyncClient(timeout=60.0) as client:
                r2 = await client.post(url, headers=headers, json={"model": model_name, "messages": current})
                r2.raise_for_status()
            return sanitize_ai_response(r2.json()["choices"][0]["message"]["content"])

        return sanitize_ai_response(msg.get("content", ""))

    except httpx.HTTPStatusError as e:
        logging.error(f"Groq HTTP error: {e.response.status_code} — {e.response.text}")
        return None
    except Exception as e:
        logging.error(f"Groq unexpected error: {e}")
        return None


# ─────────────────────────────────────────────────────────────
#  STREAMING GROQ  (SSE)
# ─────────────────────────────────────────────────────────────
async def stream_groq_reply(messages: list, model_name: str) -> AsyncGenerator[str, None]:
    """Yields SSE-formatted chunks for a streaming endpoint."""
    if not GROQ_API_KEY:
        yield "data: {\"chunk\": \"Groq API key not configured.\"}\n\n"
        yield "data: [DONE]\n\n"
        return

    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    payload = {"model": model_name, "messages": messages, "stream": True}

    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            async with client.stream("POST", "https://api.groq.com/openai/v1/chat/completions",
                                      headers=headers, json=payload) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data:"):
                        continue
                    raw = line[5:].strip()
                    if raw == "[DONE]":
                        yield "data: [DONE]\n\n"
                        return
                    try:
                        data   = json.loads(raw)
                        delta  = data["choices"][0].get("delta", {})
                        chunk  = delta.get("content", "")
                        if chunk:
                            yield f"data: {json.dumps({'chunk': chunk})}\n\n"
                    except Exception:
                        continue
    except Exception as e:
        yield f"data: {json.dumps({'error': str(e)})}\n\n"
        yield "data: [DONE]\n\n"


# ─────────────────────────────────────────────────────────────
#  LOCAL AI
# ─────────────────────────────────────────────────────────────
async def get_local_ai_reply(user_message: str, model_name: str) -> str | None:
    paths = get_local_ai_paths(model_name)
    if not os.path.exists(paths["model_path"]) or not os.path.exists(paths["responses_path"]):
        return None
    try:
        pipeline_model = joblib.load(paths["model_path"])
        with open(paths["responses_path"], "r", encoding="utf-8") as f:
            resp_map = json.load(f)
        probs = pipeline_model.predict_proba([user_message])
        best_prob = probs.max()
        if best_prob < LOCAL_AI_CONFIDENCE:
            return None
        label = str(pipeline_model.predict([user_message])[0])
        return resp_map.get(label)
    except Exception as e:
        logging.error(f"Local AI error: {e}")
        return None

async def train_local_ai(prompt: str, reply: str, model_name: str):
    paths = get_local_ai_paths(model_name)
    df = pd.read_csv(paths["data_path"], dtype={"label": str}) if os.path.exists(paths["data_path"]) else pd.DataFrame(columns=["prompt","label"])
    resp_map = json.load(open(paths["responses_path"])) if os.path.exists(paths["responses_path"]) else {}

    label = next((k for k, v in resp_map.items() if v == reply), None)
    if label is None:
        label = str(len(resp_map))
        resp_map[label] = reply

    df = pd.concat([df, pd.DataFrame([{"prompt": prompt, "label": label}])], ignore_index=True)
    df.to_csv(paths["data_path"], index=False)
    with open(paths["responses_path"], "w", encoding="utf-8") as f:
        json.dump(resp_map, f, ensure_ascii=False, indent=2)

    if len(df["label"].unique()) >= 2:
        pipeline_model = Pipeline([("tfidf", TfidfVectorizer()), ("clf", SGDClassifier(loss="modified_huber", random_state=42))])
        pipeline_model.fit(df["prompt"], df["label"])
        joblib.dump(pipeline_model, paths["model_path"])
        logging.info(f"Local model '{model_name}' retrained ({len(df)} samples).")


# ─────────────────────────────────────────────────────────────
#  AUTO PERSONA SELECTOR
# ─────────────────────────────────────────────────────────────
def auto_select_persona(user_message: str, user_id: str | None = None) -> str:
    msg = user_message.lower()
    scores: dict[str, int] = {}
    rules = [
        (["teach","learn","explain","guide","wisdom","how to"], "sensei", 3),
        (["hate","stupid","annoying","whatever","idiot"],       "tsundere", 3),
        (["cute","kawaii","nya","uwu","sparkle","adorable"],    "kawaii", 3),
        (["encourage","motivate","senpai","cheer","help me"],   "senpai", 3),
        (["dark","goth","mystery","shadow","moon","melancholy"],"goth", 3),
        (["battle","fight","game","win","warrior","combat"],    "battle_ai", 3),
        (["mine","forever","obsess","only you","yandere"],      "yandere", 3),
        (["robot","mecha","future","tech","hero","protect"],    "mecha_pilot", 3),
    ]
    for keywords, persona, weight in rules:
        if any(k in msg for k in keywords):
            scores[persona] = scores.get(persona, 0) + weight
    return max(scores, key=scores.get) if scores else "default"


# ─────────────────────────────────────────────────────────────
#  FACT EXTRACTION (background task)
# ─────────────────────────────────────────────────────────────
async def extract_and_save_facts(user_id: str, messages: list):
    last_user = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "")
    if not last_user or len(last_user.strip()) < 5:
        return

    prompt = f"""Extract concrete facts about the user from this message.
Return ONLY a flat JSON object. If no facts, return {{}}.
Message: "{last_user}"
Extract: name, location, age, occupation, hobby, language, preferences.
Strict JSON only, no markdown."""

    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            r = await client.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"},
                json={
                    "model": "llama-3.1-8b-instant",
                    "messages": [{"role":"user","content":prompt}],
                    "temperature": 0.1,
                    "max_tokens": 150,
                    "response_format": {"type": "json_object"},
                },
            )
            r.raise_for_status()
        raw = r.json()["choices"][0]["message"]["content"].strip()
        facts = {k: v for k, v in json.loads(raw).items()
                 if v and str(v).strip().lower() not in ("", "none", "null", "unknown")}
        if facts:
            facts["last_updated"] = datetime.now(timezone.utc)
            long_term_memory_col.update_one({"user_id": user_id}, {"$set": facts}, upsert=True)
            logging.info(f"Saved facts for {user_id}: {facts}")
    except Exception as e:
        logging.warning(f"Fact extraction failed: {e}")


# ─────────────────────────────────────────────────────────────
#  DAILY LIMIT
# ─────────────────────────────────────────────────────────────
def get_user_timezone(user_id: str, ip: str) -> str:
    mem = long_term_memory_col.find_one({"user_id": user_id}) or {}
    if "timezone" in mem:
        return mem["timezone"]
    try:
        r = requests.get(TIMEZONE_API_URL.format(ip=ip), timeout=5)
        tz = r.json().get("timezone", "UTC")
        long_term_memory_col.update_one({"user_id": user_id}, {"$set": {"timezone": tz}}, upsert=True)
        return tz
    except Exception:
        return "UTC"

def has_reached_daily_limit(user_id: str, ip: str) -> bool:
    tz_str = get_user_timezone(user_id, ip)
    try:
        tz = pytz.timezone(tz_str)
    except Exception:
        tz = pytz.UTC
    now_local       = datetime.now(tz)
    today_start_utc = now_local.replace(hour=0, minute=0, second=0, microsecond=0).astimezone(pytz.UTC)
    count = chat_history_col.count_documents({"user_id": user_id, "timestamp": {"$gte": today_start_utc}})
    return count >= DAILY_MESSAGE_LIMIT


# ─────────────────────────────────────────────────────────────
#  FOLLOW-UP SUGGESTION GENERATOR  (Claude-like feature)
# ─────────────────────────────────────────────────────────────
async def generate_follow_up_suggestions(user_message: str, ai_reply: str) -> list[str]:
    """
    Generate 3 smart follow-up question suggestions based on the conversation.
    Runs as background task — frontend can poll or receive via SSE.
    """
    if not GROQ_API_KEY:
        return []
    prompt = f"""Given this conversation exchange, suggest 3 short follow-up questions the user might want to ask next.
Return ONLY a JSON array of 3 strings. No markdown, no explanation.

User asked: "{user_message[:200]}"
AI replied: "{ai_reply[:300]}"

JSON array:"""
    try:
        async with httpx.AsyncClient(timeout=12.0) as client:
            r = await client.post(
                "https://api.groq.com/openai/v1/chat/completions",
                headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"},
                json={
                    "model": "llama-3.1-8b-instant",
                    "messages": [{"role":"user","content":prompt}],
                    "temperature": 0.7,
                    "max_tokens": 120,
                },
            )
            r.raise_for_status()
        content = r.json()["choices"][0]["message"]["content"].strip()
        # parse JSON array
        match = re.search(r'\[.*\]', content, re.DOTALL)
        if match:
            return json.loads(match.group())[:3]
    except Exception as e:
        logging.warning(f"Follow-up suggestions failed: {e}")
    return []


# ─────────────────────────────────────────────────────────────
#  IMAGE HELPERS
# ─────────────────────────────────────────────────────────────
def preprocess_image(image_bytes: bytes) -> np.ndarray:
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((224, 224))
    arr = tf.keras.preprocessing.image.img_to_array(img)
    arr = np.expand_dims(arr, axis=0)
    return tf.keras.applications.mobilenet_v2.preprocess_input(arr)

def build_image_interpretation(predictions: list, ocr_text: str) -> dict:
    if not predictions:
        return {"description": "I was unable to identify the contents of this image."}
    label_map = {"Windsor tie":"a Windsor tie","bow tie":"a bow tie","jersey":"a jersey",
                 "trench coat":"a trench coat","lab coat":"a lab coat","suit":"a suit",
                 "sunglasses":"sunglasses","helmet":"a helmet","jean":"jeans"}
    def clean(label):
        return label_map.get(label.replace("_"," ").strip(), label.replace("_"," ").strip())

    top  = predictions[0]
    tlbl = clean(top["description"])
    tprb = top["probability"]
    supp = [clean(p["description"]) for p in predictions[1:] if p["probability"] >= 0.08]

    if tprb >= 0.60:
        desc = f"This image appears to show {tlbl}" + (f", along with {', '.join(supp)}" if supp else "") + "."
    elif tprb >= 0.35:
        desc = f"This image likely contains {tlbl}" + (f", possibly {', '.join(supp)}" if supp else "") + ". Moderately confident."
    else:
        all_labels = [clean(p["description"]) for p in predictions if p["probability"] >= 0.05]
        desc = (f"I'm not very confident, but this may contain: {', '.join(all_labels)}." if all_labels
                else "I could not confidently identify this image.")

    if ocr_text and len(ocr_text.strip()) > 2:
        desc += f' Text found in image: "{ocr_text.strip()[:300]}"'
    return {"description": desc}


# ─────────────────────────────────────────────────────────────
#  PYDANTIC REQUEST MODELS
# ─────────────────────────────────────────────────────────────
class ChatMessage(BaseModel):
    user_id:         str
    message:         str
    instructions:    str  = ""
    autoPersonality: bool = False
    additionalInfor: str  = ""
    model:           AIModel = DEFAULT_MODEL
    force_groq:      bool    = False
    persona:         Optional[str] = None
    deep_think:      bool    = False
    deep_search:     bool    = False
    response_length: ResponseLength = ResponseLength.BALANCED
    tone:            ToneStyle      = ToneStyle.DEFAULT
    json_mode:       bool    = False   # NEW: structured JSON output

class TranslateRequest(BaseModel):
    user_id:  str
    text:     str
    target_language: str  # e.g. "French", "Zulu", "Japanese"

class SummariseRequest(BaseModel):
    user_id: str
    text:    str
    style:   str = "bullet"   # bullet | paragraph | tldr

class ToneRewriteRequest(BaseModel):
    user_id: str
    text:    str
    tone:    ToneStyle

class BranchRequest(BaseModel):
    user_id:     str
    branch_name: str
    from_message_index: int   # fork chat at this message index

class CodeRunRequest(BaseModel):
    user_id: str
    code:    str
    language: str = "python"


# ─────────────────────────────────────────────────────────────
#  ════════════════════════════════════════════════════════════
#   ENDPOINTS
# ════════════════════════════════════════════════════════════
# ─────────────────────────────────────────────────────────────

# ── HEALTH ────────────────────────────────────────────────────
@app.get("/health")
def health_check():
    return {
        "status": "ok",
        "models_loaded": list(ml_models.keys()),
        "bs4_available": BS4_AVAILABLE,
        "free_search": True,
        "serpapi_fallback": bool(SERPAPI_API_KEY),
    }


# ══════════════════════════════════════════════════════════════
#  MAIN CHAT ENDPOINT  /chat/
# ══════════════════════════════════════════════════════════════
@app.post("/chat/")
async def chat(payload: ChatMessage, request: Request):
    user_id     = payload.user_id
    user_msg    = payload.message.strip()
    ip          = request.client.host if request.client else "127.0.0.1"

    # Guards
    if is_rate_limited(user_id):
        return {"response": "⚡ Slow down! You're sending messages too quickly. Please wait a moment."}
    if has_reached_daily_limit(user_id, ip):
        return {"response": f"😔 Daily limit of {DAILY_MESSAGE_LIMIT} messages reached. Try again tomorrow."}
    if is_inappropriate(user_msg):
        return {"response": "😏 Sorry, I can't respond to that type of message."}

    # Persona
    selected_persona = payload.persona
    if payload.autoPersonality:
        selected_persona = auto_select_persona(user_msg, user_id)

    # --- DEEP SEARCH MODE ---
    if payload.deep_search:
        search_results = await web_search_free(user_msg)
        synthesis_msgs = [
            {"role": "system",  "content": "You are a web search summarizer. Answer based ONLY on the provided search results."},
            {"role": "user",    "content": f"Search results:\n{search_results}\n\nUser question: {user_msg}"},
        ]
        reply = await get_groq_reply(synthesis_msgs, AIModel.GROQ_70B.value)
        if reply:
            reply = sanitize_ai_response(reply)
            asyncio.create_task(train_local_ai(user_msg, reply, payload.model.value))
            save_user_memory(user_id, user_msg, reply)
            suggestions = await generate_follow_up_suggestions(user_msg, reply)
            return {"response": inject_ad(reply, user_id), "follow_up_suggestions": suggestions}
        return {"response": "😅 Search failed. Please try again."}

    # --- STANDARD / DEEP THINK CHAT ---
    deep_think_mode = DeepThinkMode.ADVANCED if payload.deep_think else DeepThinkMode.STANDARD
    location        = load_user_location(user_id)

    system_prompt = get_system_prompt(
        user_id=user_id,
        persona=selected_persona,
        deep_think=deep_think_mode,
        location=location,
        instructions=payload.instructions or None,
        response_length=payload.response_length,
        tone=payload.tone,
    )

    # Build message list for LLM
    memory = load_user_memory(user_id)
    messages_for_llm = [{"role": "system", "content": system_prompt}]
    for m in memory[-10:]:
        messages_for_llm.append({"role": "user",      "content": m["user"][:200]})
        messages_for_llm.append({"role": "assistant",  "content": m["ai"][:250]})
    messages_for_llm.append({"role": "user", "content": user_msg[:400]})

    # JSON mode — wrap system prompt
    if payload.json_mode:
        messages_for_llm[0]["content"] += "\nIMPORTANT: Respond ONLY with valid JSON. No markdown, no explanation."

    final_reply    = None
    groq_fallback  = False

    # Try local model first (unless force_groq or deep_think)
    if not payload.force_groq and not payload.deep_think and payload.model in (AIModel.NEURONES_SELF, AIModel.NEURONES_SELF_3):
        final_reply = await get_local_ai_reply(user_msg, payload.model.value)
        if not final_reply:
            groq_fallback = True

    # Groq with tool calling
    if groq_fallback or not final_reply:
        groq_model = AIModel.GROQ_70B.value if payload.deep_think else AIModel.GROQ_8B.value
        final_reply = await get_groq_reply_with_tools(messages_for_llm, groq_model, user_id)
        if final_reply and is_high_quality_response(final_reply):
            asyncio.create_task(train_local_ai(user_msg, final_reply, payload.model.value))

    if not final_reply:
        fallback = "😅 NeuraPrompt AI couldn't respond. Please retry in a moment."
        save_user_memory(user_id, user_msg, fallback)
        return {"response": fallback}

    final_reply = sanitize_ai_response(final_reply)

    # Background tasks
    asyncio.create_task(extract_and_save_facts(user_id, messages_for_llm))
    suggestions = await generate_follow_up_suggestions(user_msg, final_reply)
    save_user_memory(user_id, user_msg, final_reply)

    resp = {
        "response": inject_ad(final_reply, user_id),
        "follow_up_suggestions": suggestions,
    }
    if payload.autoPersonality and selected_persona:
        resp["auto_selected_persona"] = selected_persona
    return resp


# ══════════════════════════════════════════════════════════════
#  STREAMING ENDPOINT  /chat/stream/
#  Frontend: EventSource('/chat/stream/?...') or fetch with ReadableStream
# ══════════════════════════════════════════════════════════════
@app.post("/chat/stream/")
async def chat_stream(payload: ChatMessage, request: Request):
    """
    Server-Sent Events streaming endpoint.
    Each SSE chunk: data: {"chunk": "..."}
    Final:          data: [DONE]
    """
    user_id  = payload.user_id
    user_msg = payload.message.strip()
    ip       = request.client.host if request.client else "127.0.0.1"

    if is_rate_limited(user_id):
        async def rate_error():
            yield "data: {\"chunk\": \"⚡ Rate limited. Please slow down.\"}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(rate_error(), media_type="text/event-stream")

    selected_persona = payload.persona
    if payload.autoPersonality:
        selected_persona = auto_select_persona(user_msg, user_id)

    location     = load_user_location(user_id)
    system_prompt = get_system_prompt(
        user_id=user_id, persona=selected_persona,
        deep_think=DeepThinkMode.ADVANCED if payload.deep_think else DeepThinkMode.STANDARD,
        location=location, instructions=payload.instructions or None,
        response_length=payload.response_length, tone=payload.tone,
    )

    memory = load_user_memory(user_id)
    messages_for_llm = [{"role": "system", "content": system_prompt}]
    for m in memory[-8:]:
        messages_for_llm.append({"role": "user",     "content": m["user"][:200]})
        messages_for_llm.append({"role": "assistant", "content": m["ai"][:250]})
    messages_for_llm.append({"role": "user", "content": user_msg[:400]})

    groq_model = AIModel.GROQ_70B.value if payload.deep_think else AIModel.GROQ_8B.value

    async def event_generator():
        full_reply = []
        async for chunk in stream_groq_reply(messages_for_llm, groq_model):
            yield chunk
            if chunk.startswith("data: {"):
                try:
                    data = json.loads(chunk[6:].strip())
                    full_reply.append(data.get("chunk", ""))
                except Exception:
                    pass

        # After streaming, save to memory and train
        complete = "".join(full_reply)
        if complete:
            save_user_memory(user_id, user_msg, complete)
            if is_high_quality_response(complete):
                asyncio.create_task(train_local_ai(user_msg, complete, payload.model.value))
            asyncio.create_task(extract_and_save_facts(user_id, messages_for_llm))

    return StreamingResponse(event_generator(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


# ══════════════════════════════════════════════════════════════
#  FREE SEARCH ENDPOINT  /search/
# ══════════════════════════════════════════════════════════════
@app.get("/search/")
async def search_endpoint(q: str = Query(..., description="Search query")):
    """
    Public search endpoint — uses DuckDuckGo, no API key needed.
    Returns raw formatted search results.
    """
    if not q.strip():
        raise HTTPException(status_code=400, detail="Query cannot be empty")
    results = await web_search_free(q.strip())
    return {"query": q, "results": results}


# ══════════════════════════════════════════════════════════════
#  TRANSLATE  /translate/
# ══════════════════════════════════════════════════════════════
@app.post("/translate/")
async def translate(req: TranslateRequest):
    """
    Translate any text to a target language using Groq.
    No extra API key — uses existing Groq key.
    """
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Text is required")
    messages = [
        {"role": "system",  "content": f"You are a translator. Translate the user's text to {req.target_language}. Return ONLY the translated text, no explanations."},
        {"role": "user",    "content": req.text},
    ]
    result = await get_groq_reply(messages, AIModel.GROQ_8B.value)
    return {"original": req.text, "translated": result or "Translation failed.", "language": req.target_language}


# ══════════════════════════════════════════════════════════════
#  SUMMARISE  /summarise/
# ══════════════════════════════════════════════════════════════
@app.post("/summarise/")
async def summarise(req: SummariseRequest):
    """
    Summarise long text in different styles.
    Also accepts plain-text content from documents.
    """
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Text is required")

    style_prompts = {
        "bullet":    "Summarise as concise bullet points.",
        "paragraph": "Summarise in 2-3 clear paragraphs.",
        "tldr":      "Give a TL;DR in 1-2 sentences.",
    }
    style_instruction = style_prompts.get(req.style, style_prompts["bullet"])

    messages = [
        {"role": "system",  "content": f"You are a summarisation expert. {style_instruction}"},
        {"role": "user",    "content": f"Summarise this:\n\n{req.text[:4000]}"},
    ]
    result = await get_groq_reply(messages, AIModel.GROQ_8B.value)
    return {"summary": result or "Summarisation failed.", "style": req.style}


# ══════════════════════════════════════════════════════════════
#  PDF SUMMARISE  /summarise-pdf/
# ══════════════════════════════════════════════════════════════
@app.post("/summarise-pdf/")
async def summarise_pdf(user_id: str = Form(...), file: UploadFile = File(...), style: str = Form("bullet")):
    """Upload a PDF and get a summary — no external service needed."""
    if not PDF_AVAILABLE:
        raise HTTPException(status_code=501, detail="PyPDF2 not installed. pip install PyPDF2")
    try:
        raw = await file.read()
        reader = PyPDF2.PdfReader(io.BytesIO(raw))
        text   = "\n".join(page.extract_text() or "" for page in reader.pages[:15])  # limit pages
        if not text.strip():
            raise HTTPException(status_code=400, detail="Could not extract text from PDF.")
        req = SummariseRequest(user_id=user_id, text=text, style=style)
        return await summarise(req)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"PDF processing failed: {e}")


# ══════════════════════════════════════════════════════════════
#  TONE REWRITE  /rewrite-tone/
# ══════════════════════════════════════════════════════════════
@app.post("/rewrite-tone/")
async def rewrite_tone(req: ToneRewriteRequest):
    """
    Rewrite the given text in a different tone.
    Claude-like feature: make it formal / casual / friendly / bullet points.
    """
    tone_map = {
        ToneStyle.FORMAL:   "Rewrite this text in a formal, professional tone.",
        ToneStyle.CASUAL:   "Rewrite this text in a casual, relaxed conversational tone.",
        ToneStyle.FRIENDLY: "Rewrite this text in a warm, friendly and encouraging tone.",
        ToneStyle.BULLET:   "Convert this text into concise bullet points.",
        ToneStyle.DEFAULT:  "Clean up and improve this text while keeping the same tone.",
    }
    instruction = tone_map.get(req.tone, tone_map[ToneStyle.DEFAULT])
    messages = [
        {"role": "system", "content": instruction + " Return ONLY the rewritten text."},
        {"role": "user",   "content": req.text},
    ]
    result = await get_groq_reply(messages, AIModel.GROQ_8B.value)
    return {"original": req.text, "rewritten": result or "Rewrite failed.", "tone": req.tone}


# ══════════════════════════════════════════════════════════════
#  SAFE CODE INTERPRETER  /run-code/
#  Python only, heavily sandboxed — no filesystem, no imports
# ══════════════════════════════════════════════════════════════
SAFE_BUILTINS = {
    "print": print, "len": len, "range": range, "int": int, "float": float,
    "str": str, "list": list, "dict": dict, "set": set, "tuple": tuple,
    "sum": sum, "max": max, "min": min, "abs": abs, "round": round,
    "sorted": sorted, "enumerate": enumerate, "zip": zip, "map": map,
    "filter": filter, "bool": bool, "type": type, "isinstance": isinstance,
}

@app.post("/run-code/")
async def run_code(req: CodeRunRequest):
    """
    Execute simple Python code in a restricted sandbox.
    Only safe builtins allowed — no imports, no file access, no network.
    Equivalent to Claude's code execution feature.
    """
    if req.language.lower() != "python":
        return {"output": None, "error": f"Language '{req.language}' not yet supported. Only Python."}

    # Block dangerous patterns
    dangerous = ["import", "__import__", "exec(", "eval(", "open(", "os.", "sys.", "subprocess", "socket"]
    for pattern in dangerous:
        if pattern in req.code:
            return {"output": None, "error": f"Blocked: '{pattern}' is not allowed in sandbox."}

    import io as _io, contextlib

    stdout_buffer = _io.StringIO()
    result_vars   = {}
    try:
        with contextlib.redirect_stdout(stdout_buffer):
            exec(compile(req.code, "<sandbox>", "exec"), {"__builtins__": SAFE_BUILTINS}, result_vars)
        output = stdout_buffer.getvalue()
        # If no print output, show last expression value
        if not output.strip() and result_vars:
            last_val = list(result_vars.values())[-1]
            output   = repr(last_val)
        return {"output": output or "(no output)", "error": None, "variables": {k: repr(v) for k, v in result_vars.items() if not k.startswith("_")}}
    except Exception as e:
        return {"output": None, "error": str(e)}


# ══════════════════════════════════════════════════════════════
#  CONVERSATION BRANCHING  /branch/
# ══════════════════════════════════════════════════════════════
@app.post("/branch/create/")
async def create_branch(req: BranchRequest):
    """
    Fork a conversation at a specific message point.
    Creates a new named branch so the user can explore different directions.
    """
    messages = list(chat_history_col.find(
        {"user_id": req.user_id}
    ).sort("timestamp", 1).limit(req.from_message_index * 2))

    if not messages:
        raise HTTPException(status_code=404, detail="No chat history found.")

    branch_id = hashlib.md5(f"{req.user_id}{req.branch_name}{time.time()}".encode()).hexdigest()[:12]
    branches_col.insert_one({
        "branch_id":    branch_id,
        "user_id":      req.user_id,
        "branch_name":  req.branch_name,
        "messages":     [{"role": m["role"], "content": m["content"]} for m in messages],
        "created_at":   datetime.now(timezone.utc),
    })
    return {"branch_id": branch_id, "branch_name": req.branch_name, "message_count": len(messages)}


@app.get("/branch/list/")
async def list_branches(user_id: str = Query(...)):
    """List all conversation branches for a user."""
    branches = list(branches_col.find({"user_id": user_id}, {"_id": 0, "messages": 0}))
    return {"branches": branches}


@app.get("/branch/load/")
async def load_branch(user_id: str = Query(...), branch_id: str = Query(...)):
    """Load a specific conversation branch."""
    branch = branches_col.find_one({"user_id": user_id, "branch_id": branch_id})
    if not branch:
        raise HTTPException(status_code=404, detail="Branch not found.")
    return {"branch_name": branch["branch_name"], "messages": branch["messages"]}


# ══════════════════════════════════════════════════════════════
#  IMAGE ANALYSIS  /image-analysis/  (updated with question param)
# ══════════════════════════════════════════════════════════════
@app.post("/image-analysis/")
async def image_analysis(
    user_id:  str        = Form(...),
    file:     UploadFile = File(...),
    question: str        = Form(""),   # NEW: optional follow-up question
):
    try:
        file_bytes = await file.read()
        file_size_kb = round(len(file_bytes) / 1024, 2)

        # Save to GridFS
        try:
            image_id = fs.put(file_bytes, filename=file.filename, contentType=file.content_type, user_id=user_id)
        except Exception as e:
            return {"status": "error", "message": f"Storage failed: {e}"}

        # OCR
        ocr_text = ""
        if TESSERACT_AVAILABLE:
            try:
                ocr_text = pytesseract.image_to_string(Image.open(io.BytesIO(file_bytes))).strip()
            except Exception as e:
                logging.warning(f"OCR failed: {e}")

        # Classification
        try:
            preprocessed = preprocess_image(file_bytes)
            model        = ml_models.get("image_analyzer")
            preds        = model.predict(preprocessed)
            decoded      = tf.keras.applications.mobilenet_v2.decode_predictions(preds, top=3)[0]
            predictions  = [{"description": lbl, "probability": float(prob)} for (_, lbl, prob) in decoded]
            confident    = max(p["probability"] for p in predictions) > 0.6
        except Exception as e:
            return {"status": "error", "message": f"Classification failed: {e}"}

        # Save metadata
        images_col.insert_one({
            "user_id": user_id, "file_id": image_id,
            "filename": file.filename, "content_type": file.content_type,
            "size_kb": file_size_kb, "predictions": predictions,
            "ocr_text": ocr_text, "question": question, "user_feedback": None,
        })

        interpretation = build_image_interpretation(predictions, ocr_text)

        # If a question was provided, ask Groq to answer it in context
        follow_up_answer = ""
        if question.strip() and GROQ_API_KEY:
            context = interpretation["description"]
            if ocr_text:
                context += f" Text in image: {ocr_text[:200]}"
            qa_msgs = [
                {"role": "system",  "content": "You are an image analysis assistant. Answer the user's question about the image based on the analysis provided."},
                {"role": "user",    "content": f"Image analysis: {context}\n\nUser question: {question.strip()}"},
            ]
            follow_up_answer = await get_groq_reply(qa_msgs, AIModel.GROQ_8B.value) or ""

        payload_out = {
            "status":         "success" if confident else "uncertain",
            "metadata":       {"filename": file.filename, "content_type": file.content_type, "size_kb": file_size_kb},
            "predictions":    predictions,
            "ocr_text":       ocr_text if ocr_text else None,
            "interpretation": interpretation["description"],
            "follow_up":      follow_up_answer or "How can I assist you further with this image? 😊",
        }
        if not confident:
            payload_out["needs_user_verification"] = True
            payload_out["message"] = "I'm not very confident — could you describe what's in this image to help me learn?"
        return payload_out

    except Exception as e:
        return {"status": "error", "message": f"Unexpected failure: {e}"}


# ══════════════════════════════════════════════════════════════
#  EXISTING ENDPOINTS (kept intact from v4)
# ══════════════════════════════════════════════════════════════

@app.post("/image-feedback/")
async def image_feedback(image_id: str = Form(...), feedback: str = Form(...)):
    result = images_col.update_one({"file_id": image_id}, {"$set": {"user_feedback": feedback}})
    if result.modified_count == 0:
        raise HTTPException(status_code=404, detail="Image not found")
    return {"status": "success", "message": f"Feedback saved: {feedback}"}

@app.post("/submit-labeled-image/")
async def submit_labeled_image(user_id: str = Form(...), label: str = Form(...), image_file: UploadFile = File(...)):
    img_bytes = await image_file.read()
    pending_images_col.insert_one({
        "user_id": user_id, "user_label": label.strip().lower(),
        "filename": image_file.filename, "content_type": image_file.content_type,
        "image_data": img_bytes, "status": "pending",
        "timestamp": datetime.now(timezone.utc),
    })
    return {"status": "success", "message": "Thank you! Your feedback will help the AI learn."}

@app.post("/admin/approve-image/{image_id}")
async def approve_image(image_id: str):
    doc = pending_images_col.find_one({"_id": ObjectId(image_id)})
    if not doc:
        raise HTTPException(status_code=404, detail="Pending image not found.")
    label  = re.sub(r'[^a-zA-Z0-9_-]', '', doc["user_label"].replace(" ", "_"))
    target = pathlib.Path(DATASET_PATH) / label
    target.mkdir(parents=True, exist_ok=True)
    (target / f"{int(time.time())}_{doc['filename']}").write_bytes(doc["image_data"])
    pending_images_col.delete_one({"_id": ObjectId(image_id)})
    return {"status": "success", "message": f"Image approved for class '{label}'."}

@app.post("/admin/reject-image/{image_id}")
async def reject_image(image_id: str):
    result = pending_images_col.delete_one({"_id": ObjectId(image_id)})
    if result.deleted_count == 0:
        raise HTTPException(status_code=404, detail="Image not found.")
    return {"status": "success", "message": "Image rejected."}

@app.post("/admin/reset-ai/")
async def reset_ai_data(model_name: AIModel = Query(AIModel.NEURONES_SELF)):
    if APP_MODE == "production":
        raise HTTPException(status_code=403, detail="Reset disabled in production.")
    if "groq" in model_name.value or "openai" in model_name.value:
        raise HTTPException(status_code=400, detail="Cannot reset external Groq models.")
    paths = get_local_ai_paths(model_name.value)
    for p in paths.values():
        if os.path.exists(p):
            os.remove(p)
    return {"message": f"Model '{model_name.value}' data cleared."}

@app.post("/admin/train/")
async def manual_train(
    prompt: str = Form(...),
    reply:  str = Form(...),
    model_name: AIModel = Form(AIModel.NEURONES_SELF),
):
    if "openai" in model_name.value:
        raise HTTPException(status_code=400, detail="Cannot train external models.")
    await train_local_ai(prompt, reply, model_name.value)
    return {"message": f"Model '{model_name.value}' trained."}

@app.get("/api/loadshedding/status")
async def get_loadshedding_status():
    url = f"https://developer.sepush.co.za/business/2.0/status?token={ESKOM_API_KEY}"
    try:
        r = requests.get(url, timeout=15)
        r.raise_for_status()
        return r.json()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/loadshedding/areas_search")
async def search_loadshedding_areas(text: str = Query(...)):
    url = f"https://developer.sepush.co.za/business/2.0/areas_search?text={text}&token={ESKOM_API_KEY}"
    try:
        r = requests.get(url, timeout=15)
        r.raise_for_status()
        return r.json()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/log-ad-click")
def log_ad_click(user_id: str = Query(...), ad_id: str = Query(...)):
    from ai_ads import log_ad_click as _log
    _log(user_id, ad_id)
    return {"message": "Logged."}

@app.get("/tools/date")
def get_date():
    return get_current_date_internal()

@app.get("/tools/weather")
async def get_weather_endpoint(city: str = Query(...)):
    return await get_weather_internal(city)

@app.get("/tools/news")
async def get_news_endpoint():
    return await get_latest_news_internal()

@app.get("/tools/search")
async def search_tool_endpoint(q: str = Query(...)):
    """Test the free search tool directly."""
    return {"results": await web_search_free(q)}

@app.post("/api/verify-crypto")
async def verify_crypto(receiver: str = Form(...), amount: float = Form(...)):
    result = check_crypto_payment(receiver, amount)
    if result.get("success"):
        return result
    raise HTTPException(status_code=404, detail=result.get("message", "Payment not found."))

# Static files — must be last
app.mount("/", StaticFiles(directory="/data/static", html=True), name="static")