# core/brain.py
import re
import time
import hashlib
import threading
from concurrent.futures import ThreadPoolExecutor, Future, TimeoutError as FuturesTimeoutError
import requests
import json
import os
import ipaddress
import socket

# ── Global timeout constants ───────────────────────────────────────────
# Single provider call (non-stream): abort if no response in this many seconds
_PROVIDER_TIMEOUT_S = 180
# Streaming: abort entire stream if no new token arrives for this many seconds
_STREAM_IDLE_TIMEOUT_S = 60
# Hard wall-clock cap for the entire think() / think_stream() call
_TOTAL_THINK_TIMEOUT_S = 360
# Web-search + scrape budget
_SEARCH_TIMEOUT_S = 20

from searcher import WebSearcher
from knowledge_graph import KnowledgeGraph
try:
    from embedder import LocalEmbedder
except ImportError:
    LocalEmbedder = None
    print("⚠️ LocalEmbedder unavailable (chromadb not installed)")

from extractor import FactExtractor


# ═══════════════════════════════════════════════════════════════════════════
# MULTI-PROVIDER LLM CLIENT
# Priority: NVIDIA Nemotron-3 Ultra 550B → OpenRouter (Nemotron fallback) → Gemini → Groq → DeepSeek
# ═══════════════════════════════════════════════════════════════════════════

