backups/app2.py

import gradio as gr
from google import genai
from google.genai import types
from google.genai.types import Tool, GoogleSearch, FunctionDeclaration
from PIL import Image
import io
import traceback
import datetime
import re
import importlib
import os
import sys
from typing import List, Dict, Any, Optional
from pathlib import Path

# Add current directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from base_extension import BaseExtension

DESCRIPTION = """
# GemiWine 
**Powered by Gemini 2.5 Flash + Google Search Grounding + Agent Extensions**
"""

BASE_SYSTEM_PROMPT = """
You are GemiWine, a helpful AI assistant with extensible capabilities.
Your core abilities include conversation, web search, and image understanding.

When users enable extensions, you gain additional tools and capabilities.
Always use the available tools when they would be helpful to the user.
Be proactive about suggesting when an extension might be useful.
"""


def log(msg: str):
    now = datetime.datetime.now().strftime("%H:%M:%S")
    print(f"[{now}] {msg}", flush=True)


def get_mime_type(file_path: str) -> str:
    """Determine MIME type from file extension"""
    ext = Path(file_path).suffix.lower()
    mime_types = {
        # Images
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.gif': 'image/gif',
        '.webp': 'image/webp',
        '.heic': 'image/heic',
        '.heif': 'image/heif',
        # Documents
        '.pdf': 'application/pdf',
        '.txt': 'text/plain',
        '.html': 'text/html',
        '.md': 'text/markdown',
        # Videos
        '.mp4': 'video/mp4',
        '.mpeg': 'video/mpeg',
        '.mov': 'video/mov',
        '.avi': 'video/avi',
        '.flv': 'video/x-flv',
        '.mpg': 'video/mpg',
        '.webm': 'video/webm',
        '.wmv': 'video/wmv',
        '.3gpp': 'video/3gpp',
    }
    return mime_types.get(ext, 'application/octet-stream')


def process_uploaded_file(client: genai.Client, file_path: str) -> types.Part:
    """Process an uploaded file and return a Part object"""
    mime_type = get_mime_type(file_path)
    file_size = Path(file_path).stat().st_size
    
    log(f"📎 Processing file: {Path(file_path).name} ({mime_type}, {file_size/1024:.1f}KB)")
    
    # For files > 20MB or videos, use File API
    if file_size > 20 * 1024 * 1024 or mime_type.startswith('video/'):
        log(f"📤 Uploading large file via File API...")
        uploaded_file = client.files.upload(file=file_path)
        log(f"✅ File uploaded: {uploaded_file.name}")
        return uploaded_file
    else:
        # For smaller files, pass inline
        with open(file_path, 'rb') as f:
            file_bytes = f.read()
        log(f"✅ File loaded inline")
        return types.Part.from_bytes(data=file_bytes, mime_type=mime_type)


