import hashlib
import io
import json
import math
import operator
import os
import re
import spaces  # Add this import at the top of app.py
import sys
import textwrap
import time
import warnings
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field, asdict
from datetime import datetime
from functools import wraps
from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Annotated

import gradio as gr
import requests
import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline as hf_pipeline

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver
from langchain_huggingface import HuggingFacePipeline

warnings.filterwarnings("ignore")

# ==============================================================================
# 1. CONFIGURATION
# ==============================================================================
@dataclass
class SwarmConfig:
    model_id: str               = "Qwen/Qwen2.5-7B-Instruct"
    torch_dtype: str            = "float16"
    device_map: str             = "auto"
    trust_remote_code: bool     = True

    max_new_tokens: int         = 1024
    temperature: float          = 0.15
    top_p: float                = 0.92
    repetition_penalty: float   = 1.1

    max_iterations: int         = 4
    max_workers: int            = 4
    worker_timeout_sec: int     = 240

    context_budget_tokens: int  = 6000
    memory_dim: int             = 512
    memory_top_k: int           = 5

    enable_tools: bool          = True
    enable_self_update: bool    = True
    search_timeout_sec: int     = 10

    approval_threshold: float   = 0.75
    min_confidence: float       = 0.65

    log_level: str              = "INFO"
    cache_system_prompts: bool  = True

CONFIG = SwarmConfig()

QUERY_CLASSES = [
    "factual", "analytical", "creative", "coding", 
    "math", "conversational", "multi_step", "comparison"
]

VERIFIER_WEIGHTS: Dict[str, Dict[str, float]] = {
    "factual":        {"accuracy": 0.40, "completeness": 0.30, "safety": 0.10, "relevance": 0.15, "format": 0.05},
    "analytical":     {"accuracy": 0.30, "completeness": 0.25, "safety": 0.10, "relevance": 0.25, "format": 0.10},
    "creative":       {"accuracy": 0.05, "completeness": 0.20, "safety": 0.15, "relevance": 0.35, "format": 0.25},
    "coding":         {"accuracy": 0.40, "completeness": 0.25, "safety": 0.15, "relevance": 0.15, "format": 0.05},
    "math":           {"accuracy": 0.50, "completeness": 0.25, "safety": 0.05, "relevance": 0.15, "format": 0.05},
    "conversational": {"accuracy": 0.10, "completeness": 0.15, "safety": 0.20, "relevance": 0.35, "format": 0.20},
    "multi_step":     {"accuracy": 0.30, "completeness": 0.35, "safety": 0.10, "relevance": 0.15, "format": 0.10},
    "comparison":     {"accuracy": 0.30, "completeness": 0.30, "safety": 0.10, "relevance": 0.20, "format": 0.10},
}

CLASS_WORKER_PROFILE: Dict[str, Dict] = {
    "factual":        {"count": 3, "strategy": "parallel",    "tools": ["web_search"]},
    "analytical":     {"count": 2, "strategy": "sequential",  "tools": ["web_search"]},
    "creative":       {"count": 1, "strategy": "sequential",  "tools": []},
    "coding":         {"count": 1, "strategy": "sequential",  "tools": ["python"]},
    "math":           {"count": 2, "strategy": "sequential",  "tools": ["calculator", "python"]},
    "conversational": {"count": 1, "strategy": "sequential",  "tools": []},
    "multi_step":     {"count": 4, "strategy": "sequential",  "tools": ["web_search", "calculator", "python"]},
    "comparison":     {"count": 2, "strategy": "parallel",    "tools": ["web_search"]},
}

