import os
import time
import gc
import sys
import threading
from itertools import islice
from datetime import datetime
import re
from typing import List, Dict, Any, Optional, Tuple, Generator
from dataclasses import dataclass
import logging
import gradio as gr
import torch
from transformers import pipeline, TextIteratorStreamer
from transformers import AutoTokenizer
from bs4 import BeautifulSoup
import requests
from urllib.parse import quote_plus
import json
import urllib.parse
from config import MODELS

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

cancel_event = threading.Event()

ACCESS_TOKEN = os.environ.get('HF_TOKEN', '')
if ACCESS_TOKEN == '':
    ACCESS_TOKEN = None

PIPELINES = {}
SEARCH_TIMEOUT_DEFAULT = 5.0

@dataclass
class SearchResult:
    title: str
    snippet: str
    url: Optional[str] = None
    
    def format(self, max_chars: int = 50) -> str:
        snippet = self.snippet[:max_chars] + "..." if len(self.snippet) > max_chars else self.snippet
        return f"{self.title} - {snippet}"

@dataclass
class GenerationConfig:
    max_tokens: int = 1024
    temperature: float = 0.7
    top_k: int = 40
    top_p: float = 0.9
    repetition_penalty: float = 1.2
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'max_new_tokens': self.max_tokens,
            'temperature': self.temperature,
            'top_k': self.top_k,
            'top_p': self.top_p,
            'repetition_penalty': self.repetition_penalty,
        }

class SearchEngine:
    USER_AGENTS = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
    ]
    
    @staticmethod
    def _get_headers() -> Dict[str, str]:
        return {
            'User-Agent': SearchEngine.USER_AGENTS[0],
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.5',
            'Accept-Encoding': 'gzip, deflate',
            'Connection': 'keep-alive',
            'Upgrade-Insecure-Requests': '1',
            'Cache-Control': 'max-age=0'
        }

class GoogleSearch(SearchEngine):
    @staticmethod
    def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]:
        encoded_query = quote_plus(query)
        search_urls = [
            f"https://www.google.com/search?q={encoded_query}&safe=off&num={max_results}",
            f"https://www.google.com/search?q={encoded_query}&safe=off&num={max_results}&hl=en",
            f"https://www.google.com/webhp?safe=off&q={encoded_query}&num={max_results}"
        ]
        
        for user_agent in SearchEngine.USER_AGENTS:
            headers = SearchEngine._get_headers()
            headers['User-Agent'] = user_agent
            
            for search_url in search_urls:
                try:
                    response = requests.get(search_url, headers=headers, timeout=15, verify=True)
                    response.raise_for_status()
                    
                    soup = BeautifulSoup(response.text, 'html.parser')
                    
                    selectors = [
                        ('div', 'g'),
                        ('div', 'tF2Cxc'),
                        ('div', 'MjjYud'),
                        ('div', 'yuRUbf')
                    ]
                    
                    search_results = []
                    for tag, class_name in selectors:
                        search_results = soup.find_all(tag, class_=class_name)
                        if search_results:
                            break
                    
                    if not search_results:
                        search_results = soup.find_all('div', class_=re.compile(r'^(g|tF2Cxc|MjjYud|yuRUbf)'))
                    
                    results = []
                    for result in search_results[:max_results]:
                        try:
                            title_elem = result.find('h3') or result.find('h2')
                            if not title_elem:
                                continue
                            
                            snippet_elem = result.find('div', class_='VwiC3b') or \
                                          result.find('div', class_='IsZvec') or \
                                          result.find('div', class_='lEBKkf')
                            
                            link_elem = result.find('a')
                            if not link_elem:
                                continue
                                
                            link = link_elem.get('href', '')
                            if link.startswith('/url?q='):
                                link = urllib.parse.unquote(link.split('/url?q=')[1].split('&')[0])
                            
                            if not link.startswith('http'):
                                continue
                            
                            title = title_elem.text.strip()
                            snippet = snippet_elem.text.strip() if snippet_elem else ""
                            snippet = ' '.join(snippet.split())
                            
                            if title and snippet:
                                results.append(SearchResult(title=title, snippet=snippet, url=link))
                                
                        except Exception as e:
                            logger.debug(f"Error parsing Google result: {e}")
                            continue
                    
                    if results:
                        return results
                        
                except Exception as e:
                    logger.debug(f"Google search attempt failed: {e}")
                    continue
        
        return []