class ExtensionManager:
    """Manages loading and interfacing with extensions"""
    
    def __init__(self):
        self.extensions: Dict[str, BaseExtension] = {}
        self.load_extensions()
    
    def load_extensions(self):
        """Dynamically load all extensions from extensions/ folder"""
        extensions_dir = Path("extensions")
        if not extensions_dir.exists():
            log("⚠️ Extensions directory not found, creating it...")
            extensions_dir.mkdir()
            return
        
        log(f"🔍 Scanning for extensions in {extensions_dir.absolute()}")
        
        for file in extensions_dir.glob("*.py"):
            if file.name.startswith("_"):
                log(f"⏭️ Skipping {file.name} (starts with _)")
                continue
            
            try:
                log(f"📦 Attempting to load: {file.name}")
                module_name = file.stem
                spec = importlib.util.spec_from_file_location(module_name, file)
                module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(module)
                
                # Find Extension class in module
                found_extension = False
                for attr_name in dir(module):
                    attr = getattr(module, attr_name)
                    if (isinstance(attr, type) and 
                        issubclass(attr, BaseExtension) and 
                        attr != BaseExtension):
                        ext = attr()
                        self.extensions[ext.name] = ext
                        log(f"✅ Loaded extension: {ext.display_name} ({ext.name})")
                        found_extension = True
                        break
                
                if not found_extension:
                    log(f"⚠️ No extension class found in {file.name}")
                    
            except Exception as e:
                log(f"❌ Failed to load {file.name}: {e}")
                traceback.print_exc()
        
        log(f"📊 Total extensions loaded: {len(self.extensions)}")
    
    def get_extension(self, name: str) -> Optional[BaseExtension]:
        return self.extensions.get(name)
    
    def get_all_extensions(self) -> List[BaseExtension]:
        return list(self.extensions.values())
    
    def get_enabled_extensions(self, user_id: str, enabled_list: List[str]) -> List[BaseExtension]:
        """Get list of enabled extension objects"""
        return [ext for name, ext in self.extensions.items() if name in enabled_list]
    
    def build_system_prompt(self, enabled_list: List[str]) -> str:
        """Build system prompt with enabled extension contexts"""
        prompt = BASE_SYSTEM_PROMPT
        
        enabled_exts = self.get_enabled_extensions("", enabled_list)
        if enabled_exts:
            prompt += "\n\n# ENABLED EXTENSIONS\nYou currently have these extensions enabled:\n\n"
            for ext in enabled_exts:
                prompt += f"## {ext.display_name}\n{ext.get_system_context()}\n\n"
        
        return prompt
    
    def get_all_tools(self, enabled_list: List[str]) -> List[types.Tool]:
        """Get all tools from enabled extensions (no search tool here)"""
        tools = []
        
        # Add extension tools only
        for ext_name in enabled_list:
            ext = self.get_extension(ext_name)
            if ext:
                tools.extend(ext.get_tools())
        
        return tools
    
    def get_search_tool(self) -> types.Tool:
        """Get Google Search tool separately"""
        return types.Tool(google_search=types.GoogleSearch())
    
    def handle_function_calls(self, user_id: str, enabled_list: List[str], function_calls: List) -> List:
        """Process function calls from Gemini and return results"""
        results = []
        
        for fc in function_calls:
            function_name = fc.name
            args = fc.args if hasattr(fc, 'args') else {}
            
            # Find which extension owns this function
            handled = False
            for ext_name in enabled_list:
                ext = self.get_extension(ext_name)
                if ext:
                    # Check if this function is in the extension's tools
                    for tool in ext.get_tools():
                        if hasattr(tool, 'function_declarations'):
                            for func_decl in tool.function_declarations:
                                if func_decl.name == function_name:
                                    result = ext.handle_tool_call(user_id, function_name, args)
                                    # Keep result as dict, don't convert to string yet
                                    results.append(result)
                                    handled = True
                                    break
                        if handled:
                            break
                if handled:
                    break
            
            if not handled:
                results.append({"error": f"Unknown function {function_name}"})
        
        return results