# ==============================================================================
# 2. RESILIENCE MANAGER & CONTEXT WINDOW
# ==============================================================================
class ResilienceManager:
    def __init__(self, max_retries: int = 3, base_delay: float = 0.8):
        self.max_retries = max_retries
        self.base_delay  = base_delay

    def retry(self, func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exc = None
            for attempt in range(self.max_retries + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as exc:
                    last_exc = exc
                    if attempt == self.max_retries: break
                    time.sleep(self.base_delay * (2 ** attempt))
            raise last_exc
        return wrapper

    @staticmethod
    def extract_json(text: str) -> Optional[Dict]:
        if not text: return None
        text = text.strip()
        m = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text)
        if m:
            try: return json.loads(m.group(1))
            except: pass
        start = text.find("{")
        if start != -1:
            depth = 0
            for i in range(start, len(text)):
                if text[i] == "{": depth += 1
                elif text[i] == "}":
                    depth -= 1
                    if depth == 0:
                        try: return json.loads(text[start:i+1])
                        except: break
        return None

resilience = ResilienceManager()

class ContextWindowManager:
    CHARS_PER_TOKEN = 3.7
    def __init__(self, budget_tokens: int): self.budget = budget_tokens
    def to_tokens(self, text: str) -> int: return max(1, int(len(text) / self.CHARS_PER_TOKEN))
    def fit(self, system: str, messages: List[Dict], reserve: int = 300) -> List[Dict]:
        avail = max(self.budget - self.to_tokens(system) - reserve, 0)
        trimmed, used = [], 0
        for msg in reversed(messages):
            tok = self.to_tokens(msg.get("content", ""))
            if used + tok > avail: break
            trimmed.insert(0, msg)
            used += tok
        return trimmed

ctx_mgr = ContextWindowManager(CONFIG.context_budget_tokens)

# ==============================================================================
# 3. SEMANTIC MEMORY & TOOL REGISTRY
# ==============================================================================
class SemanticMemory:
    def __init__(self, dim: int = CONFIG.memory_dim):
        self.dim = dim
        self._store: Dict[str, List[Tuple[str, np.ndarray]]] = defaultdict(list)
        rng = np.random.RandomState(42)
        self._proj = rng.randn(1024, dim).astype(np.float32)

    def _embed(self, text: str) -> np.ndarray:
        vec = np.zeros(1024, dtype=np.float32)
        text = text.lower()
        for n in (2, 3):
            for i in range(len(text) - n + 1):
                vec[abs(hash(text[i:i+n])) % 1024] += 1.0
        emb = vec @ self._proj
        norm = np.linalg.norm(emb)
        return emb / norm if norm > 0 else emb

    def add(self, text: str, scope: str = "general"):
        if len(text) >= 20: self._store[scope].append((text, self._embed(text)))

    def query(self, query: str, scope: str = "general", top_k: int = CONFIG.memory_top_k) -> List[str]:
        pool = self._store.get(scope, []) + self._store.get("general", [])
        if not pool: return []
        q_vec = self._embed(query)
        scored = sorted(pool, key=lambda x: float(np.dot(q_vec, x[1])), reverse=True)
        results = []
        for text, _ in scored[:top_k]:
            if text not in results: results.append(text)
        return results

_MEMORY: Dict[str, SemanticMemory] = {}
def get_memory(thread_id: str) -> SemanticMemory:
    if thread_id not in _MEMORY: _MEMORY[thread_id] = SemanticMemory()
    return _MEMORY[thread_id]

@dataclass
class ToolResult:
    tool: str
    params: str
    output: str
    success: bool
    latency_ms: float

class ToolRegistry:
    def __init__(self):
        self._tools: Dict[str, Callable] = {}
        self._descriptions: Dict[str, str] = {}
        self._register_defaults()

    def _register_defaults(self):
        self.register("web_search",  self._web_search, "Perform a live web search.")
        self.register("calculator",  self._calculator, "Evaluate a math expression safely.")
        self.register("python",      self._python_exec, "Execute Python code snippets.")

    def register(self, name: str, fn: Callable, description: str = ""):
        self._tools[name] = fn
        self._descriptions[name] = description or "No description provided."

    def register_dynamic_code(self, name: str, code: str, description: str = "") -> Tuple[bool, str]:
        try:
            local_scope: Dict[str, Any] = {}
            safe_builtins = {k: v for k, v in __builtins__.items() if k not in ["eval", "exec", "input"]}
            exec(code, {"__builtins__": safe_builtins, "np": np, "re": re, "math": math, "json": json}, local_scope)
            
            target_fn = local_scope.get(name)
            if not target_fn or not callable(target_fn):
                callables = [v for k, v in local_scope.items() if callable(v) and not k.startswith("_")]
                if not callables: return False, f"No callable function found for '{name}'."
                target_fn = callables[0]

            def wrapped_tool(params: str) -> str:
                try: return str(target_fn(params))
                except Exception as err: return f"Execution Error ({name}): {err}"

            self.register(name, wrapped_tool, description)
            return True, f"Feature tool '{name}' successfully registered."
        except Exception as exc:
            return False, f"Compilation error: {exc}"

    def available(self) -> List[str]: return list(self._tools.keys())
    def get_descriptions(self) -> str: return "\n".join(f"- **{k}**: {self._descriptions.get(k, '')}" for k in self._tools)

    def execute(self, name: str, params: str) -> ToolResult:
        t0 = time.time()
        if name not in self._tools: return ToolResult(name, params, f"Unknown tool: {name}", False, 0.0)
        try:
            out = self._tools[name](params)
            return ToolResult(name, params, str(out), True, (time.time()-t0)*1000)
        except Exception as exc:
            return ToolResult(name, params, f"Tool error ({name}): {exc}", False, (time.time()-t0)*1000)

    def _web_search(self, query: str) -> str:
        api_key = os.environ.get("TAVILY_API_KEY")
        if not api_key: return "[NO_RESULTS] TAVILY_API_KEY missing."
        try:
            res = requests.post("https://api.tavily.com/search", json={"api_key": api_key, "query": query, "max_results": 3}, timeout=CONFIG.search_timeout_sec)
            data = res.json()
            snippets = [f"{r['title']}\n{r['content']}" for r in data.get("results", [])]
            return f"[SEARCH: {query}]\n\n" + "\n\n---\n\n".join(snippets)
        except Exception as exc: return f"[NO_RESULTS] Search failed: {exc}"

    def _calculator(self, expr: str) -> str:
        safe_ns = {k: v for k, v in np.__dict__.items() if not k.startswith("_")}
        safe_ns["__builtins__"] = {}
        try: return str(eval(expr, safe_ns, {}))
        except Exception as exc: return f"Calculation error: {exc}"

    def _python_exec(self, code: str) -> str:
        old_out, sys.stdout = sys.stdout, io.StringIO()
        try:
            local = {}
            exec(code, {"__builtins__": {"print": print, "range": range, "len": len, "str": str, "int": int, "float": float, "list": list, "dict": dict, "sum": sum, "abs": abs, "round": round}}, local)
            out = sys.stdout.getvalue()
            return out.strip() or str({k: v for k, v in local.items() if not k.startswith("_")})
        except Exception as exc: return f"Exec error: {exc}"
        finally: sys.stdout = old_out

tools = ToolRegistry()

# ==============================================================================
# 4. MODEL ROUTER
# ==============================================================================
class ModelRouter:
    def __init__(self):
        self.tokenizer = None
        self.pipe      = None
        self._load()

    def _load(self):
        dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
        dtype = dtype_map.get(CONFIG.torch_dtype, torch.float16)

        self.tokenizer = AutoTokenizer.from_pretrained(CONFIG.model_id, trust_remote_code=CONFIG.trust_remote_code)
        if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            CONFIG.model_id, device_map=CONFIG.device_map, torch_dtype=dtype, trust_remote_code=CONFIG.trust_remote_code
        )
        model.eval()

        raw_pipe = hf_pipeline(
            "text-generation", model=model, tokenizer=self.tokenizer,
            max_new_tokens=CONFIG.max_new_tokens, temperature=CONFIG.temperature,
            top_p=CONFIG.top_p, repetition_penalty=CONFIG.repetition_penalty,
            do_sample=True, return_full_text=False, pad_token_id=self.tokenizer.pad_token_id,
        )
        self.pipe = HuggingFacePipeline(pipeline=raw_pipe)

    @resilience.retry
    def query(self, messages: List[Dict]) -> str:
        system_content, chat_msgs = "", []
        for m in messages:
            if m.get("role") == "system": system_content = m["content"]
            else: chat_msgs.append(m)

        fitted = ctx_mgr.fit(system_content, chat_msgs)
        if system_content: fitted = [{"role": "system", "content": system_content}] + fitted

        prompt = self.tokenizer.apply_chat_template(fitted, tokenize=False, add_generation_prompt=True)
        response = self.pipe.invoke(prompt)
        return response.strip() if response else ""