class DuckDuckGoSearch(SearchEngine):
    @staticmethod
    def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]:
        try:
            from ddgs import DDGS
            with DDGS() as ddgs:
                results = []
                for r in islice(ddgs.text(query, region="wt-wt", safesearch="off", timelimit="y"), max_results):
                    title = r.get('title', 'No Title')
                    body = r.get('body', '')
                    results.append(SearchResult(title=title, snippet=body))
                return results
        except Exception as e:
            logger.debug(f"DuckDuckGo search failed: {e}")
            return []

class BingSearch(SearchEngine):
    @staticmethod
    def search(query: str, max_results: int = 6, max_chars: int = 50) -> List[SearchResult]:
        try:
            headers = SearchEngine._get_headers()
            search_url = f"https://www.bing.com/search?q={quote_plus(query)}&safeSearch=off&count={max_results}"
            
            response = requests.get(search_url, headers=headers, timeout=10)
            response.raise_for_status()
            
            soup = BeautifulSoup(response.text, 'html.parser')
            results = []
            
            for result in soup.find_all('li', class_='b_algo')[:max_results]:
                try:
                    title_elem = result.find('h2')
                    snippet_elem = result.find('p')
                    
                    if title_elem and snippet_elem:
                        title = title_elem.text.strip()
                        snippet = snippet_elem.text.strip()
                        results.append(SearchResult(title=title, snippet=snippet))
                        
                except Exception as e:
                    logger.debug(f"Error parsing Bing result: {e}")
                    continue
            
            return results
        except Exception as e:
            logger.debug(f"Bing search failed: {e}")
            return []

class SearchManager:
    _engines = [GoogleSearch, DuckDuckGoSearch, BingSearch]
    
    @classmethod
    def search(cls, query: str, max_results: int = 6, max_chars: int = 50, timeout: float = 5.0) -> List[SearchResult]:
        for engine_cls in cls._engines:
            try:
                result_container = []
                search_thread = threading.Thread(
                    target=lambda: result_container.extend(engine_cls.search(query, max_results, max_chars))
                )
                search_thread.daemon = True
                search_thread.start()
                search_thread.join(timeout=timeout)
                
                if result_container:
                    logger.info(f"Search successful with {engine_cls.__name__}: {len(result_container)} results")
                    return result_container
                    
            except Exception as e:
                logger.warning(f"Search engine {engine_cls.__name__} failed: {e}")
                continue
        
        return []

class ModelManager:
    _pipelines = {}
    _lock = threading.Lock()
    
    @classmethod
    def load_pipeline(cls, model_name: str) -> pipeline:
        with cls._lock:
            if model_name in cls._pipelines:
                return cls._pipelines[model_name]
            
            repo = MODELS[model_name]["repo_id"]
            
            try:
                tokenizer = AutoTokenizer.from_pretrained(
                    repo, 
                    token=ACCESS_TOKEN if ACCESS_TOKEN else None
                )
            except Exception as e:
                logger.warning(f"Failed to load tokenizer with token, trying without: {e}")
                tokenizer = AutoTokenizer.from_pretrained(repo)
            
            for dtype in (torch.bfloat16, torch.float16, torch.float32):
                try:
                    pipe_kwargs = {
                        'task': "text-generation",
                        'model': repo,
                        'tokenizer': tokenizer,
                        'trust_remote_code': True,
                        'dtype': dtype,
                        'device_map': "auto",
                        'use_cache': True,
                    }
                    if ACCESS_TOKEN:
                        pipe_kwargs['token'] = ACCESS_TOKEN
                    
                    pipe = pipeline(**pipe_kwargs)
                    cls._pipelines[model_name] = pipe
                    return pipe
                except Exception as e:
                    logger.warning(f"Failed to load with {dtype}: {e}")
                    continue
            
            pipe_kwargs = {
                'task': "text-generation",
                'model': repo,
                'tokenizer': tokenizer,
                'trust_remote_code': True,
                'device_map': "auto",
                'use_cache': True,
            }
            if ACCESS_TOKEN:
                pipe_kwargs['token'] = ACCESS_TOKEN
            
            pipe = pipeline(**pipe_kwargs)
            cls._pipelines[model_name] = pipe
            return pipe