class AgentOrchestrator:
    """Orchestrates multiple specialized agents"""
    
    def __init__(self, client, chat, extension_manager, enabled_extensions):
        self.client = client
        self.chat = chat  # Multi-turn chat session
        self.extension_manager = extension_manager
        self.enabled_extensions = enabled_extensions
        # Create a separate chat session for search (to isolate it from tool calls)
        self.search_chat = client.chats.create(model="gemini-2.5-flash")
        
    def call_search_agent(self, query: str, file_parts: List = None) -> tuple:
        """Call specialized search agent using streaming - returns (text, citations)"""
        log("🔍 Calling Search Agent...")
        
        grounding_tool = types.Tool(google_search=types.GoogleSearch())
        config = types.GenerateContentConfig(
            system_instruction="You are a search specialist. Use Google Search to find relevant, accurate information. Provide concise, well-cited answers.",
            tools=[grounding_tool],
            temperature=0.7,
            max_output_tokens=2048
        )
        
        try:
            # Build message content with files if provided
            content_parts = []
            if file_parts:
                content_parts.extend(file_parts)
            content_parts.append(query)
            
            # Use streaming like the working example
            result_text = ""
            last_chunk = None
            
            stream = self.search_chat.send_message_stream(content_parts, config=config)
            for chunk in stream:
                last_chunk = chunk
                if hasattr(chunk, 'candidates') and chunk.candidates:
                    candidate = chunk.candidates[0]
                    if hasattr(candidate, 'content') and candidate.content:
                        if hasattr(candidate.content, 'parts') and candidate.content.parts:
                            for part in candidate.content.parts:
                                if hasattr(part, 'text') and part.text:
                                    result_text += part.text
            
            # Extract citations from the last chunk using the working function
            citations = None
            if last_chunk and hasattr(last_chunk, 'candidates') and last_chunk.candidates:
                log(f"🔍 Extracting citations from search response...")
                citations = insert_citations_from_grounding(last_chunk.candidates)
                if citations:
                    log(f"✅ Citations extracted successfully")
                else:
                    log(f"⚠️ No citations found in grounding metadata")
            
            if result_text:
                log(f"✅ Search Agent returned {len(result_text)} chars")
            else:
                log(f"⚠️ Search Agent returned empty result")
            
            return result_text, citations
            
        except Exception as e:
            log(f"⚠️ Search Agent error: {e}")
            traceback.print_exc()
        
        return "", None
    
    def call_tool_agent(self, query: str, search_context: str = "", reasoning_budget: int = -1, file_parts: List = None) -> tuple:
        """Call tool execution agent with function calling - uses multi-turn chat"""
        log("🛠️ Calling Tool Agent...")
        
        # Build prompt with context if needed
        prompt = query
        if search_context:
            prompt = f"[Context from Search]\n{search_context}\n\n[User Request]\n{query}"
        
        # Get extension tools
        tools = self.extension_manager.get_all_tools(self.enabled_extensions)
        
        system_prompt = self.extension_manager.build_system_prompt(self.enabled_extensions)
        system_prompt += """

CRITICAL INSTRUCTIONS FOR TOOL USAGE:
- You have PERSISTENT STATE across all conversations in this chat session
- Timers, tasks, notes, and other data remain even after responses
- When users ask about "the timer", "the alarm", "my tasks", etc., they're referring to items created earlier
- ALWAYS use your tools (list_timers, list_tasks, check_timer, etc.) when asked about status
- Don't say you can't access information - use your available tools first
- Be proactive: if user mentions checking something, use the appropriate tool immediately

If search context is provided, incorporate it naturally.
When images, PDFs, videos, or other files are provided, analyze them thoroughly and reference them in your response."""
        
        config = types.GenerateContentConfig(
            system_instruction=system_prompt,
            tools=tools,
            temperature=0.7,
            max_output_tokens=4096,
            thinking_config=types.ThinkingConfig(
                include_thoughts=True,
                thinking_budget=reasoning_budget,
            )
        )
        
        try:
            # Build message content with files if provided
            content_parts = []
            if file_parts:
                content_parts.extend(file_parts)
            content_parts.append(prompt)
            
            # Use the chat session's send_message (maintains conversation history automatically)
            response = self.chat.send_message(
                content_parts,
                config=config
            )
            
            function_calls = []
            text_response = ""
            thoughts = ""
            
            if response.candidates and response.candidates[0].content:
                for part in response.candidates[0].content.parts:
                    if hasattr(part, 'function_call') and part.function_call:
                        function_calls.append(part.function_call)
                        log(f"🔧 Tool call: {part.function_call.name}")
                    
                    if getattr(part, "text", None):
                        if getattr(part, "thought", False):
                            thoughts += part.text
                        else:
                            text_response += part.text
            
            return function_calls, text_response, thoughts
            
        except Exception as e:
            log(f"⚠️ Tool Agent error: {e}")
            traceback.print_exc()
            return [], "", ""
    
    def synthesize_response(self, query: str, search_results: str, tool_results: list, search_citations: Optional[str] = None, file_parts: List = None) -> tuple:  
        """Synthesize final response from all sources - returns (text, images_html)"""
        log("✨ Synthesizing final response...")
    
        synthesis_prompt = f"[Original Query]\n{query}\n\n"
    
        if search_results:
            synthesis_prompt += f"[Web Search Results]\n{search_results}\n\n"
        
        # Collect any generated images from tool results
        generated_images = []
        if tool_results:
            synthesis_prompt += "[Tool Execution Results]\n"
            for tool_name, result in tool_results:
                if result is None:
                    result = "(no result)"
                
                # Check if result contains a generated chart/image
                if isinstance(result, dict) and 'image_base64' in result:
                    generated_images.append({
                        'base64': result['image_base64'],
                        'title': result.get('message', 'Generated visualization'),
                        'filepath': result.get('filepath', '')
                    })
                    # Don't include base64 in the synthesis prompt (too long)
                    result_clean = dict(result)
                    result_clean.pop('image_base64', None)
                    synthesis_prompt += f"- {tool_name}: {result_clean.get('message', '')} (Chart created and will be displayed)\n"
                else:
                    synthesis_prompt += f"- {tool_name}: {result}\n"
            synthesis_prompt += "\n"
    
        synthesis_prompt += "Provide a comprehensive answer that incorporates all available information above. Be natural and conversational."
        
        # If files were provided, reference them in the context
        if file_parts:
            synthesis_prompt += "\n\nNote: The user has provided files (images/documents/videos) with their query. Make sure to reference and discuss the content of these files in your response."
    
        config = types.GenerateContentConfig(
            system_instruction="You are a synthesis specialist. Combine information from multiple sources into coherent, helpful responses. When files are provided, analyze and reference them in your answer.",
            temperature=0.7,
            max_output_tokens=4096
        )
    
        try:
            # Build content parts with files if provided
            content_parts = []
            if file_parts:
                content_parts.extend(file_parts)
            content_parts.append(types.Part(text=synthesis_prompt))
            
            response = self.client.models.generate_content(
                model="gemini-2.5-flash",
                contents=[types.Content(role="user", parts=content_parts)],
                config=config
            )
        
            result_text = ""
            if response.candidates and response.candidates[0].content:
                for part in response.candidates[0].content.parts:
                    if getattr(part, "text", None):
                        result_text += part.text
            
            return result_text, generated_images
        except Exception as e:
            log(f"⚠️ Synthesis error: {e}")
    
        return "I encountered an error synthesizing the response.", []