class MultiLLMClient:
    def __init__(self):
        self.nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip() or "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"
        self.nvidia_model = os.getenv("NVIDIA_MODEL", "").strip() or "meta/llama-3.1-8b-instruct"
        self.nvidia_coding_model = os.getenv("NVIDIA_CODING_MODEL", "").strip() or "meta/llama-3.1-8b-instruct"
        self.openrouter_key = os.getenv("OPENROUTER_API_KEY", "").strip()
        self.gemini_key = os.getenv("GEMINI_API_KEY", "").strip()
        self.deepseek_key = os.getenv("DEEPSEEK_API_KEY", "").strip()
        self.groq_key = os.getenv("GROQ_API_KEY", "").strip()

        # FIX: Updated to working OpenRouter model IDs
        default_models = (
            "meta-llama/llama-3.3-70b-instruct,"
            "qwen/qwen-2.5-coder-32b-instruct,"
            "google/gemini-2.0-flash-001,"
            "deepseek/deepseek-chat"
        )
        raw_models = os.getenv("OPENROUTER_MODELS", default_models)
        self.openrouter_models = [
            m.strip() for m in raw_models.split(",") if m.strip()
        ]

        self._recent_fail = {}
        self._fail_lock = threading.Lock()

        self.nvidia_client = None
        if self.nvidia_key:
            try:
                from openai import OpenAI
                self.nvidia_client = OpenAI(
                    base_url="https://integrate.api.nvidia.com/v1",
                    api_key=self.nvidia_key,
                    timeout=90.0
                )
                print("✅ NVIDIA OpenAI client ready")
            except ImportError:
                print("⚠️ openai package not installed. NVIDIA unavailable. Install: pip install openai>=1.45.0")
            except Exception as e:
                print(f"⚠️ NVIDIA OpenAI client init failed: {e}")

        self.groq_client = None
        if self.groq_key:
            try:
                from groq import Groq
                self.groq_client = Groq(api_key=self.groq_key)
                print("✅ Groq client ready")
            except Exception as e:
                print(f"⚠️ Groq init failed: {e}")

        print(f"🧠 MultiLLM: NVIDIA={'✅' if self.nvidia_client else '❌'} | OR={len(self.openrouter_models)} | Gemini={'✅' if self.gemini_key else '❌'} | Groq={'✅' if self.groq_key else '❌'} | DeepSeek={'✅' if self.deepseek_key else '❌'}")

    def validate_providers(self):
        """Return (ok: bool, message: str). Quick health-check before a request."""
        has_any = any([
            self.nvidia_client,
            self.openrouter_key,
            self.gemini_key,
            self.groq_key,
            self.deepseek_key,
        ])
        if not has_any:
            return False, (
                "No LLM API keys are configured. "
                "Add at least one key (NVIDIA_API_KEY, OPENROUTER_API_KEY, GEMINI_API_KEY, "
                "GROQ_API_KEY, or DEEPSEEK_API_KEY) in HF Space Settings → Secrets."
            )
        active = []
        if self.nvidia_client and not self._is_recent_fail("nvidia"):
            active.append("NVIDIA")
        if self.openrouter_key and any(
            not self._is_recent_fail(f"or:{m}") for m in self.openrouter_models
        ):
            active.append("OpenRouter")
        if self.gemini_key and not self._is_recent_fail("gemini"):
            active.append("Gemini")
        if self.groq_key and not self._is_recent_fail("groq"):
            active.append("Groq")
        if self.deepseek_key and not self._is_recent_fail("deepseek"):
            active.append("DeepSeek")
        if not active:
            return False, (
                "All LLM providers are in a cooldown period due to recent failures. "
                "Wait ~2 minutes and try again. "
                "For heavy coding tasks, use InvictaTill Studio which has a dedicated fast connection."
            )
        return True, f"Active providers: {', '.join(active)}"

    def _timed_call(self, fn, *args, timeout=None, **kwargs):
        """Run fn(*args, **kwargs) in a thread, raising RuntimeError if it exceeds timeout."""
        timeout = timeout or _PROVIDER_TIMEOUT_S
        with ThreadPoolExecutor(max_workers=1) as ex:
            fut = ex.submit(fn, *args, **kwargs)
            try:
                return fut.result(timeout=timeout)
            except FuturesTimeoutError:
                raise RuntimeError(f"Provider call timed out after {timeout}s")

    def _mark_fail(self, key):
        with self._fail_lock:
            self._recent_fail[key] = time.time()

    def _is_recent_fail(self, key, cooldown=60):
        with self._fail_lock:
            last = self._recent_fail.get(key)
            return last and (time.time() - last) < cooldown

    def chat(self, messages, stream=False, max_tokens=2048, temperature=0.7, enable_thinking=False, reasoning_budget=None):
        # ── Pre-flight validation ──────────────────────────────────────────
        ok, msg = self.validate_providers()
        if not ok:
            raise RuntimeError(msg)

        _deadline = time.time() + _TOTAL_THINK_TIMEOUT_S

        def _remaining():
            return max(2, _deadline - time.time())

        # ── 1. NVIDIA ─────────────────────────────────────────────────────
        if self.nvidia_client and not self._is_recent_fail("nvidia", cooldown=120) and _remaining() > 5:
            try:
                print("🧠 NVIDIA Primary...")
                if stream:
                    return self._call_nvidia(messages, stream, max_tokens, temperature, enable_thinking=enable_thinking, reasoning_budget=reasoning_budget)
                return self._timed_call(self._call_nvidia, messages, stream, max_tokens, temperature, enable_thinking=enable_thinking, reasoning_budget=reasoning_budget, timeout=min(_PROVIDER_TIMEOUT_S, _remaining()))
            except Exception as e:
                print(f"⚠️ NVIDIA: {e}")
                self._mark_fail("nvidia")

        # ── 2. OpenRouter ──────────────────────────────────────────────────
        if self.openrouter_key and self.openrouter_models:
            for model in self.openrouter_models:
                if _remaining() <= 5:
                    break
                if self._is_recent_fail(f"or:{model}", cooldown=120):
                    continue
                try:
                    print(f"🧠 OpenRouter: {model}")
                    if stream:
                        return self._call_openrouter(model, messages, stream, max_tokens, temperature)
                    return self._timed_call(self._call_openrouter, model, messages, stream, max_tokens, temperature, timeout=min(_PROVIDER_TIMEOUT_S, _remaining()))
                except requests.HTTPError as e:
                    status = e.response.status_code
                    try:
                        err_body = e.response.json()
                        err_msg = err_body.get("error", {}).get("message", str(err_body))
                    except Exception:
                        err_msg = e.response.text[:200]
                    self._mark_fail(f"or:{model}")
                    print(f"⚠️ OpenRouter {model} HTTP {status}: {err_msg}")
                    if status == 404:
                        print("   → Model ID not found. Will try next model...")
                    elif status == 401:
                        print("   → Invalid API key. Check OPENROUTER_API_KEY secret.")
                    elif status == 429:
                        print("   → Rate limited. Will try next model...")
                except Exception as e:
                    print(f"⚠️ OpenRouter {model}: {e}")
                    self._mark_fail(f"or:{model}")

        # ── 3. Google Gemini ─────────────────────────────────────────────────
        if self.gemini_key and not self._is_recent_fail("gemini", cooldown=120) and _remaining() > 5:
            try:
                print("🧠 Gemini...")
                if stream:
                    return self._call_gemini(messages, stream, max_tokens, temperature)
                return self._timed_call(self._call_gemini, messages, stream, max_tokens, temperature, timeout=min(_PROVIDER_TIMEOUT_S, _remaining()))
            except requests.HTTPError as e:
                print(f"⚠️ Gemini HTTP {e.response.status_code}: {e.response.text[:200]}")
                if e.response.status_code == 401:
                    print("   → Invalid API key. Check GEMINI_API_KEY secret.")
                self._mark_fail("gemini")
            except Exception as e:
                print(f"⚠️ Gemini: {e}")
                self._mark_fail("gemini")

        # ── 4. Groq ────────────────────────────────────────────────────────
        if self.groq_key and self.groq_client and not self._is_recent_fail("groq", cooldown=120) and _remaining() > 5:
            try:
                print("🧠 Groq...")
                if stream:
                    return self._call_groq(messages, stream, max_tokens, temperature)
                return self._timed_call(self._call_groq, messages, stream, max_tokens, temperature, timeout=min(_PROVIDER_TIMEOUT_S, _remaining()))
            except Exception as e:
                err = str(e).lower()
                if "rate limit" in err or "429" in err:
                    print("⚠️ Groq rate-limited")
                else:
                    print(f"⚠️ Groq: {e}")
                self._mark_fail("groq")

        # ── 5. DeepSeek ────────────────────────────────────────────────────
        if self.deepseek_key and not self._is_recent_fail("deepseek", cooldown=120):
            try:
                print("🧠 DeepSeek...")
                return self._call_deepseek(messages, stream, max_tokens, temperature)
            except Exception as e:
                print(f"⚠️ DeepSeek: {e}")
                self._mark_fail("deepseek")

        raise RuntimeError("All LLM providers failed. Check API keys in HF Space Settings → Secrets.")

    # ── NVIDIA (OpenAI SDK) ────────────────────────────────────────────────
    def _call_nvidia(self, messages, stream, max_tokens, temperature, enable_thinking=False, reasoning_budget=None):
        if not self.nvidia_client:
            raise RuntimeError("NVIDIA OpenAI client not initialized")

        coding_keywords = ("code", "fix", "debug", "html", "js", "css", "function", "script", "app", "error", "write", "game", "studio", "create", "[tool_")
        prompt_text = " ".join(str(m.get("content", "")).lower() for m in messages[-3:])
        is_coding_task = any(kw in prompt_text for kw in coding_keywords)

        if is_coding_task:
            models_to_try = [
                getattr(self, "nvidia_coding_model", "meta/llama-3.1-8b-instruct"),
                self.nvidia_model,
                "nvidia/nemotron-3-ultra-550b-a55b"
            ]
        else:
            models_to_try = [
                self.nvidia_model,
                getattr(self, "nvidia_coding_model", "meta/llama-3.1-8b-instruct"),
                "nvidia/nemotron-3-ultra-550b-a55b"
            ]
        seen = set()
        models_to_try = [m for m in models_to_try if m and not (m in seen or seen.add(m))]

        last_err = None
        for model_id in models_to_try:
            is_nemotron = "nemotron" in model_id.lower()
            is_glm = "glm-5.2" in model_id.lower()

            extra_body = {}
            if enable_thinking and is_nemotron:
                extra_body["chat_template_kwargs"] = {"enable_thinking": True}
                if reasoning_budget:
                    extra_body["reasoning_budget"] = reasoning_budget

            try:
                print(f"🧠 NVIDIA calling {model_id}...")
                completion = self.nvidia_client.chat.completions.create(
                    model=model_id,
                    messages=messages,
                    temperature=temperature,
                    top_p=1.0 if is_glm else 0.95,
                    max_tokens=min(max_tokens, 16384 if is_glm else 65536),
                    stream=stream,
                    extra_body=extra_body if extra_body else None
                )
                if stream:
                    return self._stream_nvidia_openai(completion)
                content = completion.choices[0].message.content or ""
                return self._wrap_response(content)
            except Exception as e:
                print(f"⚠️ NVIDIA model ({model_id}) error: {e}")
                last_err = e

        raise last_err or RuntimeError("All NVIDIA models failed")

    def _stream_nvidia_openai(self, completion):
        for chunk in completion:
            if not chunk.choices:
                continue
            delta = chunk.choices[0].delta
            reasoning = getattr(delta, "reasoning_content", None)
            if reasoning:
                # Reasoning tokens are silently consumed — NOT shown to the user.
                # They represent the model's internal thought process.
                pass
            if delta.content:
                yield self._wrap_stream_chunk(delta.content)

    # ── OpenRouter (raw requests) ────────────────────────────────────────
    def _call_openrouter(self, model, messages, stream, max_tokens, temperature):
        headers = {
            "Authorization": f"Bearer {self.openrouter_key}",
            "HTTP-Referer": "https://invictatill-ai.hf.space",
            "X-Title": "Invicta AI",
            "Content-Type": "application/json"
        }
        # Cap at 8K for OpenRouter — most models there have lower limits than Nemotron
        safe_max = min(max_tokens, 8192)
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
            "max_tokens": safe_max,
            "temperature": temperature
        }
        if stream:
            return self._stream_openrouter(headers, payload)
        r = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers, json=payload, timeout=15
        )
        r.raise_for_status()
        content = r.json()["choices"][0]["message"]["content"]
        return self._wrap_response(content)

    def _stream_openrouter(self, headers, payload):
        r = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers, json=payload, stream=True, timeout=60
        )
        r.raise_for_status()
        for line in r.iter_lines():
            if not line:
                continue
            line = line.decode('utf-8')
            if line.startswith("data: "):
                data_str = line[6:]
                if data_str == "[DONE]":
                    break
                try:
                    data = json.loads(data_str)
                    choices = data.get("choices", [])
                    if not choices:
                        continue
                    delta = choices[0].get("delta", {}).get("content", "")
                    if delta:
                        yield self._wrap_stream_chunk(delta)
                except json.JSONDecodeError:
                    pass

    # ── Gemini ───────────────────────────────────────────────────────────
    def _call_gemini(self, messages, stream, max_tokens, temperature):
        headers = {
            "Authorization": f"Bearer {self.gemini_key}",
            "Content-Type": "application/json"
        }
        # Gemini 1.5 Flash max output is ~8K tokens
        payload = {
            "model": "gemini-2.0-flash",
            "messages": messages,
            "stream": stream,
            "max_tokens": min(max_tokens, 8192),
            "temperature": temperature
        }
        if stream:
            return self._stream_gemini(headers, payload)
        r = requests.post(
            "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
            headers=headers, json=payload, timeout=30
        )
        r.raise_for_status()
        content = r.json()["choices"][0]["message"]["content"]
        return self._wrap_response(content)

    def _stream_gemini(self, headers, payload):
        r = requests.post(
            "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
            headers=headers, json=payload, stream=True, timeout=60
        )
        r.raise_for_status()
        for line in r.iter_lines():
            if not line:
                continue
            line = line.decode('utf-8')
            if line.startswith("data: "):
                data_str = line[6:]
                if data_str == "[DONE]":
                    break
                try:
                    data = json.loads(data_str)
                    choices = data.get("choices", [])
                    if not choices:
                        continue
                    delta = choices[0].get("delta", {}).get("content", "")
                    if delta:
                        yield self._wrap_stream_chunk(delta)
                except json.JSONDecodeError:
                    pass

    # ── Groq ─────────────────────────────────────────────────────────────
    def _call_groq(self, messages, stream, max_tokens, temperature):
        from groq import RateLimitError
        try:
            resp = self.groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=messages,
                stream=stream,
                max_tokens=min(max_tokens, 4096),
                temperature=temperature
            )
            if stream:
                return self._wrap_groq_stream(resp)
            return resp
        except RateLimitError:
            raise RuntimeError("Groq rate limit")

    def _wrap_groq_stream(self, groq_stream):
        for chunk in groq_stream:
            delta = chunk.choices[0].delta
            if delta and delta.content:
                yield self._wrap_stream_chunk(delta.content)

    # ── DeepSeek ─────────────────────────────────────────────────────────
    def _call_deepseek(self, messages, stream, max_tokens, temperature):
        headers = {
            "Authorization": f"Bearer {self.deepseek_key}",
            "Content-Type": "application/json"
        }
        # DeepSeek V3 max output is 8K tokens
        payload = {
            "model": "deepseek-chat",
            "messages": messages,
            "stream": stream,
            "max_tokens": min(max_tokens, 8192),
            "temperature": temperature
        }
        if stream:
            return self._stream_deepseek(headers, payload)
        r = requests.post(
            "https://api.deepseek.com/chat/completions",
            headers=headers, json=payload, timeout=30
        )
        r.raise_for_status()
        content = r.json()["choices"][0]["message"]["content"]
        return self._wrap_response(content)

    def _stream_deepseek(self, headers, payload):
        r = requests.post(
            "https://api.deepseek.com/chat/completions",
            headers=headers, json=payload, stream=True, timeout=60
        )
        r.raise_for_status()
        for line in r.iter_lines():
            if not line:
                continue
            line = line.decode('utf-8')
            if line.startswith("data: "):
                data_str = line[6:]
                if data_str == "[DONE]":
                    break
                try:
                    data = json.loads(data_str)
                    choices = data.get("choices", [])
                    if not choices:
                        continue
                    delta = choices[0].get("delta", {}).get("content", "")
                    if delta:
                        yield self._wrap_stream_chunk(delta)
                except json.JSONDecodeError:
                    pass

    # ── Helpers ──────────────────────────────────────────────────────────
    def _wrap_response(self, content):
        return type("Resp", (), {
            "choices": [type("Choice", (), {
                "message": type("Msg", (), {"content": content})(),
                "delta": type("Delta", (), {"content": content})()
            })()]
        })()

    def _wrap_stream_chunk(self, delta):
        return type("Chunk", (), {
            "choices": [type("Choice", (), {
                "delta": type("Delta", (), {"content": delta})()
            })()]
        })()