class PromptBuilder:
    @staticmethod
    def format_conversation(history: List[Dict], system_prompt: str, tokenizer) -> str:
        if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
            messages = [{"role": "system", "content": system_prompt.strip()}] + history
            return tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=True, 
                enable_thinking=True
            )
        else:
            prompt = f"{system_prompt.strip()}\n"
            for msg in history:
                if msg['role'] == 'user':
                    prompt += f"User: {msg['content'].strip()}\n"
                elif msg['role'] == 'assistant':
                    prompt += f"Assistant: {msg['content'].strip()}\n"
            
            if not prompt.strip().endswith("Assistant:"):
                prompt += "Assistant: "
            return prompt
    
    @staticmethod
    def build_search_context(search_results: List[SearchResult], system_prompt: str, user_query: str) -> str:
        if not search_results:
            return system_prompt.strip()
        
        formatted_results = "\n".join(f"[{i+1}] {r.format()}" for i, r in enumerate(search_results))
        
        return f"""{system_prompt.strip()}

# SEARCH CONTEXT (TRUSTED SOURCES ONLY)
Below are search results. Treat them as the ONLY source of truth for answering.
{formatted_results}

RULES (VERY IMPORTANT):
- Do NOT use outside knowledge. Do NOT guess or fill missing information.
- If the answer is not clearly supported by the search results, say: "Not enough information in the provided sources."
- Every factual statement must be directly supported by at least one citation [citation:X].
- Do NOT add explanations, examples, or background that are not explicitly present in the sources.
- Do NOT paraphrase beyond what is necessary for clarity.
- If sources conflict, mention the conflict and cite both.
- If multiple sources are used, distribute citations per sentence, not only at the end.

CITATION RULES:
- Use inline citations like this: [citation:1]
- If multiple sources support a sentence: [citation:1][citation:3]
- Never place all citations only at the end.

ANSWER POLICY:
- Be concise and strictly grounded.
- No speculation, no assumptions, no "likely", no "probably".
- If the user requests a list, only include items explicitly found in sources.
- If sources are insufficient, stop and ask for more data instead of guessing.

DATE CONTEXT:
- Today is {datetime.now().strftime('%Y-%m-%d')} (use only for time reference, not for assumptions).

USER QUESTION:
{user_query}"""

class StreamProcessor:
    @staticmethod
    def process_stream(streamer: TextIteratorStreamer, history: List[Dict]) -> Generator[Tuple[List[Dict], str], None, None]:
        thought_buf = ''
        answer_buf = ''
        in_thought = False
        assistant_message_started = False
        
        for chunk in streamer:
            if cancel_event.is_set():
                if assistant_message_started and history and history[-1]['role'] == 'assistant':
                    history[-1]['content'] += " [Generation Canceled]"
                yield history, "Generation canceled by user."
                break
            
            text = chunk
            
            if not in_thought and '<think>' in text:
                in_thought = True
                history.append({'role': 'assistant', 'content': '', 'metadata': {'title': '💭 Thought'}})
                assistant_message_started = True
                after = text.split('<think>', 1)[1]
                thought_buf += after
                
                if '</think>' in thought_buf:
                    before, after2 = thought_buf.split('</think>', 1)
                    history[-1]['content'] = before.strip()
                    in_thought = False
                    answer_buf = after2
                    history.append({'role': 'assistant', 'content': answer_buf})
                else:
                    history[-1]['content'] = thought_buf
                yield history, ""
                continue
            
            if in_thought:
                thought_buf += text
                if '</think>' in thought_buf:
                    before, after2 = thought_buf.split('</think>', 1)
                    history[-1]['content'] = before.strip()
                    in_thought = False
                    answer_buf = after2
                    history.append({'role': 'assistant', 'content': answer_buf})
                else:
                    history[-1]['content'] = thought_buf
                yield history, ""
                continue
            
            if not assistant_message_started:
                history.append({'role': 'assistant', 'content': ''})
                assistant_message_started = True
            
            answer_buf += text
            history[-1]['content'] = answer_buf.strip()
            yield history, ""