def determine_needs_search(chat, query: str) -> bool:
    """Determine if query needs web search - uses chat session for reliability"""
    
    # Simple heuristic first - if query explicitly asks to search
    search_keywords = ['search', 'find online', 'look up online', 'google', 'search online', 'check online']
    if any(keyword in query.lower() for keyword in search_keywords):
        log(f"🔍 Search triggered by explicit keyword")
        return True
    
    # For questions about recommendations, comparisons, "best" items - likely needs search
    recommendation_keywords = ['best', 'recommend', 'top', 'which', 'what are good', 'compare']
    if any(keyword in query.lower() for keyword in recommendation_keywords):
        log(f"🔍 Search triggered by recommendation question")
        return True
    
    # Default to no search for timer/task management queries
    internal_keywords = ['timer', 'alarm', 'task', 'note', 'how much time']
    if any(keyword in query.lower() for keyword in internal_keywords):
        log(f"❌ No search - internal tool query")
        return False
    
    log(f"❌ No search - general query")
    return False


# Global instances
EXTENSION_MANAGER = ExtensionManager()
CHAT_SESSIONS: Dict[str, Dict[str, Any]] = {}


def get_or_create_session(api_key: str):
    if not api_key:
        return None, None
    if api_key in CHAT_SESSIONS:
        return (CHAT_SESSIONS[api_key]["client"], 
                CHAT_SESSIONS[api_key]["chat"])
    try:
        client = genai.Client(api_key=api_key)
        # Create a chat session for multi-turn conversations
        chat = client.chats.create(model="gemini-2.5-flash")
        CHAT_SESSIONS[api_key] = {
            "client": client,
            "chat": chat
        }
        log("✅ Created new Gemini session with multi-turn chat.")
        return client, chat
    except Exception as e:
        log(f"❌ Error creating Gemini client: {e}")
        return None, None