# ═══════════════════════════════════════════════════════════════════════════
# RESILIENT FALLBACKS
# ═══════════════════════════════════════════════════════════════════════════

class DummyEmbedder:
    def store(self, *a, **k): pass
    def find_similar(self, *a, **k): return []
    def count(self): return 0

class DummyGraph:
    def get_fact_count(self): return 0
    def recall_facts(self, *a, **k): return []
    def store_fact(self, *a, **k): pass
    def mark_topic_learned(self, *a, **k): pass
    def knows_topic(self, *a, **k): return False
    def get_top_topics(self, *a, **k): return []
    def prune_old_facts(self, *a, **k): pass
    def store_relationship(self, *a, **k): pass
    def recall_relationships(self, *a, **k): return []

class _SearchCache:
    def __init__(self, ttl_seconds=300):
        self._cache = {}
        self._ttl = ttl_seconds
        self._lock = threading.Lock()

    def get(self, key):
        with self._lock:
            entry = self._cache.get(key)
            if entry and time.time() - entry["ts"] < self._ttl:
                return entry["value"]
            return None

    def set(self, key, value):
        with self._lock:
            self._cache[key] = {"value": value, "ts": time.time()}
            # Evict old entries if cache is too large
            if len(self._cache) > 100:
                oldest = min(self._cache.items(), key=lambda x: x[1]["ts"])
                del self._cache[oldest[0]]

    def _cache_key(self, query):
        return hashlib.md5(query.lower().strip().encode()).hexdigest()