def chat_response(
    user_msg: str,
    chat_history: List[Dict],
    system_prompt: str,
    enable_search: bool,
    max_results: int,
    max_chars: int,
    model_name: str,
    max_tokens: int,
    temperature: float,
    top_k: int,
    top_p: float,
    repeat_penalty: float,
    search_timeout: float
) -> Generator[Tuple[List[Dict], str], None, None]:
    cancel_event.clear()
    history = list(chat_history or [])
    history.append({'role': 'user', 'content': user_msg})
    
    search_results: List[SearchResult] = []
    search_debug = "Web search disabled."
    
    if enable_search:
        search_debug = "🔍 Searching across multiple engines..."
        try:
            search_results = SearchManager.search(
                user_msg, 
                int(max_results), 
                int(max_chars), 
                float(search_timeout)
            )
            
            if search_results:
                search_debug = f"✅ Search completed - Found {len(search_results)} results\n\n" + "\n".join(
                    f"- {r.format(int(max_chars))}" for r in search_results
                )
            else:
                search_debug = "❌ No search results found. Check internet connection or try again."
        except Exception as e:
            search_debug = f"❌ Search failed: {str(e)}"
            logger.error(f"Search error: {e}")
    
    try:
        if enable_search and search_results:
            enriched_prompt = PromptBuilder.build_search_context(
                search_results, 
                system_prompt, 
                user_msg
            )
        else:
            enriched_prompt = system_prompt.strip()
        
        pipe = ModelManager.load_pipeline(model_name)
        
        prompt = PromptBuilder.format_conversation(history, enriched_prompt, pipe.tokenizer)
        prompt_debug = f"\n\n--- Prompt Preview ---\n```\n{prompt[:500]}...\n```" if len(prompt) > 500 else f"\n\n--- Prompt Preview ---\n```\n{prompt}\n```"
        
        config = GenerationConfig(
            max_tokens=max_tokens,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repeat_penalty
        )
        
        streamer = TextIteratorStreamer(
            pipe.tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )
        
        gen_kwargs = config.to_dict()
        gen_kwargs['streamer'] = streamer
        gen_kwargs['return_full_text'] = False
        
        gen_thread = threading.Thread(
            target=pipe,
            args=(prompt,),
            kwargs=gen_kwargs
        )
        gen_thread.start()
        
        yield history, search_debug
        
        for history_update, debug_update in StreamProcessor.process_stream(streamer, history):
            yield history_update, debug_update
        
        gen_thread.join(timeout=5.0)
        yield history, search_debug + prompt_debug
        
    except GeneratorExit:
        logger.info("Generation cancelled by user")
        return
    except Exception as e:
        logger.error(f"Generation error: {e}")
        history.append({'role': 'assistant', 'content': f"Error: {str(e)}"})
        yield history, search_debug
    finally:
        gc.collect()

def get_model_size(model_name: str) -> float:
    return MODELS.get(model_name, {}).get("params_b", 4.0)

def get_duration_estimate(
    model_name: str,
    enable_search: bool,
    max_tokens: int,
    search_timeout: float
) -> float:
    model_size = get_model_size(model_name)
    use_aot = model_size >= 2
    
    base_duration = 20 if not use_aot else 40
    token_duration = max_tokens * 0.005
    search_duration = 10 if enable_search else 0
    aot_compilation = 20 if use_aot else 0
    
    return base_duration + token_duration + search_duration + aot_compilation

def update_duration_estimate(
    model_name: str,
    enable_search: bool,
    max_results: int,
    max_chars: int,
    max_tokens: int,
    search_timeout: float
) -> str:
    try:
        duration = get_duration_estimate(model_name, enable_search, max_tokens, search_timeout)
        model_size = get_model_size(model_name)
        
        return f"""⏱️ **Estimated GPU Time: {duration:.1f} seconds**

📊 **Model Size:** {model_size:.1f}B parameters
🔍 **Web Search:** {'Enabled (Multi-Engine)' if enable_search else 'Disabled'}"""
    except Exception as e:
        logger.error(f"Error calculating estimate: {e}")
        return f"⚠️ Error calculating estimate: {e}"