def insert_citations_from_grounding(candidates):
    """Extract citations from grounding metadata - using chunk titles as display names"""
    try:
        if not candidates:
            log("⚠️ No candidates for citation extraction")
            return None
        
        cand = candidates[0]
        
        # Check if grounding metadata exists
        grounding = getattr(cand, "grounding_metadata", None)
        if not grounding:
            log("⚠️ No grounding_metadata found")
            return None
        
        # Get chunks
        chunks = getattr(grounding, "grounding_chunks", None) or []
        
        if not chunks:
            log("⚠️ No grounding_chunks found")
            return None
        
        # Build citation list from chunks
        citations = []
        seen_titles = set()
        
        for idx, chunk in enumerate(chunks):
            if hasattr(chunk, 'web') and chunk.web:
                uri = getattr(chunk.web, "uri", None)
                title = getattr(chunk.web, "title", None)
                
                # Use title as the clickable text since it shows the actual domain
                if uri and title and title not in seen_titles:
                    seen_titles.add(title)
                    citations.append(f"[{title}]({uri})")
                elif uri:
                    citations.append(f"[Source {idx+1}]({uri})")
        
        if citations:
            citation_text = "\n\n📚 **Sources:** " + " • ".join(citations)
            log(f"✅ Created {len(citations)} citations with source domains")
            return citation_text
        else:
            log("⚠️ No valid citations could be created")
            return None
            
    except Exception as e:
        log(f"⚠️ Citation extraction failed: {e}")
        traceback.print_exc()
        return None


def reasoning_budget(level: str) -> int:
    level = (level or "Dynamic").lower()
    if level == "none":
        return 0
    elif level == "concise":
        return 256
    elif level == "strong":
        return 2048
    elif level == "dynamic":
        return -1
    return -1