router = ModelRouter()

# ==============================================================================
# 5. GRAPH STATE & NODES
# ==============================================================================
@dataclass
class WorkerOutput:
    agent: str
    result: str
    tool_calls: List[ToolResult] = field(default_factory=list)
    latency_ms: float = 0.0
    error: Optional[str] = None
    def to_dict(self) -> Dict:
        d = asdict(self)
        d["tool_calls"] = [asdict(t) for t in self.tool_calls]
        return d

@dataclass
class WorkerLedger:
    agent: str
    iteration: int
    status: str
    task: str
    succeeded_at: List[str]
    failed_at: List[str]
    partial_knowledge: str
    tools_tried: List[str]
    tools_succeeded: List[str]
    suggested_fix: str
    error_detail: Optional[str]
    def to_dict(self) -> Dict: return asdict(self)

class SwarmState(TypedDict):
    user_request: str
    thread_id: str
    query_class: str
    meta_prompts: Dict[str, Any]
    orchestrator_plan: Dict[str, Any]
    worker_outputs: Annotated[List[Dict], operator.add]
    worker_ledger: Annotated[List[Dict], operator.add]
    verification_report: Dict[str, Any]
    conversation_history: Annotated[List[Dict], operator.add]
    retrieved_memories: List[str]
    created_features: Annotated[List[Dict], operator.add]
    current_iteration: int