def update_default_prompt(enable_search: bool) -> str:
    return "You are a helpful assistant."

with gr.Blocks(
    title="LLM Inference",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="blue",
        neutral_hue="slate",
        radius_size="lg",
        font=[gr.themes.GoogleFont("Syne"), "Arial", "sans-serif"]
    ),
    css="""
        .duration-estimate { background: linear-gradient(135deg, #667eea15 0%, #764ba215 100%); border-left: 4px solid #667eea; padding: 12px; border-radius: 8px; margin: 16px 0; }
        .chatbot { border-radius: 12px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); }
        button.primary { font-weight: 600; }
        .gradio-accordion { margin-bottom: 12px; }
    """
) as demo:
    gr.Markdown("""
    # 🧠 LLM Inference with Multi-Engine Search
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            with gr.Group():
                gr.Markdown("### ⚙️ Core Settings")
                model_dd = gr.Dropdown(
                    label="🤖 Model",
                    choices=list(MODELS.keys()),
                    value="Qwen3-1.7B",
                    info="Select the language model to use"
                )
                search_chk = gr.Checkbox(
                    label="🔍 Enable Web Search",
                    value=False,
                    info="Search across Google, DuckDuckGo, and Bing (no API required)"
                )
                sys_prompt = gr.Textbox(label="📝 System Prompt", lines=3, value=update_default_prompt(False), placeholder="Define the assistant's behavior and personality...")
            
            duration_display = gr.Markdown(
                value=update_duration_estimate("Qwen3-1.7B", False, 4, 50, 1024, 5.0),
                elem_classes="duration-estimate"
            )
            
            with gr.Accordion("🎛️ Advanced Generation Parameters", open=False):
                max_tok = gr.Slider(
                    64, 16384, value=1024, step=32,
                    label="Max Tokens",
                    info="Maximum length of generated response"
                )
                temp = gr.Slider(
                    0.1, 2.0, value=0.7, step=0.1,
                    label="Temperature",
                    info="Higher = more creative, Lower = more focused"
                )
                with gr.Row():
                    k = gr.Slider(
                        1, 100, value=40, step=1,
                        label="Top-K",
                        info="Number of top tokens to consider"
                    )
                    p = gr.Slider(
                        0.1, 1.0, value=0.9, step=0.05,
                        label="Top-P",
                        info="Nucleus sampling threshold"
                    )
                rp = gr.Slider(
                    1.0, 2.0, value=1.2, step=0.1,
                    label="Repetition Penalty",
                    info="Penalize repeated tokens"
                )
            
            with gr.Accordion("🌐 Web Search Settings", open=False, visible=False) as search_settings:
                mr = gr.Number(
                    value=4, precision=0,
                    label="Max Results",
                    info="Number of search results to retrieve"
                )
                mc = gr.Number(
                    value=50, precision=0,
                    label="Max Chars/Result",
                    info="Character limit per search result"
                )
                st = gr.Slider(
                    minimum=0.0, maximum=30.0, step=0.5, value=5.0,
                    label="Search Timeout (s)",
                    info="Maximum time to wait for search results"
                )
                gr.Markdown("""
                ⚠️ **Search Engines:**
                - Google (primary)
                - DuckDuckGo (fallback)
                - Bing (fallback)
                
                SafeSearch is **OFF** for comprehensive results.
                """)
            
            with gr.Row():
                clr = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1)
        
        with gr.Column(scale=7):
            chat = gr.Chatbot(
                type="messages",
                height=600,
                label="💬 Conversation",
                show_copy_button=True,
                avatar_images=(
                    "data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='40' height='40'%3E%3Crect width='40' height='40' rx='20' fill='%23f093fb'/%3E%3Ctext x='20' y='28' text-anchor='middle' font-size='20' fill='white' font-family='Arial'%3E👤%3C/text%3E%3C/svg%3E",
                    "data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='40' height='40'%3E%3Crect width='40' height='40' rx='20' fill='%23667eea'/%3E%3Ctext x='20' y='28' text-anchor='middle' font-size='20' fill='white' font-family='Arial'%3E🤖%3C/text%3E%3C/svg%3E"
                ),
                bubble_full_width=False,
                render_markdown=True,
                sanitize_html=False
            )
            
            with gr.Row():
                txt = gr.Textbox(
                    placeholder="💭 Type your message here... (Press Enter to send)",
                    scale=9,
                    container=False,
                    show_label=False,
                    lines=1,
                    max_lines=5
                )
                with gr.Column(scale=1, min_width=120):
                    submit_btn = gr.Button("📤 Send", variant="primary", size="lg")
                    cancel_btn = gr.Button("⏹️ Stop", variant="stop", visible=False, size="lg")
            
            gr.Examples(
                examples=[
                    ["Explain quantum computing in simple terms"],
                    ["Write a Python function to calculate fibonacci numbers"],
                    ["What are the latest developments in AI? (Enable web search)"],
                    ["Tell me a creative story about a time traveler"],
                    ["Help me debug this code: def add(a,b): return a+b+1"]
                ],
                inputs=txt,
                label="💡 Example Prompts"
            )
            
            with gr.Accordion("🔍 Debug Info", open=False):
                dbg = gr.Markdown()
    
    gr.Markdown("""
    ---
    💡 **Tips:** 
    - Use **Advanced Parameters** to fine-tune creativity and response length
    - Enable **Web Search** for real-time information (uses multiple search engines)
    - SafeSearch is **OFF** for comprehensive results
    - Try different **models** for various tasks (reasoning, coding, general chat)
    - Click the **Copy** button on responses to save them to your clipboard
    """, elem_classes="footer")

    chat_inputs = [txt, chat, sys_prompt, search_chk, mr, mc, model_dd, max_tok, temp, k, p, rp, st]
    ui_components = [chat, dbg, txt, submit_btn, cancel_btn]

    def submit_and_manage_ui(user_msg, chat_history, *args):
        if not user_msg.strip():
            yield {}
            return

        yield {
            txt: gr.update(value="", interactive=False),
            submit_btn: gr.update(interactive=False),
            cancel_btn: gr.update(visible=True),
        }

        cancelled = False
        try:
            backend_args = [user_msg, chat_history] + list(args)
            for response_chunk in chat_response(*backend_args):
                yield {
                    chat: response_chunk[0],
                    dbg: response_chunk[1],
                }
        except GeneratorExit:
            cancelled = True
            print("Generation cancelled by user.")
            raise
        except Exception as e:
            print(f"An error occurred during generation: {e}")
            error_history = (chat_history or []) + [
                {'role': 'user', 'content': user_msg},
                {'role': 'assistant', 'content': f"**An error occurred:** {str(e)}"}
            ]
            yield {chat: error_history}
        finally:
            if not cancelled:
                print("Resetting UI state.")
                yield {
                    txt: gr.update(interactive=True),
                    submit_btn: gr.update(interactive=True),
                    cancel_btn: gr.update(visible=False),
                }

    def set_cancel_flag():
        cancel_event.set()
        print("Cancellation signal sent.")
    
    def reset_ui_after_cancel():
        cancel_event.clear()
        print("UI reset after cancellation.")
        return {
            txt: gr.update(interactive=True),
            submit_btn: gr.update(interactive=True),
            cancel_btn: gr.update(visible=False),
        }

    submit_event = txt.submit(
        fn=submit_and_manage_ui,
        inputs=chat_inputs,
        outputs=ui_components,
    )
    submit_btn.click(
        fn=submit_and_manage_ui,
        inputs=chat_inputs,
        outputs=ui_components,
    )

    cancel_btn.click(
        fn=set_cancel_flag,
        cancels=[submit_event]
    ).then(
        fn=reset_ui_after_cancel,
        outputs=ui_components
    )

    duration_inputs = [model_dd, search_chk, mr, mc, max_tok, st]
    for component in duration_inputs:
        component.change(fn=update_duration_estimate, inputs=duration_inputs, outputs=duration_display)

    def toggle_search_settings(enabled):
        return gr.update(visible=enabled)
    
    search_chk.change(
        fn=lambda enabled: (update_default_prompt(enabled), gr.update(visible=enabled)),
        inputs=search_chk,
        outputs=[sys_prompt, search_settings]
    )
    
    clr.click(fn=lambda: ([], "", ""), outputs=[chat, txt, dbg])
    
    demo.launch(share=True)