def chat_with_gemini(api_key, chat_history_msgs, multimodal_input, show_thoughts, reasoning_level, enabled_extensions):
    log("=== chat_with_gemini CALLED ===")

    if not api_key:
        chat_history_msgs = chat_history_msgs or []
        chat_history_msgs.append({
            "role": "assistant", 
            "content": "🔑 Please enter your Gemini API key first."
        })
        yield chat_history_msgs
        return

    client, chat = get_or_create_session(api_key)
    if not client:
        chat_history_msgs.append({
            "role": "assistant", 
            "content": "⚠️ Could not create Gemini session."
        })
        yield chat_history_msgs
        return

    user_text = (multimodal_input or {}).get("text", "") or ""
    uploaded_files = (multimodal_input or {}).get("files", []) or []

    if chat_history_msgs is None:
        chat_history_msgs = []
    
    # Process uploaded files
    file_parts = []
    if uploaded_files:
        log(f"📎 Processing {len(uploaded_files)} uploaded file(s)...")
        for file_path in uploaded_files:
            try:
                file_part = process_uploaded_file(client, file_path)
                file_parts.append(file_part)
            except Exception as e:
                log(f"❌ Error processing file {file_path}: {e}")
                traceback.print_exc()
    
    chat_history_msgs.append({"role": "user", "content": user_text})
    yield chat_history_msgs

    assistant_base_index = len(chat_history_msgs)
    
    # Setup thinking display if enabled
    if show_thoughts:
        thought_index = assistant_base_index
        chat_history_msgs.append({"role": "assistant", "content": "<em>💭 Thinking...</em>"})
        answer_index = thought_index + 1
        chat_history_msgs.append({"role": "assistant", "content": "🤔 Processing..."})
    else:
        thought_index = None
        answer_index = assistant_base_index
        chat_history_msgs.append({"role": "assistant", "content": "🤔 Processing..."})
    
    yield chat_history_msgs

    try:
        # Initialize variables at function scope
        search_citations = None
        
        # AGENT ORCHESTRATION APPROACH
        if enabled_extensions:
            log("🎭 Using multi-agent orchestration with multi-turn chat")
            orchestrator = AgentOrchestrator(client, chat, EXTENSION_MANAGER, enabled_extensions)
            
            budget = reasoning_budget(reasoning_level)
            thoughts_accumulated = ""
            
            # Step 1: Determine if search is needed
            needs_search = determine_needs_search(chat, user_text)
            log(f"📊 Search needed: {needs_search}")
            
            # Step 2: Call search agent if needed
            search_results = ""
            if needs_search:
                chat_history_msgs[answer_index]["content"] = "🔍 Searching the web..."
                yield chat_history_msgs
                
                search_results, search_citations = orchestrator.call_search_agent(user_text, file_parts)
                log(f"📋 After search: search_citations = {search_citations[:100] if search_citations else 'None'}")
                
                if search_results:
                    chat_history_msgs[answer_index]["content"] = "✅ Found information online\n\n🛠️ Now processing with tools..."
                    yield chat_history_msgs
            
            # Step 3: Call tool agent (with files)
            function_calls, tool_response, tool_thoughts = orchestrator.call_tool_agent(
                user_text, search_results, budget, file_parts
            )
            
            # Show thoughts if available
            if tool_thoughts and show_thoughts:
                thoughts_accumulated += tool_thoughts
                chat_history_msgs[thought_index]["content"] = (
                    f"<details open>"  
                    f"<summary><strong>💭 GemiWine's Thinking</strong></summary>"
                    f"<div style='white-space:pre-wrap;background:inherit;color:inherit;"
                    f"padding:8px;border-radius:8px;border:1px solid var(--border-color);'>"
                    f"{thoughts_accumulated.strip()}</div>"
                    f"</details>"
                )
                yield chat_history_msgs
            
            # Step 4: Execute function calls if any
            tool_results = []
            if function_calls:
                chat_history_msgs[answer_index]["content"] = "⚙️ Executing tools..."
                yield chat_history_msgs
                
                user_id = api_key
                results = EXTENSION_MANAGER.handle_function_calls(
                    user_id, enabled_extensions, function_calls
                )
                
                for fc, result in zip(function_calls, results):
                    tool_results.append((fc.name, result))
                    log(f"✅ {fc.name}: {result}")
            
            # Step 5: Synthesize final response
            if search_results or tool_results or tool_response:
                chat_history_msgs[answer_index]["content"] = "✨ Synthesizing answer..."
                yield chat_history_msgs
                final_answer, generated_images = orchestrator.synthesize_response(user_text, search_results, tool_results, search_citations, file_parts)
            else:
                final_answer = tool_response or "I couldn't process that request."
                generated_images = []
            
            # Build the final content with citations if available
            final_content = (
                f"<div><strong>🍇 Final Answer</strong>"
                f"<div style='white-space:pre-wrap;background:inherit;color:inherit;"
                f"padding:8px;border-radius:8px;border:1px solid var(--border-color);'>"
                f"{final_answer.strip()}</div></div>"
            )
            
            # Add generated images/charts
            if generated_images:
                log(f"📊 Adding {len(generated_images)} generated visualizations to response")
                for img_data in generated_images:
                    final_content += f"\n\n<div style='margin-top:16px;'>"
                    final_content += f"<strong>📊 {img_data['title']}</strong><br/>"
                    final_content += f"<img src='data:image/png;base64,{img_data['base64']}' style='max-width:100%;border-radius:8px;box-shadow:0 2px 8px rgba(0,0,0,0.1);'/>"
                    if img_data['filepath']:
                        final_content += f"<br/><small style='color:#666;'>Saved to: {img_data['filepath']}</small>"
                    final_content += "</div>"
            
            # Append citations if they exist
            if search_citations:
                final_content += "\n\n" + search_citations
                log(f"✅ Appended citations to final answer")
            
            chat_history_msgs[answer_index]["content"] = final_content
            yield chat_history_msgs
            
        else:
            # No extensions - simple streaming with search
            log("📺 Using simple streaming mode")
            
            # Build parts for message with files
            parts = []
            if file_parts:
                parts.extend(file_parts)
            parts.append(user_text)
            
            budget = reasoning_budget(reasoning_level)
            grounding_tool = types.Tool(google_search=types.GoogleSearch())
            config = types.GenerateContentConfig(
                system_instruction=BASE_SYSTEM_PROMPT,
                tools=[grounding_tool],
                temperature=0.7,
                top_p=0.9,
                max_output_tokens=8192,
                thinking_config=types.ThinkingConfig(
                    include_thoughts=True,
                    thinking_budget=budget,
                )
            )
            
            stream = chat.send_message_stream(parts, config=config)
            
            answer = ""
            thoughts = ""
            last_chunk = None
            
            # Add thinking placeholder if needed
            if show_thoughts:
                thought_index = answer_index
                chat_history_msgs[answer_index]["content"] = "<em>💭 Thinking...</em>"
                answer_index = len(chat_history_msgs)
                chat_history_msgs.append({"role": "assistant", "content": ""})
                yield chat_history_msgs
            
            for chunk in stream:
                last_chunk = chunk
                if not getattr(chunk, "candidates", None):
                    continue
                candidate = chunk.candidates[0]
                
                if getattr(candidate, "content", None):
                    for part in candidate.content.parts:
                        if not getattr(part, "text", None):
                            continue
                        
                        if getattr(part, "thought", False):
                            thoughts += part.text
                            if show_thoughts:
                                chat_history_msgs[thought_index]["content"] = (
                                    f"<details open>"  
                                    f"<summary><strong>💭 GemiWine's Thinking</strong></summary>"
                                    f"<div style='white-space:pre-wrap;background:inherit;color:inherit;"
                                    f"padding:8px;border-radius:8px;border:1px solid var(--border-color);'>"
                                    f"{thoughts.strip()}</div>"
                                    f"</details>"
                                )
                                yield chat_history_msgs
                        else:
                            answer += part.text
                            chat_history_msgs[answer_index]["content"] = (
                                f"<div><strong>🍇 Final Answer</strong>"
                                f"<div style='white-space:pre-wrap;background:inherit;color:inherit;"
                                f"padding:8px;border-radius:8px;border:1px solid var(--border-color);'>"
                                f"{answer.strip()}</div></div>"
                            )
                            yield chat_history_msgs
            
            # Add citations
            if last_chunk:
                citations = insert_citations_from_grounding(last_chunk.candidates)
                if citations:
                    chat_history_msgs[answer_index]["content"] += "\n\n" + citations
                    yield chat_history_msgs
        
        log("✅ Response complete.")
        return
        
    except Exception as e:
        log(f"❌ Error: {e}")
        traceback.print_exc()
        chat_history_msgs[answer_index]["content"] = f"⚠️ Error: {e}"
        yield chat_history_msgs
        return