def classify_node(state: SwarmState) -> Dict:
    raw = router.query([
        {"role": "system", "content": "Classify into EXACTLY ONE: factual | analytical | creative | coding | math | conversational | multi_step | comparison. Output ONLY the word."},
        {"role": "user", "content": state["user_request"]},
    ]).strip().lower()
    raw = re.sub(r"[^a-z_]", "", raw)
    return {"query_class": raw if raw in QUERY_CLASSES else "multi_step"}

def meta_prompt_node(state: SwarmState) -> Dict:
    qc = state["query_class"]
    raw = router.query([
        {"role": "system", "content": f"Generate short JSON instructions for {qc}. Format: {{'orchestrator_instruction': '...', 'worker_instruction': '...', 'verifier_instruction': '...', 'synthesis_instruction': '...', 'expected_format': 'prose'}}. No markdown."},
        {"role": "user", "content": f"Request: {state['user_request']}"},
    ])
    meta = resilience.extract_json(raw) or {
        "orchestrator_instruction": "Plan tasks.", "worker_instruction": "Execute task.",
        "verifier_instruction": "Score outputs.", "synthesis_instruction": "Merge.", "expected_format": "prose"
    }
    return {"meta_prompts": meta}

def memory_node(state: SwarmState) -> Dict:
    memories = get_memory(state["thread_id"]).query(state["user_request"], scope=state["query_class"])
    return {"retrieved_memories": memories}

_FEATURE_BUILDER_SYSTEM = """You are the Swarm Architecture Self-Update Engine.
Synthesize a NEW Python feature tool that the swarm lacks.
Output ONLY valid JSON matching:
{
  "feature_name": "lowercase_tool_name",
  "description": "Clear description of tool",
  "python_code": "def feature_name(params: str) -> str:\\n    return result_string"
}
Constraint: Accept single string `params` and return string."""

def feature_builder_node(state: SwarmState) -> Dict:
    if not CONFIG.enable_self_update: return {"created_features": []}
    ledger = state.get("worker_ledger", [])
    failed_gaps = [e.get("suggested_fix", "") for e in ledger if e.get("status") in ("failed", "partial")]
    
    prompt_context = f"Request: {state['user_request']}\nAvailable Tools: {tools.available()}\nGaps: {failed_gaps}"
    raw = router.query([
        {"role": "system", "content": _FEATURE_BUILDER_SYSTEM},
        {"role": "user", "content": prompt_context}
    ])
    spec = resilience.extract_json(raw)
    if not spec or "feature_name" not in spec or "python_code" not in spec: return {"created_features": []}

    fname = re.sub(r"[^a-z0-9_]", "", spec["feature_name"].lower())
    success, msg = tools.register_dynamic_code(fname, spec["python_code"], spec.get("description", ""))
    return {"created_features": [{"feature_name": fname, "description": spec.get("description", ""), "code": spec["python_code"], "registered": success}]}