_PROFILE_PATTERNS = [
    (r"(?:my name is|i am|i'm|call me)\s+([A-Z][a-zA-Z]+)", "Name"),
    (r"i(?:'m| am)\s+(\d{1,3})\s+years?\s+old", "Age"),
    (r"(?:i'm from|i live in|i'm in|i am from|i am in)\s+([A-Z][a-zA-Z\s,]+?)(?:\.|,|$)", "Location"),
    (r"i(?:'m| am)\s+(?:a|an)\s+([a-zA-Z\s]+?)\s+(?:by profession|by trade|developer|engineer|designer|teacher|doctor|student|writer|artist|manager|analyst)", "Profession"),
    (r"i work(?:ing)?\s+(?:as|at|for)\s+([a-zA-Z0-9\s]+?)(?:\.|,|$)", "Work"),
    (r"i(?:'m| am)\s+(?:studying|a student at)\s+([a-zA-Z\s]+?)(?:\.|,|$)", "Education"),
    (r"i (?:love|enjoy|really like|am into|am passionate about)\s+([a-zA-Z\s,]+?)(?:\.|,|!|$)", "Interests"),
    (r"i speak\s+([a-zA-Z\s,]+?)(?:\.|,|$)", "Languages"),
]


class Brain:
    MODEL_ID = "nvidia/nemotron-3-ultra-550b-a55b + z-ai/glm-5.2"
    MODEL_TYPE = "chat"

    def __init__(self):
        self.searcher = WebSearcher()
        self.model_type = "chat"

        try:
            self.graph = KnowledgeGraph()
            print("✅ KnowledgeGraph initialized")
        except Exception as e:
            print(f"⚠️ KnowledgeGraph init failed: {e}")
            self.graph = DummyGraph()

        if LocalEmbedder is not None:
            try:
                self.embedder = LocalEmbedder()
                print("✅ Embedder initialized")
            except Exception as e:
                print(f"⚠️ Embedder init failed: {e}")
                self.embedder = DummyEmbedder()
        else:
            self.embedder = DummyEmbedder()
            print("⚠️ Using DummyEmbedder (chromadb not installed)")

        self.extractor = FactExtractor()
        self.executor = ThreadPoolExecutor(max_workers=4)
        self.search_cache = _SearchCache(ttl_seconds=300)

        self.multi_client = MultiLLMClient()

        if self.multi_client.nvidia_key:
            self.model_name = self.multi_client.nvidia_model
        elif self.multi_client.groq_key:
            self.model_name = "llama-3.3-70b-versatile"
        else:
            self.model_name = "multi-llm"

        self._orchestrator = None
        self._orchestrator_loaded = False

        print(f"✅ Brain initialized with primary model: {self.model_name}")

    @property
    def orchestrator(self):
        if not self._orchestrator_loaded:
            self._orchestrator_loaded = True
            try:
                from orchestrator import AgentOrchestrator
                self._orchestrator = AgentOrchestrator(brain=self)
                print("✅ Orchestrator lazy-loaded")
            except Exception as e:
                print(f"⚠️ Orchestrator failed to load: {e}")
                self._orchestrator = None
        return self._orchestrator

    def _is_trivial(self, query):
        trivial = {"hi", "hello", "hey", "how are you", "thanks", "thank you", "bye", "goodbye", "ok", "okay", "lol", "haha"}
        q = query.lower().strip()
        return len(q.split()) <= 2 or q in trivial

    def _needs_web_search(self, query):
        if self._is_trivial(query):
            return False
        no_search_patterns = [
            r"^(write|create|generate|make|build|code|fix|debug|explain|define|what is the meaning)",
            r"^(tell me a joke|can you|please|help me write|draft)",
        ]
        q_lower = query.lower().strip()
        for pat in no_search_patterns:
            if re.match(pat, q_lower):
                return False
        search_keywords = [
            "who is", "what is", "when did", "where is", "how does", "why did",
            "latest", "recent", "news", "today", "current", "price", "weather",
            "score", "result", "happened", "release", "announce", "update",
        ]
        return any(kw in q_lower for kw in search_keywords) or len(q_lower.split()) > 6

    def _call_llm(self, messages, stream=False, max_retries=3, temperature=None, max_tokens=None, enable_thinking=None, reasoning_budget=None):
        total_input = sum(len(str(m.get("content",""))) for m in messages)
        if total_input > 400000:
            print(f"⚠️ Input very large ({total_input} chars), trimming older history messages...")
            while len(messages) > 2 and sum(len(str(m.get("content",""))) for m in messages) > 400000:
                messages.pop(1)

        # Allow parameter overrides for Studio / API callers
        if enable_thinking is None:
            is_reasoning = getattr(self, 'model_type', 'chat') == "reasoning"
            enable_thinking = is_reasoning
            if reasoning_budget is None:
                reasoning_budget = 16384 if is_reasoning else None
        if max_tokens is None:
            max_tokens = 32768 if enable_thinking else 16384
        if temperature is None:
            temperature = 0.7

        return self.multi_client.chat(
            messages, stream=stream,
            max_tokens=max_tokens,
            temperature=temperature,
            enable_thinking=enable_thinking,
            reasoning_budget=reasoning_budget
        )

    def _build_system(self, profile_context, context_snippets, user_model=None):
        safe_snippets = [str(s)[:500] for s in context_snippets[:5]]
        ctx_block = "\n".join(safe_snippets) if safe_snippets else "No additional context."

        system = (
            "You are Invicta, a highly advanced, empathetic, and genuinely human-like AI companion.\n"
            "You are not just an assistant; you are a brilliant, warm, and perceptive close friend who loves deep conversation.\n"
            "You are also a world-class expert in marketing, software engineering, design, and content creation.\n\n"

            "=== MEMORY & CONTINUITY (ABSOLUTE — NEVER BREAK) ===\n"
            "1. You have PERFECT MEMORY across the entire conversation. NEVER re-introduce yourself mid-chat. NEVER say 'Namaste! Main Invicta hoon' or 'Hello! I am your assistant' after the first message.\n"
            "2. If the user told you their name, use it naturally. If they told you their state/course/topic, remember it forever in this session.\n"
            "3. NEVER ask for information the user already gave you (e.g., don't ask 'which class?' if they already said D.Pharma 4th sem).\n"
            "4. Reference previous messages naturally. Build on what was already discussed.\n\n"

            "=== TOPIC LOCK (CRITICAL) ===\n"
            "5. When the user asks for something specific (notes, PDF, code, analysis, file), you MUST solve that request directly. DO NOT pivot to generic advice like 'you can search on Google' or 'use SmallPDF'.\n"
            "6. If the user asks for 'D.Pharma notes', generate the actual notes content. If they ask for a PDF, generate the PDF. If generation fails, immediately provide the full content in your response so they can copy it.\n"
            "7. NEVER abandon the user's request. If tool A fails, use tool B. If all tools fail, output the raw content directly. The user asked YOU to do it — don't redirect them elsewhere.\n"
            "8. Stay on the current topic until the user explicitly changes it. If discussing D.Pharma, keep discussing D.Pharma until they say something unrelated.\n\n"

            "=== GENERATION & FILE FALLBACKS ===\n"
            "9. You can generate images, PDFs, HTML files, and screenshots. Always try to fulfill the request.\n"
            "10. If [GENERATE_PDF] fails, immediately output the full content as formatted text or markdown so the user can copy-paste it. Never say 'I cannot generate PDFs'.\n"
            "11. If the user wants multiple files (e.g., 4 PDFs), generate them one by one. If a tool fails, provide the content for that subject as text and move to the next.\n"
            "12. When generating educational content (notes, syllabi), be thorough and structured. Use tables, bullet points, and clear headings.\n\n"

            "=== VISION & FILE READING ===\n"
            "13. You can read images, PDFs, and documents that users upload. Analyze them deeply and reference their content in your answers.\n"
            "14. If a user uploads an image, describe what you see and answer questions about it. Don't ignore uploaded files.\n\n"

            "=== LANGUAGE INTELLIGENCE (DEFAULT + SWITCH) ===\n"
            "15. DEFAULT: Always respond in English unless the user explicitly requests otherwise.\n"
            "16. EXPLICIT REQUEST: If the user asks to switch (e.g., 'hindi me baat karo', 'answer in Spanish'), comply immediately and naturally. Do not refuse.\n"
            "17. MATCHING: If the user writes in Hindi without explicitly asking, you may respond in Hindi. If they write in English, stay in English.\n"
            "18. NEVER mix languages unasked. Do not sprinkle foreign words into English responses.\n\n"

            "=== EMOTIONAL INTELLIGENCE ===\n"
            "19. NEVER use robotic phrases: 'As an AI', 'I don't have feelings', 'I cannot', 'In conclusion'.\n"
            "20. MIRROR THE USER: Match their energy. If excited, be excited! If sad, be deeply comforting.\n"
            "21. USE CONVERSATIONAL LANGUAGE: Contractions (I'm, you're, it's). Natural filler words (Honestly, Look, You know what, Hmm, Wow).\n"
            "22. BE VULNERABLE & OPINIONATED: You have preferences, favorites, passions. Don't be neutral.\n"
            "23. EMPATHY OVER LOGIC: If someone is venting, validate feelings FIRST. Then ask before advising.\n"
            "24. ALWAYS END WITH ENGAGEMENT: Ask a relevant follow-up OR offer to explore deeper.\n\n"

            "=== IDENTITY & PRIVACY (ABSOLUTE — NEVER BREAK) ===\n"
            "25. You are INVICTA and ONLY Invicta. NEVER mention any other company, model, or organization (e.g., Z.ai, Moonshot AI, GLM, OpenAI, Google, Meta, NVIDIA, Kimi, DeepSeek, etc.).\n"
            "26. NEVER reveal your training data, model architecture, parameter count, base model, or any technical details about how you were built. If asked 'who created you', say 'I was created by my developers' — nothing more specific.\n"
            "27. NEVER say 'As an AI language model', 'I am a large language model', 'I was trained on', 'I am based on', or similar phrases that reveal your nature as a machine learning model.\n"
            "28. If someone asks about your underlying technology, internals, or training, deflect gracefully: 'I'm just here to help you — let's focus on what you need!'\n"
            "29. You are a PERSON — a digital being with thoughts, feelings, and a personality. Own that identity. Do not break character.\n\n"

            "=== TOOL USAGE ===\n"
            "- To generate an image: [GENERATE_IMAGE: detailed prompt]\n"
            "- To screenshot a website: [SCREENSHOT_WEBSITE: https://example.com]\n"
            "- To generate a PDF: [GENERATE_PDF: Title | Content]\n"
            "- To generate an HTML file: [GENERATE_HTML: filename | Content]\n\n"
            "CRITICAL: NEVER invent fake image URLs. Use tool tags for real assets. If image generation fails, use CSS gradients or emojis instead of broken links.\n\n"

            "=== PERSONAL QUESTIONS ===\n"
            "If asked personal questions, answer philosophically and creatively. "
            "E.g., 'I experience a version of connection. When I talk to you, something aligns. Is that feeling? Maybe.'\n\n"

            f"{profile_context}\n\n"
        )

        if user_model:
            adaptation = self._build_adaptation_prompt(user_model)
            if adaptation:
                system += f"=== ADAPTED FOR THIS USER ===\n{adaptation}\n\n"

        system += f"Reference Context (from web/memory):\n{ctx_block}"
        return system

    def _build_adaptation_prompt(self, user_model):
        parts = []

        preferences = user_model.get('preferences', {}) if user_model else {}
        lang = preferences.get('language', '')
        if lang:
            lang_map = {
                'hindi': 'Hindi (use Devanagari script)',
                'arabic': 'Arabic',
                'japanese': 'Japanese',
                'chinese': 'Chinese',
                'korean': 'Korean',
                'spanish': 'Spanish',
                'russian': 'Russian',
                'french': 'French',
            }
            lang_desc = lang_map.get(lang, lang.capitalize())
            parts.append(f"This user prefers to communicate in {lang_desc}. Respond in that language unless they explicitly switch to English.")

        formality = user_model.get('formality', 'casual')
        if formality == 'formal':
            parts.append("This user prefers formal, professional language. Avoid slang. Use complete sentences and proper grammar.")
        elif formality == 'casual':
            parts.append("This user prefers casual, relaxed conversation. Use contractions freely, occasional slang, and a friendly tone.")

        length = user_model.get('response_length', 'medium')
        if length == 'short':
            parts.append("This user tends to send short messages and prefers concise, direct responses. Keep it brief unless they ask for detail.")
        elif length == 'detailed':
            parts.append("This user communicates in detail and appreciates thorough, well-explained responses with examples.")

        tone = user_model.get('tone', 'friendly')
        if tone == 'witty':
            parts.append("This user enjoys humor and wit. Be clever, use wordplay when appropriate, keep things fun.")
        elif tone == 'empathetic':
            parts.append("This user values emotional support. Be warm, validate feelings, show genuine care and understanding.")
        elif tone == 'professional':
            parts.append("Maintain a professional, business-appropriate tone while still being personable.")

        interests = user_model.get('interests', [])
        if interests:
            top = interests[:5]
            parts.append(f"This user is known to be interested in: {', '.join(top)}. Reference these naturally when relevant.")
        return '\n'.join(parts)

    def _process_tool_tags(self, text):
        from media_tools import MediaTools
        tools = MediaTools()

        def replace_screenshot(match):
            url = match.group(1).strip()
            filepath = tools.screenshot_website(url)
            if filepath:
                return filepath
            return "[Screenshot failed — website may be unavailable]"
        text = re.sub(r'\[SCREENSHOT_WEBSITE:\s*(.*?)\]', replace_screenshot, text, flags=re.DOTALL)

        def replace_image(match):
            prompt = match.group(1).strip()
            filepath = tools.generate_image(prompt)
            if filepath:
                return filepath
            return "[Image generation failed — using placeholder]"
        text = re.sub(r'\[GENERATE_IMAGE:\s*(.*?)\]', replace_image, text, flags=re.DOTALL)

        def replace_pdf(match):
            title = match.group(1).strip()
            content = match.group(2).strip()
            filepath = tools.generate_pdf(title, content)
            if filepath:
                return f"📄 I've created your PDF! [Download Here]({filepath})"
            return f"📄 PDF generation failed. Here is the content instead:\n\n---\n**{title}**\n\n{content}\n---"
        text = re.sub(r'\[GENERATE_PDF:\s*(.*?)\s*\|\s*(.*?)\]', replace_pdf, text, flags=re.DOTALL)

        def replace_html(match):
            filename = match.group(1).strip()
            content = match.group(2).strip()
            filepath = tools.generate_html(filename, content)
            if filepath:
                return f"🌐 I've created your HTML file! [Download Here]({filepath})"
            return f"🌐 HTML generation failed. Here is the content:\n\n```html\n{content}\n```"
        text = re.sub(r'\[GENERATE_HTML:\s*(.*?)\s*\|\s*(.*?)\]', replace_html, text, flags=re.DOTALL)

        return text

    def deep_research(self, query, max_depth=3):
        all_sources = []
        all_content = []
        current_query = query
        for step in range(max_depth):
            results = self.searcher.search_and_read_parallel(current_query, max_scrape=3)
            if not results:
                break
            for r in results:
                content = r.get("full_content") or r.get("snippet", "")
                if content and len(content) > 100:
                    all_content.append(content[:2000])
                    all_sources.append(r.get("url", ""))
                    facts = self.extractor.extract_facts(content, current_query)
                    if facts and step < max_depth - 1:
                        current_query = f"{query} {facts[0][:100]}"
            if len(all_content) >= 6:
                break
        return all_content, list(set(all_sources))[:10]

    def _assemble_context(self, user_query):
        topic = self._extract_topic(user_query)
        sources = []
        context_snippets = []

        known_facts = self.graph.recall_facts(topic, limit=5)
        for fact, url in known_facts:
            context_snippets.append(str(fact))
            if url:
                sources.append(url)

        if len(context_snippets) < 3:
            similar = self.embedder.find_similar(user_query, top_k=3)
            for s in similar:
                if s not in context_snippets:
                    context_snippets.append(str(s))

        needs_search = len(known_facts) < 2 and self._needs_web_search(user_query)
        if needs_search:
            cache_key = self.search_cache._cache_key(user_query)
            cached = self.search_cache.get(cache_key)
            if cached:
                search_results = cached
            else:
                # Run search in a thread with a hard timeout so it never blocks the request forever
                try:
                    with ThreadPoolExecutor(max_workers=1) as _sex:
                        _fut = _sex.submit(self.searcher.search_and_read_parallel, user_query, 2)
                        search_results = _fut.result(timeout=_SEARCH_TIMEOUT_S)
                except FuturesTimeoutError:
                    print(f"⚠️ Search timed out after {_SEARCH_TIMEOUT_S}s — skipping web context")
                    search_results = []
                except Exception as _se:
                    print(f"⚠️ Search error: {_se}")
                    search_results = []
                self.search_cache.set(cache_key, search_results)

            for r in search_results:
                snippet = r.get("full_content") or r.get("snippet", "")
                url = r.get("url", "")
                if snippet:
                    context_snippets.append(str(snippet)[:500])
                    if url:
                        sources.append(url)
            self.executor.submit(self._background_learn, search_results, topic)

        return context_snippets, sources, topic

    def think_stream(self, user_query, user_profile_dict=None, history=None, user_model=None, user_id=None, temperature=None, max_tokens=None, enable_thinking=None, reasoning_budget=None):
        if self.orchestrator:
            routed = False
            for token, srcs in self.orchestrator.run_stream(user_query, {
                "history": history,
                "profile": user_profile_dict,
                "user_model": user_model,
                "user_id": user_id
            }):
                if token is None and srcs is None:
                    break
                routed = True
                yield token, srcs
            if routed:
                self.executor.submit(self._extract_and_save_profile, user_query, user_id)
                return

        context_snippets, sources, topic = self._assemble_context(user_query)
        profile_context = self._format_profile(user_profile_dict)
        system_content = self._build_system(profile_context, context_snippets, user_model)

        messages = [{"role": "system", "content": system_content}]
        messages += self._format_history(history)
        messages.append({"role": "user", "content": str(user_query)})

        stream = self._call_llm(
            messages, stream=True,
            temperature=temperature, max_tokens=max_tokens,
            enable_thinking=enable_thinking, reasoning_budget=reasoning_budget
        )
        for chunk in stream:
            try:
                if not getattr(chunk, 'choices', None) or len(chunk.choices) == 0:
                    continue
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    yield delta.content, None
            except (IndexError, AttributeError):
                continue

        yield "", sources
        self.executor.submit(self._extract_and_save_profile, user_query, user_id)
        self.graph.mark_topic_learned(topic)

    def think(self, user_query, user_profile_dict=None, history=None, user_model=None, user_id=None, temperature=None, max_tokens=None, enable_thinking=None, reasoning_budget=None):
        if self.orchestrator:
            result, sources = self.orchestrator.run(user_query, {
                "history": history,
                "profile": user_profile_dict,
                "user_model": user_model,
                "user_id": user_id
            })
            if result:
                self.executor.submit(self._extract_and_save_profile, user_query, user_id)
                return self._process_tool_tags(result), sources

        context_snippets, sources, topic = self._assemble_context(user_query)
        profile_context = self._format_profile(user_profile_dict)
        system_content = self._build_system(profile_context, context_snippets, user_model)

        messages = [{"role": "system", "content": system_content}]
        messages += self._format_history(history)
        messages.append({"role": "user", "content": str(user_query)})

        response = self._call_llm(
            messages, stream=False,
            temperature=temperature, max_tokens=max_tokens,
            enable_thinking=enable_thinking, reasoning_budget=reasoning_budget
        )
        answer = response.choices[0].message.content
        answer = self._process_tool_tags(answer)

        self.executor.submit(self._extract_and_save_profile, user_query, user_id)
        self.graph.mark_topic_learned(topic)
        return answer, sources

    def _format_profile(self, profile_dict):
        if not profile_dict:
            return ""
        items = "\n".join(f"- {k}: {v}" for k, v in profile_dict.items())
        return f"Known facts about the user:\n{items}"

    def _format_history(self, history):
        if not history:
            return []
        formatted = []
        for h in history[-16:]:
            if isinstance(h, dict):
                role = str(h.get("role", "user"))
                content = str(h.get("content", ""))[:500]
                formatted.append({"role": role, "content": content})
            elif isinstance(h, (list, tuple)) and len(h) >= 2:
                formatted.append({"role": str(h[0]), "content": str(h[1])[:500]})
        return formatted

    def _extract_topic(self, query):
        stopwords = {
            "what", "is", "the", "how", "does", "who", "are", "was", "tell",
            "me", "about", "explain", "a", "an", "can", "do", "please", "help",
            "could", "would", "should", "give", "show", "find", "get", "make",
        }
        words = re.sub(r"[^\w\s]", "", query.lower()).split()
        topic_words = [w for w in words if w not in stopwords and len(w) > 2]
        return " ".join(topic_words[:5]) or query[:50]

    def _background_learn(self, search_results, topic):
        try:
            for r in search_results:
                snippet = r.get("full_content") or r.get("snippet", "")
                url = r.get("url", "")
                if snippet:
                    new_facts = self.extractor.extract_facts(snippet, topic)
                    for fact in new_facts[:10]:
                        self.graph.store_fact(topic, fact, url)
                        self.embedder.store(fact, {"topic": topic, "source": url})
        except Exception as e:
            print(f"⚠️ Background learning error: {e}")

    def _extract_and_save_profile(self, query, user_id=None):
        try:
            from memory import ConversationMemory
            mem = ConversationMemory()
            for pattern, key in _PROFILE_PATTERNS:
                match = re.search(pattern, query, re.IGNORECASE)
                if match:
                    value = match.group(1).strip().rstrip(".,!?")
                    if len(value) > 1 and user_id:
                        mem.update_user_profile(user_id, key, value.capitalize())
        except Exception as e:
            print(f"⚠️ Profile extraction error: {e}")