def build_extension_ui():
    """Build the extension toggle UI"""
    extensions = EXTENSION_MANAGER.get_all_extensions()
    
    if not extensions:
        return gr.Markdown("No extensions available"), []
    
    checkboxes = []
    with gr.Accordion("🔌 Agent Extensions", open=True):
        gr.Markdown("Enable extensions to give the agent additional capabilities:")
        gr.Markdown("✨ **Agentic Mode:** When extensions are enabled, the agent uses multi-step reasoning with search + tools")
        for ext in extensions:
            cb = gr.Checkbox(
                label=f"{ext.icon} {ext.display_name}",
                info=ext.description,
                value=False
            )
            checkboxes.append((ext.name, cb))
    
    return checkboxes


with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="purple", secondary_hue="blue"),
    title="GemiWine",
    fill_width=True
) as demo:
    gr.HTML("""
    <style>
        .gradio-container { padding-top: 1.5rem; padding-bottom: 1.5rem; }
        .chat-panel {
            background: rgba(255, 255, 255, 0.05);
            border-radius: 16px !important;
            padding: 1.5rem;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
            border: 1px solid rgba(255, 255, 255, 0.1);
        }
        .message-input {
            border-radius: 12px !important;
            border: 1px solid rgba(0,0,0,0.1);
        }
    </style>
    """)

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            gr.Markdown("## ⚙️ Settings & Controls")
            api_key = gr.Textbox(
                label="🔑 Gemini API Key",
                placeholder="Paste your Gemini API key here...",
                type="password",
            )
            reasoning_level = gr.Radio(
                ["None", "Concise", "Strong", "Dynamic"],
                label="🧠 Reasoning Level",
                value="Dynamic",
                info="Controls the model's thinking depth.",
            )
            show_thoughts = gr.Checkbox(
                label="💭 Show Thinking",
                value=True,
                info="Display reasoning process before answers.",
            )
            
            # Build extension checkboxes
            extension_checkboxes = build_extension_ui()

        with gr.Column(scale=4):
            with gr.Group(elem_classes="chat-panel"):
                chatbot = gr.Chatbot(
                    label="🍇 Chat with GemiWine",
                    height=650,
                    show_copy_button=True,
                    type="messages",  
                    avatar_images=(None, "https://i.imgur.com/Q2EMk2N.png"),
                )
                multimodal_msg = gr.MultimodalTextbox(
                    file_types=[
                        "image", "video", "audio",  # Gradio presets
                        ".pdf", ".txt", ".md", ".html", ".xml",  # Documents
                        ".doc", ".docx", ".csv", ".json"  # Additional formats
                    ],
                    placeholder="Ask anything, upload images/PDFs/videos, or let extensions help you...",
                    label="Your Message",
                    elem_classes="message-input",
                    autofocus=True
                )
    
    # Hidden state to track enabled extensions
    enabled_extensions_state = gr.State([])

    def clear_box():
        return {"text": "", "files": []}

    def handle_chat(api_key_input, chat_history_msgs, multimodal_dict, thinking_flag, reasoning_lvl, *extension_states):
        # Convert extension checkbox states to list of enabled extension names
        enabled = []
        for (ext_name, _), is_enabled in zip(extension_checkboxes, extension_states):
            if is_enabled:
                enabled.append(ext_name)
        
        log(f"Enabled extensions: {enabled}")
        
        yield from chat_with_gemini(
            api_key_input, chat_history_msgs, multimodal_dict, 
            thinking_flag, reasoning_lvl, enabled
        )

    def check_timers(api_key_input, chat_history, enabled_exts):
        """Background function to check for completed timers"""
        if not api_key_input or 'timer' not in enabled_exts:
            return chat_history
        
        timer_ext = EXTENSION_MANAGER.get_extension('timer')
        if not timer_ext:
            return chat_history
        
        user_id = api_key_input
        timer_ext.initialize_state(user_id)
        state = timer_ext.get_state(user_id)
        
        import datetime as dt
        now = dt.datetime.now()
        newly_completed = []
        
        for timer in state.get("timers", []):
            if timer.get("active") and not timer.get("notified", False):
                end_time = dt.datetime.fromisoformat(timer["end_time"])
                if now >= end_time:
                    newly_completed.append(timer)
                    timer["notified"] = True
        
        if newly_completed:
            timer_ext.update_state(user_id, state)
            
            # Add notification to chat
            if chat_history is None:
                chat_history = []
            
            for timer in newly_completed:
                notification = f"⏰ **Timer Complete!** Your timer '{timer['name']}' has finished!"
                chat_history.append({"role": "assistant", "content": notification})
                log(f"⏰ Timer notification sent: {timer['name']}")
        
        return chat_history

    # Get just the checkbox components for inputs
    checkbox_components = [cb for _, cb in extension_checkboxes]
    
    # Main chat submission
    multimodal_msg.submit(
        fn=handle_chat,
        inputs=[api_key, chatbot, multimodal_msg, show_thoughts, reasoning_level] + checkbox_components,
        outputs=[chatbot],
        queue=True,
    ).then(fn=clear_box, outputs=[multimodal_msg])
    
    # Background timer check - runs every 10 seconds
    timer_check = gr.Timer(value=10, active=True)
    
    def update_enabled_state(*extension_states):
        enabled = []
        for (ext_name, _), is_enabled in zip(extension_checkboxes, extension_states):
            if is_enabled:
                enabled.append(ext_name)
        return enabled
    
    # Update enabled extensions state whenever checkboxes change
    for _, cb in extension_checkboxes:
        cb.change(
            fn=update_enabled_state,
            inputs=checkbox_components,
            outputs=[enabled_extensions_state]
        )
    
    # Timer polling
    timer_check.tick(
        fn=check_timers,
        inputs=[api_key, chatbot, enabled_extensions_state],
        outputs=[chatbot]
    )


if __name__ == "__main__":
    log(f"===== GemiWine with Extensions started at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} =====")
    demo.launch()