def orchestrator_node(state: SwarmState) -> Dict:
    meta = state["meta_prompts"]
    system_prompt = f"""{meta['orchestrator_instruction']}\nAvailable Tools:\n{tools.get_descriptions()}\nOutput ONLY JSON: {{"subagents": [{{"name": "Worker1", "role": "desc", "task": "specific task", "required_tools": ["python"], "depends_on": []}}], "execution_strategy": "parallel"}}"""
    plan = resilience.extract_json(router.query([{"role": "system", "content": system_prompt}, {"role": "user", "content": state["user_request"]}]))
    if not plan or not plan.get("subagents"):
        plan = {"subagents": [{"name": "Worker", "role": "Agent", "task": state["user_request"], "required_tools": [], "depends_on": []}], "execution_strategy": "parallel"}
    return {"orchestrator_plan": plan, "worker_outputs": []}

def _run_worker(spec: Dict, meta: Dict, prior_outputs: Dict[str, str], iteration: int = 1) -> Tuple[WorkerOutput, WorkerLedger]:
    name, role, task = spec.get("name", "Worker"), spec.get("role", "Assistant"), spec.get("task", "")
    msgs = [
        {"role": "system", "content": f"{meta.get('worker_instruction')}\nYou are {role}. Subtask: {task}. Available tools: {tools.available()}. Call tools via <tool_call>tool: name\\nparams: params</tool_call>"},
        {"role": "user", "content": task}
    ]
    t0, tool_calls, error_msg = time.time(), [], None
    try: reply = router.query(msgs)
    except Exception as exc: error_msg, reply = str(exc), ""

    if reply and "<tool_call>" in reply:
        for match in re.finditer(r"<tool_call>([\s\S]*?)</tool_call>", reply):
            block = match.group(1)
            t_name_m = re.search(r"tool:\s*(\w+)", block)
            t_params_m = re.search(r"params:\s*([\s\S]+)", block)
            if t_name_m and t_params_m:
                t_name, t_params = t_name_m.group(1).strip(), t_params_m.group(1).strip()
                if t_name in tools.available():
                    tr = tools.execute(t_name, t_params)
                    tool_calls.append(tr)
                    msgs.extend([{"role": "assistant", "content": reply}, {"role": "user", "content": f"[Tool result: {t_name}]\n{tr.output}"}])
                    try: reply = router.query(msgs)
                    except Exception as exc2: error_msg = str(exc2)

    latency = (time.time() - t0) * 1000
    ledger = WorkerLedger(name, iteration, "success" if reply.strip() else "failed", task, [task] if reply.strip() else [], [] if reply.strip() else [task], reply[:300], [t.tool for t in tool_calls], [t.tool for t in tool_calls if t.success], "", error_msg)
    return WorkerOutput(agent=name, result=reply, tool_calls=tool_calls, latency_ms=latency, error=error_msg), ledger

def workers_node(state: SwarmState) -> Dict:
    plan, meta, subagents = state["orchestrator_plan"], state["meta_prompts"], state["orchestrator_plan"].get("subagents", [])
    results, ledgers = [], []
    for ag in subagents:
        wo, ledger = _run_worker(ag, meta, {}, state.get("current_iteration", 1))
        results.append(wo); ledgers.append(ledger)
    return {"worker_outputs": [wo.to_dict() for wo in results], "worker_ledger": [l.to_dict() for l in ledgers]}

def verifier_node(state: SwarmState) -> Dict:
    outputs = state["worker_outputs"]
    cur_iter = state.get("current_iteration", 1)
    summary = "\n".join(f"[{w['agent']}]\n{w['result'][:800]}" for w in outputs)
    raw = router.query([
        {"role": "system", "content": f"Score outputs (0.0-1.0). Threshold: {CONFIG.approval_threshold}. JSON: {{'approved': true, 'confidence': 0.9, 'missing_capability': false, 'final_synthesis': 'merged answer'}}"},
        {"role": "user", "content": f"Query: {state['user_request']}\nOutputs:\n{summary}"}
    ])
    report = resilience.extract_json(raw) or {"approved": True, "final_synthesis": outputs[0]["result"] if outputs else "No response."}
    if cur_iter >= CONFIG.max_iterations: report["approved"] = True
    return {"verification_report": report, "current_iteration": cur_iter + 1}

def validation_gate(state: SwarmState) -> str:
    rep = state["verification_report"]
    if rep.get("approved"): return "synthesis"
    if CONFIG.enable_self_update and rep.get("missing_capability"): return "feature_builder"
    return "orchestrator"

def synthesis_node(state: SwarmState) -> Dict:
    raw = router.query([
        {"role": "system", "content": "Format and polish the final synthesis. Clean out internal system tags."},
        {"role": "user", "content": f"Draft:\n{state['verification_report'].get('final_synthesis', '')}"}
    ])
    return {"verification_report": {**state["verification_report"], "final_synthesis": raw.strip()}}

def store_memory_node(state: SwarmState) -> Dict:
    final = state["verification_report"].get("final_synthesis", "")
    if final: get_memory(state["thread_id"]).add(final, scope=state.get("query_class", "general"))
    return {"conversation_history": [{"role": "user", "content": state["user_request"]}, {"role": "assistant", "content": final}]}

# ==============================================================================
# 6. BUILD SWARM ENGINE GRAPH
# ==============================================================================
def build_swarm():
    wf = StateGraph(SwarmState)
    nodes = [
        ("classify", classify_node), ("meta_prompt", meta_prompt_node),
        ("memory", memory_node), ("feature_builder", feature_builder_node),
        ("orchestrator", orchestrator_node), ("workers", workers_node),
        ("verifier", verifier_node), ("synthesis", synthesis_node),
        ("store_memory", store_memory_node)
    ]
    for n, f in nodes: wf.add_node(n, f)
    
    wf.set_entry_point("classify")
    wf.add_edge("classify", "meta_prompt")
    wf.add_edge("meta_prompt", "memory")
    wf.add_edge("memory", "feature_builder")
    wf.add_edge("feature_builder", "orchestrator")
    wf.add_edge("orchestrator", "workers")
    wf.add_edge("workers", "verifier")
    wf.add_conditional_edges("verifier", validation_gate, {
        "feature_builder": "feature_builder",
        "orchestrator": "orchestrator", 
        "synthesis": "synthesis"
    })
    wf.add_edge("synthesis", "store_memory")
    wf.add_edge("store_memory", END)
    return wf.compile(checkpointer=InMemorySaver())

class SwarmEngine:
    def __init__(self):
        self.app = build_swarm()

    def run(self, user_request: str, thread_id: str = "default") -> Dict:
        initial = {
            "user_request": user_request, "thread_id": thread_id, "query_class": "factual",
            "meta_prompts": {}, "orchestrator_plan": {}, "worker_outputs": [],
            "worker_ledger": [], "verification_report": {}, "conversation_history": [],
            "retrieved_memories": [], "created_features": [], "current_iteration": 1
        }
        final = self.app.invoke(initial, config={"configurable": {"thread_id": thread_id}})
        return {
            "response": final["verification_report"].get("final_synthesis", ""),
            "created_features": final.get("created_features", []),
        }

swarm = SwarmEngine()

# ==============================================================================
# 7. GRADIO CHAT INTERFACE FUNCTION
# ==============================================================================
@spaces.GPU  # This decorator tells Hugging Face ZeroGPU to assign VRAM on-demand
def respond(
    message: str,
    history: List[Dict[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Executes the Swarm Engine workflow and streams back the synthesized response
    along with any dynamic features self-generated during execution.
    """
    # Synchronize sliders to config
    CONFIG.max_new_tokens = max_tokens
    CONFIG.temperature = temperature
    CONFIG.top_p = top_p

    # Run query through Swarm Graph
    res = swarm.run(message)
    response_text = res["response"]

    # Append any dynamic self-generated code tools to output display
    if res.get("created_features"):
        response_text += "\n\n---\n### 🛠️ Dynamically Self-Generated Tools:"
        for feat in res["created_features"]:
            response_text += f"\n\n**Tool Name:** `{feat['feature_name']}`\n"
            response_text += f"**Description:** {feat['description']}\n"
            response_text += f"```python\n{feat['code']}\n```"

    yield response_text

# Setup Gradio ChatInterface
chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are an autonomous self-updating swarm engine.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.15, step=0.05, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.92, step=0.05, label="Top-p (nucleus sampling)"),
    ],
    title="🐝 Swarm Engine v7.0 — Self-Updating Multi-Agent Architecture",
    description="Powered by Qwen2.5-7B-Instruct & LangGraph. The swarm plans, executes, verifies, and writes its own Python tools on the fly."
)

with gr.Blocks() as demo:
    chatbot.render()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860,share=True)