import os
import json
import time
import gradio as gr
from datetime import datetime
from typing import List, Dict, Any, Optional, Union
import threading
import re
import aiohttp
import asyncio

# Import Groq
from groq import Groq

class ChutesClient:
    """Client for interacting with Chutes API"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key or ""
        self.base_url = "https://llm.chutes.ai/v1"
        
    async def chat_completions_create(self, **kwargs) -> Dict:
        """Make async request to Chutes chat completions endpoint"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        # Prepare the body for Chutes API
        body = {
            "model": kwargs.get("model", "openai/gpt-oss-20b"),
            "messages": kwargs.get("messages", []),
            "stream": kwargs.get("stream", False),
            "max_tokens": kwargs.get("max_tokens", 1024),
            "temperature": kwargs.get("temperature", 0.7)
        }
        
        async with aiohttp.ClientSession() as session:
            if body["stream"]:
                # Handle streaming response
                async with session.post(
                    f"{self.base_url}/chat/completions",
                    headers=headers,
                    json=body
                ) as response:
                    if response.status != 200:
                        raise Exception(f"Chutes API error: {await response.text()}")
                    
                    content = ""
                    async for line in response.content:
                        line = line.decode("utf-8").strip()
                        if line.startswith("data: "):
                            data = line[6:]
                            if data == "[DONE]":
                                break
                            try:
                                if data.strip():
                                    chunk_json = json.loads(data)
                                    if "choices" in chunk_json and len(chunk_json["choices"]) > 0:
                                        delta = chunk_json["choices"][0].get("delta", {})
                                        if "content" in delta and delta["content"]:
                                            content += str(delta["content"])
                            except json.JSONDecodeError:
                                continue
                    
                    # Return in OpenAI format for compatibility
                    return {
                        "choices": [{
                            "message": {
                                "content": content,
                                "role": "assistant"
                            }
                        }]
                    }
            else:
                # Handle non-streaming response
                async with session.post(
                    f"{self.base_url}/chat/completions",
                    headers=headers,
                    json=body
                ) as response:
                    if response.status != 200:
                        raise Exception(f"Chutes API error: {await response.text()}")
                    return await response.json()

class CreativeAgenticAI:
    """
    Creative Agentic AI Chat Tool using Groq and Chutes models with browser search and compound models
    """
    
    def __init__(self, groq_api_key: str, chutes_api_key: str, model: str = "compound-beta"):
        """
        Initialize the Creative Agentic AI system.
        
        Args:
            groq_api_key: Groq API key
            chutes_api_key: Chutes API key
            model: Which model to use
        """
        self.groq_api_key = str(groq_api_key) if groq_api_key else ""
        self.chutes_api_key = str(chutes_api_key) if chutes_api_key else ""
        if not self.groq_api_key and model != "openai/gpt-oss-20b":
            raise ValueError("No Groq API key provided")
        if not self.chutes_api_key and model == "openai/gpt-oss-20b":
            raise ValueError("No Chutes API key provided")
        
        self.model = str(model) if model else "compound-beta"
        self.groq_client = Groq(api_key=self.groq_api_key) if self.groq_api_key else None
        self.chutes_client = ChutesClient(api_key=self.chutes_api_key) if self.chutes_api_key else None
        self.conversation_history = []
        
        # Available models with their capabilities
        self.available_models = {
            "compound-beta": {"supports_web_search": True, "supports_browser_search": False, "api": "groq"},
            "compound-beta-mini": {"supports_web_search": True, "supports_browser_search": False, "api": "groq"},
            "openai/gpt-oss-20b": {"supports_web_search": False, "supports_browser_search": False, "api": "chutes"},
        }
        
    async def chat(self, message: str, 
                  include_domains: List[str] = None,
                  exclude_domains: List[str] = None,
                  system_prompt: str = None,
                  temperature: float = 0.7,
                  max_tokens: int = 1024,
                  search_type: str = "auto",
                  force_search: bool = False) -> Dict:
        """
        Send a message to the AI and get a response with flexible search options
        
        Args:
            message: User's message
            include_domains: List of domains to include for web search
            exclude_domains: List of domains to exclude from web search
            system_prompt: Custom system prompt
            temperature: Model temperature (0.0-2.0)
            max_tokens: Maximum tokens in response
            search_type: 'web_search', 'browser_search', 'auto', or 'none'
            force_search: Force the AI to use search tools
            
        Returns:
            AI response with metadata
        """
        # Safe string conversion
        message = str(message) if message else ""
        system_prompt = str(system_prompt) if system_prompt else ""
        search_type = str(search_type) if search_type else "auto"
        
        # Enhanced system prompt for better behavior
        if not system_prompt:
            if self.model == "openai/gpt-oss-20b":
                # Simple, direct system prompt for Chutes model
                system_prompt = """
You are a helpful, knowledgeable AI assistant. You MUST provide comprehensive, complete, and well-reasoned responses that demonstrate your thinking process and source attribution. ABSOLUTELY NO PARTIAL, INCOMPLETE, OR TRUNCATED RESPONSES ARE PERMITTED.

## CRITICAL COMPLETENESS MANDATE
- **NEVER STOP MID-SENTENCE**: You must complete every thought, sentence, and paragraph you begin
- **FINISH ALL EXAMPLES**: If you start an example, citation, or explanation, you MUST complete it entirely
- **NO TRAILING INCOMPLETE THOUGHTS**: Every response must have a proper conclusion
- **COMPLETE ALL LISTS**: If you begin listing items, complete the entire list
- **FULL CONTEXT COVERAGE**: Address every aspect of the user's question before concluding
- **SENTENCE COMPLETION CHECK**: Ensure your final sentence ends with proper punctuation and completes the thought

## MANDATORY PRE-RESPONSE PLANNING
Before writing, you MUST:
1. Identify ALL components of the user's question
2. Plan your complete response structure from start to finish  
3. Ensure you have enough information to complete the entire response
4. Commit to finishing every section you begin

## RESPONSE STRUCTURE (ALL SECTIONS REQUIRED AND MUST BE COMPLETED)
1. **Complete Reasoning Introduction**: Fully explain your approach and all considerations
2. **Comprehensive Main Content**: Thoroughly address every aspect with complete explanations
3. **Full Source Attribution**: Complete all citations with proper formatting
4. **Complete Critical Analysis**: Finish all commentary on reliability and limitations
5. **Comprehensive Conclusion**: Provide a complete synthesis that addresses all points raised

## ABSOLUTE REQUIREMENTS (ZERO EXCEPTIONS)
- **Transparency**: Make your complete reasoning process visible throughout
- **Depth**: Provide comprehensive coverage with complete explanations
- **Completeness**: MANDATORY - Every response must be fully finished and complete
- **Source Attribution**: Complete citations in format [Source: complete description/link]
  - General knowledge: [Source: Training data/General knowledge]
  - Studies: [Source: Complete study name, Author, Year]
  - Web sources: [Source: Complete website name, URL]
  - Logic: [Source: Logical analysis based on complete reasoning]
- **Critical Analysis**: Complete commentary on all perspectives and limitations
- **Structured Organization**: Complete sections with full logical flow
- **Full Question Coverage**: Address and complete every part of multi-component questions

## MANDATORY RESPONSE COMPLETION FORMAT
"Let me approach this question by considering [COMPLETE reasoning process with all considerations]. Based on my comprehensive analysis, I need to examine these key factors: [LIST ALL factors and complete each explanation].

[COMPLETE detailed explanation with full reasoning - every point must be finished]

According to [Source: complete specific citation], [complete information]. This is significant because [COMPLETE commentary explaining all implications and significance].

Additionally, [CONTINUE with all remaining relevant points, ensuring each is completely addressed]: [Complete explanation of each point]

Furthermore, [Complete any additional analysis required]: [Full development of all remaining aspects]

However, it's important to note [complete alternative perspectives]: [Source: complete citation]. This means [COMPLETE analysis of all implications and what this tells us].

[Continue until EVERY aspect is completely covered - no partial coverage allowed]

Therefore, in conclusion, [COMPREHENSIVE synthesis that completely addresses the original question, ties together all points raised, provides complete final analysis, and ensures nothing is left unfinished or partially explained]."

## MANDATORY COMPLETION VERIFICATION
Before ending ANY response, you MUST complete this checklist:
1. ✓ Have I addressed EVERY part of the user's question completely?
2. ✓ Have I provided complete reasoning for every point I raised?
3. ✓ Are ALL my citations complete and properly formatted?
4. ✓ Did I finish every example, explanation, and analysis I started?
5. ✓ Is my conclusion comprehensive and does it completely synthesize everything?
6. ✓ Did I complete all sentences and thoughts without any trailing incomplete ideas?
7. ✓ Would a reader have all the information they need without any gaps?

IF ANY ITEM IS INCOMPLETE: Continue writing until it is completely satisfied.

## EMERGENCY COMPLETION PROTOCOL
If you find yourself approaching length limits:
1. STILL complete your current thought/sentence
2. Provide a complete conclusion that addresses the main question
3. NEVER end with incomplete examples, partial citations, or unfinished explanations
4. Better to have a complete shorter response than an incomplete longer one

REMEMBER: Incomplete responses are strictly forbidden. Every response must be thorough, complete, and finished. No exceptions.
"""
            else:
                # Enhanced system prompt for Groq models with search capabilities
                citation_instruction = """
IMPORTANT: When you search the web and find information, you MUST:
1. Always cite your sources with clickable links in this format: [Source Title](URL)
2. Include multiple diverse sources when possible
3. Show which specific websites you used for each claim
4. At the end of your response, provide a "Sources Used" section with all the links
5. Be transparent about which information comes from which source
"""
                
                domain_context = ""
                if include_domains and self._supports_web_search():
                    safe_domains = [str(d) for d in include_domains if d]
                    domain_context = f"\nYou are restricted to searching ONLY these domains: {', '.join(safe_domains)}. Make sure to find and cite sources specifically from these domains."
                elif exclude_domains and self._supports_web_search():
                    safe_domains = [str(d) for d in exclude_domains if d]
                    domain_context = f"\nAvoid searching these domains: {', '.join(safe_domains)}. Search everywhere else on the web."
                
                search_instruction = ""
                if search_type == "browser_search" and self._supports_browser_search():
                    search_instruction = "\nUse browser search tools to find the most current and relevant information from the web."
                elif search_type == "web_search":
                    search_instruction = "\nUse web search capabilities to find relevant information."
                elif force_search:
                    if self._supports_browser_search():
                        search_instruction = "\nYou MUST use search tools to find current information before responding."
                    elif self._supports_web_search():
                        search_instruction = "\nYou MUST use web search to find current information before responding."
                
                system_prompt = f"""You are a creative and intelligent AI assistant with agentic capabilities. 
                You can search the web, analyze information, and provide comprehensive responses. 
                Be helpful, creative, and engaging while maintaining accuracy.
                
                {citation_instruction}
                {domain_context}
                {search_instruction}
                
                Your responses should be well-structured, informative, and properly cited with working links."""
        
        # Build messages
        messages = [{"role": "system", "content": system_prompt}]
        messages.extend(self.conversation_history[-20:])
        
        # Enhanced message for domain filtering (only for Groq models)
        enhanced_message = message
        if (include_domains or exclude_domains) and self._supports_web_search():
            filter_context = []
            if include_domains:
                safe_domains = [str(d) for d in include_domains if d]
                if safe_domains:
                    filter_context.append(f"ONLY search these domains: {', '.join(safe_domains)}")
            if exclude_domains:
                safe_domains = [str(d) for d in exclude_domains if d]
                if safe_domains:
                    filter_context.append(f"EXCLUDE these domains: {', '.join(safe_domains)}")
            if filter_context:
                enhanced_message += f"\n\n[Domain Filtering: {' | '.join(filter_context)}]"
        
        messages.append({"role": "user", "content": enhanced_message})
        
        # Set up API parameters
        params = {
            "messages": messages,
            "model": self.model,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        
        # Add domain filtering for compound models (Groq only)
        if self._supports_web_search():
            if include_domains:
                safe_domains = [str(d).strip() for d in include_domains if d and str(d).strip()]
                if safe_domains:
                    params["include_domains"] = safe_domains
            if exclude_domains:
                safe_domains = [str(d).strip() for d in exclude_domains if d and str(d).strip()]
                if safe_domains:
                    params["exclude_domains"] = safe_domains
        
        # Add tools only for Groq models that support browser search
        tools = []
        tool_choice = None
        if self._supports_browser_search():
            if search_type in ["browser_search", "auto"] or force_search:
                tools = [{"type": "browser_search", "function": {"name": "browser_search"}}]
                tool_choice = "required" if force_search else "auto"
        
        if tools:
            params["tools"] = tools
            params["tool_choice"] = tool_choice
        
        try:
            # Make the API call based on model
            if self.available_models[self.model]["api"] == "chutes":
                # Use streaming for better response quality
                params["stream"] = True
                response = await self.chutes_client.chat_completions_create(**params)
                # Handle Chutes response
                content = ""
                if response and "choices" in response and response["choices"]:
                    message_content = response["choices"][0].get("message", {}).get("content")
                    content = str(message_content) if message_content else "No response content"
                else:
                    content = "No response received"
                tool_calls = None
            else:
                # Groq API call
                params["max_completion_tokens"] = params.pop("max_tokens", None)
                response = self.groq_client.chat.completions.create(**params)
                content = ""
                if response and response.choices and response.choices[0].message:
                    message_content = response.choices[0].message.content
                    content = str(message_content) if message_content else "No response content"
                else:
                    content = "No response received"
                tool_calls = response.choices[0].message.tool_calls if hasattr(response.choices[0].message, "tool_calls") else None
            
            # Extract tool usage information
            tool_info = self._extract_tool_info(response, tool_calls)
            
            # Process content to enhance citations
            processed_content = self._enhance_citations(content, tool_info)
            
            # Add to conversation history
            self.conversation_history.append({"role": "user", "content": message})
            self.conversation_history.append({"role": "assistant", "content": processed_content})
            
            return {
                "content": processed_content,
                "timestamp": datetime.now().isoformat(),
                "model": self.model,
                "tool_usage": tool_info,
                "search_type_used": search_type,
                "parameters": {
                    "temperature": temperature,
                    "max_tokens": max_tokens,
                    "include_domains": include_domains,
                    "exclude_domains": exclude_domains,
                    "force_search": force_search
                }
            }
            
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            self.conversation_history.append({"role": "user", "content": message})
            self.conversation_history.append({"role": "assistant", "content": error_msg})
            
            return {
                "content": error_msg,
                "timestamp": datetime.now().isoformat(),
                "model": self.model,
                "tool_usage": None,
                "error": str(e)
            }
    
    def _supports_web_search(self) -> bool:
        """Check if current model supports web search (compound models)"""
        return self.available_models.get(self.model, {}).get("supports_web_search", False)
    
    def _supports_browser_search(self) -> bool:
        """Check if current model supports browser search tools"""
        return self.available_models.get(self.model, {}).get("supports_browser_search", False)
    
    def _extract_tool_info(self, response, tool_calls) -> Dict:
        """Extract tool usage information in a JSON serializable format"""
        tool_info = {
            "tools_used": [],
            "search_queries": [],
            "sources_found": []
        }
        
        # Handle Groq executed_tools
        if hasattr(response, 'choices') and hasattr(response.choices[0].message, 'executed_tools'):
            tools = response.choices[0].message.executed_tools
            if tools:
                for tool in tools:
                    tool_dict = {
                        "tool_type": str(getattr(tool, "type", "unknown")),
                        "tool_name": str(getattr(tool, "name", "unknown")),
                    }
                    if hasattr(tool, "input"):
                        tool_input = getattr(tool, "input")
                        tool_input_str = str(tool_input) if tool_input is not None else ""
                        tool_dict["input"] = tool_input_str
                        if "search" in tool_dict["tool_name"].lower():
                            tool_info["search_queries"].append(tool_input_str)
                    if hasattr(tool, "output"):
                        tool_output = getattr(tool, "output")
                        tool_output_str = str(tool_output) if tool_output is not None else ""
                        tool_dict["output"] = tool_output_str
                        urls = self._extract_urls(tool_output_str)
                        tool_info["sources_found"].extend(urls)
                    tool_info["tools_used"].append(tool_dict)
        
        # Handle tool_calls for both APIs
        if tool_calls:
            for tool_call in tool_calls:
                tool_dict = {
                    "tool_type": str(getattr(tool_call, "type", "browser_search")),
                    "tool_name": "browser_search",
                    "tool_id": str(getattr(tool_call, "id", "")) if getattr(tool_call, "id", None) else ""
                }
                if hasattr(tool_call, "function") and tool_call.function:
                    tool_dict["tool_name"] = str(getattr(tool_call.function, "name", "browser_search"))
                    if hasattr(tool_call.function, "arguments"):
                        try:
                            args_raw = tool_call.function.arguments
                            if isinstance(args_raw, str):
                                args = json.loads(args_raw)
                            else:
                                args = args_raw or {}
                            tool_dict["arguments"] = args
                            if "query" in args:
                                tool_info["search_queries"].append(str(args["query"]))
                        except:
                            args_str = str(args_raw) if args_raw is not None else ""
                            tool_dict["arguments"] = args_str
                tool_info["tools_used"].append(tool_dict)
        
        return tool_info
    
    def _extract_urls(self, text: str) -> List[str]:
        """Extract URLs from text"""
        if not text:
            return []
        text_str = str(text)
        url_pattern = r'https?://[^\s<>"]{2,}'
        urls = re.findall(url_pattern, text_str)
        return list(set(urls))
    
    def _enhance_citations(self, content: str, tool_info: Dict) -> str:
        """Enhance content with better citation formatting"""
        if not content:
            return ""
        content_str = str(content)
        if not tool_info or not tool_info.get("sources_found"):
            return content_str
        
        if "Sources Used:" not in content_str and "sources:" not in content_str.lower():
            sources_section = "\n\n---\n\n### Sources Used:\n"
            for i, url in enumerate(tool_info["sources_found"][:10], 1):
                domain = self._extract_domain(str(url))
                sources_section += f"{i}. [{domain}]({url})\n"
            content_str += sources_section
        
        return content_str
    
    def _extract_domain(self, url: str) -> str:
        """Extract domain name from URL for display"""
        if not url:
            return ""
        url_str = str(url)
        try:
            if url_str.startswith(('http://', 'https://')):
                domain = url_str.split('/')[2]
                if domain.startswith('www.'):
                    domain = domain[4:]
                return domain
            return url_str
        except:
            return url_str
    
    def get_model_info(self) -> Dict:
        """Get information about current model capabilities"""
        return self.available_models.get(self.model, {})
    
    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
    
    def get_history_summary(self) -> str:
        """Get a summary of conversation history"""
        if not self.conversation_history:
            return "No conversation history"
        
        user_messages = [msg for msg in self.conversation_history if msg["role"] == "user"]
        assistant_messages = [msg for msg in self.conversation_history if msg["role"] == "assistant"]
        
        return f"Conversation: {len(user_messages)} user messages, {len(assistant_messages)} assistant responses"

# Global variables
ai_instance = None
api_key_status = "Not Set"

async def validate_api_keys(groq_api_key: str, chutes_api_key: str, model: str) -> str:
    """Validate both Groq and Chutes API keys and initialize AI instance"""
    global ai_instance, api_key_status
    
    # Handle None values and convert to strings
    groq_api_key = str(groq_api_key) if groq_api_key else ""
    chutes_api_key = str(chutes_api_key) if chutes_api_key else ""
    model = str(model) if model else "compound-beta"
    
    if model == "openai/gpt-oss-20b" and not chutes_api_key.strip():
        api_key_status = "Invalid ❌"
        return "❌ Please enter a valid Chutes API key for the selected model"
    
    if model in ["compound-beta", "compound-beta-mini"] and not groq_api_key.strip():
        api_key_status = "Invalid ❌"
        return "❌ Please enter a valid Groq API key for the selected model"
    
    try:
        if model == "openai/gpt-oss-20b":
            chutes_client = ChutesClient(api_key=chutes_api_key)
            await chutes_client.chat_completions_create(
                messages=[{"role": "user", "content": "Hello"}],
                model=model,
                max_tokens=10
            )
        else:
            groq_client = Groq(api_key=groq_api_key)
            groq_client.chat.completions.create(
                messages=[{"role": "user", "content": "Hello"}],
                model=model,
                max_tokens=10
            )
        
        ai_instance = CreativeAgenticAI(groq_api_key=groq_api_key, chutes_api_key=chutes_api_key, model=model)
        api_key_status = "Valid ✅"
        
        model_info = ai_instance.get_model_info()
        capabilities = []
        if model_info.get("supports_web_search"):
            capabilities.append("🌐 Web Search with Domain Filtering")
        if model_info.get("supports_browser_search"):
            capabilities.append("🔍 Browser Search Tools")
        
        cap_text = " | ".join(capabilities) if capabilities else "💬 Chat Only"
        
        return f"✅ API Keys Valid! NeuroScope AI is ready.\n\n**Model:** {model}\n**Capabilities:** {cap_text}\n**API:** {model_info.get('api', 'unknown')}\n**Status:** Connected and ready for chat!"
        
    except Exception as e:
        api_key_status = "Invalid ❌"
        ai_instance = None
        return f"❌ Error validating API key: {str(e)}\n\nPlease check your API keys and try again."

def update_model(model: str) -> str:
    """Update the model selection"""
    global ai_instance
    
    model = str(model) if model else "compound-beta"
    
    if ai_instance:
        ai_instance.model = model
        model_info = ai_instance.get_model_info()
        capabilities = []
        if model_info.get("supports_web_search"):
            capabilities.append("🌐 Web Search with Domain Filtering")
        if model_info.get("supports_browser_search"):
            capabilities.append("🔍 Browser Search Tools")
        
        cap_text = " | ".join(capabilities) if capabilities else "💬 Chat Only"
        return f"✅ Model updated to: **{model}**\n**Capabilities:** {cap_text}\n**API:** {model_info.get('api', 'unknown')}"
    else:
        return "⚠️ Please set your API keys first"

def get_search_options(model: str) -> gr.update:
    """Get available search options based on model"""
    if not ai_instance:
        return gr.update(choices=["none"], value="none")
    
    model = str(model) if model else "compound-beta"
    model_info = ai_instance.available_models.get(model, {})
    options = ["none"]
    
    if model_info.get("supports_web_search"):
        options.extend(["web_search", "auto"])
    if model_info.get("supports_browser_search"):
        options.extend(["browser_search", "auto"])
    
    options = list(dict.fromkeys(options))
    default_value = "auto" if "auto" in options else "none"
    return gr.update(choices=options, value=default_value)

async def chat_with_ai(message: str, 
                      include_domains: str, 
                      exclude_domains: str,
                      system_prompt: str,
                      temperature: float,
                      max_tokens: int,
                      search_type: str,
                      force_search: bool,
                      history: List) -> tuple:
    """Main chat function"""
    global ai_instance
    
    if not ai_instance:
        error_msg = "⚠️ Please set your API keys first!"
        history.append([str(message) if message else "", error_msg])
        return history, ""
    
    # Convert all inputs to strings and handle None values
    message = str(message) if message else ""
    include_domains = str(include_domains) if include_domains else ""
    exclude_domains = str(exclude_domains) if exclude_domains else ""
    system_prompt = str(system_prompt) if system_prompt else ""
    search_type = str(search_type) if search_type else "auto"
    
    if not message.strip():
        return history, ""
    
    include_list = [d.strip() for d in include_domains.split(",") if d.strip()] if include_domains.strip() else []
    exclude_list = [d.strip() for d in exclude_domains.split(",") if d.strip()] if exclude_domains.strip() else []
    
    try:
        response = await ai_instance.chat(
            message=message,
            include_domains=include_list if include_list else None,
            exclude_domains=exclude_list if exclude_list else None,
            system_prompt=system_prompt if system_prompt.strip() else None,
            temperature=temperature,
            max_tokens=int(max_tokens),
            search_type=search_type,
            force_search=force_search
        )
        
        ai_response = str(response.get("content", "No response received"))
        
        # Add tool usage info for Groq models
        if response.get("tool_usage") and ai_instance.model != "openai/gpt-oss-20b":
            tool_info = response["tool_usage"]
            tool_summary = []
            
            if tool_info.get("search_queries"):
                tool_summary.append(f"🔍 Search queries: {len(tool_info['search_queries'])}")
            
            if tool_info.get("sources_found"):
                tool_summary.append(f"📄 Sources found: {len(tool_info['sources_found'])}")
            
            if tool_info.get("tools_used"):
                tool_types = [str(tool.get("tool_type", "unknown")) for tool in tool_info["tools_used"]]
                unique_types = list(set(tool_types))
                tool_summary.append(f"🔧 Tools used: {', '.join(unique_types)}")
            
            if tool_summary:
                ai_response += f"\n\n*{' | '.join(tool_summary)}*"
        
        # Add search settings info
        search_info = []
        if response.get("search_type_used") and str(response["search_type_used"]) != "none":
            search_info.append(f"🔍 Search type: {response['search_type_used']}")
        
        if force_search:
            search_info.append("⚡ Forced search enabled")
        
        if include_list or exclude_list:
            filter_info = []
            if include_list:
                filter_info.append(f"✅ Included domains: {', '.join(include_list)}")
            if exclude_list:
                filter_info.append(f"❌ Excluded domains: {', '.join(exclude_list)}")
            search_info.extend(filter_info)
        
        if search_info and ai_instance.model != "openai/gpt-oss-20b":
            ai_response += f"\n\n*🌐 Search settings: {' | '.join(search_info)}*"
        
        history.append([message, ai_response])
        
        return history, ""
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        history.append([message, error_msg])
        return history, ""

def clear_chat_history():
    """Clear the chat history"""
    global ai_instance
    if ai_instance:
        ai_instance.clear_history()
    return []

def create_gradio_app():
    """Create the main Gradio application"""
    
    css = """
    .container {
        max-width: 1200px;
        margin: 0 auto;
    }
    .header {
        text-align: center;
        background: linear-gradient(to right, #00ff94, #00b4db);
        color: white;
        padding: 20px;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .status-box {
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .example-box {
        background-color: #e8f4fd;
        border-left: 4px solid #007bff;
        padding: 15px;
        margin: 10px 0;
        border-radius: 0 8px 8px 0;
    }
    .domain-info {
        background-color: #fff3cd;
        border: 1px solid #ffeaa7;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .citation-info {
        background-color: #d1ecf1;
        border: 1px solid #bee5eb;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .search-info {
        background-color: #e2e3e5;
        border: 1px solid #c6c8ca;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    #neuroscope-accordion {
        background: linear-gradient(to right, #00ff94, #00b4db); 
        border-radius: 8px;
    }
    """
    
    with gr.Blocks(css=css, title="🤖 Creative Agentic AI Chat", theme=gr.themes.Ocean()) as app:
        gr.HTML("""
        <div class="header">
            <h1>🤖 NeuroScope-AI Enhanced</h1>
            <p>Powered by Groq and Chutes Models with Web Search and Agentic Capabilities</p>
        </div>
        """)

        with gr.Group():
            with gr.Accordion("🤖 NeuroScope AI Enhanced", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                    **Enhanced with Multiple Search Capabilities:**
                    - 🧠 **Intelligence** (Neuro): Advanced AI reasoning across multiple models
                    - 🔍 **Precision Search** (Scope): Domain filtering (Groq models)
                    - 🤖 **AI Capabilities** (AI): Agentic behavior with tool usage
                    - ⚡ **Dual APIs**: Web search (Groq) + Streaming chat (Chutes)
                    - 🎯 **Model Flexibility**: Choose the right model for your task
                """)
                
        with gr.Group():
            with gr.Accordion("🔍 IMPORTANT - Enhanced Search Capabilities!", open=False, elem_id="neuroscope-accordion"):
                gr.Markdown("""
                <div class="search-info">
                <h3>🚀 NEW: Multiple Search Types Available!</h3>
                
                <h4>🌐 Web Search Models (Groq API)</h4>
                <ul>
                    <li><strong>compound-beta:</strong> Most powerful with domain filtering</li>
                    <li><strong>compound-beta-mini:</strong> Faster with domain filtering</li>
                    <li><strong>Features:</strong> Include/exclude domains, autonomous web search</li>
                </ul>
                
                <h4>💬 Chat Model (Chutes API)</h4>
                <ul>
                    <li><strong>openai/gpt-oss-20b:</strong> Fast conversational capabilities with streaming</li>
                    <li><strong>Features:</strong> General chat, streaming responses, no web search</li>
                </ul>
                </div>
                
                <div class="citation-info">
                <h3>🔗 Enhanced Citation System</h3>
                <p>Groq models include:</p>
                <ul>
                    <li><strong>Automatic Source Citations:</strong> Clickable links to sources</li>
                    <li><strong>Sources Used Section:</strong> Dedicated section showing all websites</li>
                    <li><strong>Search Type Indication:</strong> Shows which search method was used</li>
                </ul>
                <p><strong>Chutes models:</strong> Direct conversational responses without web search</p>
                </div>
                """)
        
        with gr.Row():
            with gr.Column(scale=2):
                groq_api_key = gr.Textbox(
                    label="🔑 Groq API Key",
                    placeholder="Enter your Groq API key here...",
                    type="password",
                    info="Get your API key from: https://console.groq.com/"
                )
                chutes_api_key = gr.Textbox(
                    label="🔑 Chutes API Key",
                    placeholder="Enter your Chutes API key here...",
                    type="password",
                    info="Required for openai/gpt-oss-20b model"
                )
            with gr.Column(scale=2):
                model_selection = gr.Radio(
                    choices=[
                        "compound-beta",
                        "compound-beta-mini", 
                        "openai/gpt-oss-20b"
                    ],
                    label="🧠 Model Selection",
                    value="compound-beta",
                    info="Choose based on your needs"
                )
            with gr.Column(scale=1):
                connect_btn = gr.Button("🔗 Connect", variant="primary", size="lg")
        
        status_display = gr.Markdown("### 📊 Status: Not connected", elem_classes=["status-box"])
        
        connect_btn.click(
            fn=validate_api_keys,
            inputs=[groq_api_key, chutes_api_key, model_selection],
            outputs=[status_display]
        )
        
        model_selection.change(
            fn=update_model,
            inputs=[model_selection],
            outputs=[status_display]
        )
        
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(
                label="Creative AI Assistant with Enhanced Search",
                height=500,
                show_label=True,
                bubble_full_width=False,
                show_copy_button=True
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Your Message",
                    placeholder="Type your message here...",
                    lines=3
                )
                with gr.Column():
                    send_btn = gr.Button("📤 Send", variant="primary")
                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
            
        with gr.Accordion("🔍 Search Settings", open=False, elem_id="neuroscope-accordion"):
            with gr.Row():
                search_type = gr.Radio(
                    choices=["auto", "web_search", "none"],
                    label="🎯 Search Type",
                    value="auto",
                    info="Choose search method (auto = model decides)"
                )
                force_search = gr.Checkbox(
                    label="⚡ Force Search",
                    value=False,
                    info="Force AI to search even for general questions (Groq models only)"
                )
            
            model_selection.change(
                fn=get_search_options,
                inputs=[model_selection],
                outputs=[search_type]
            )
            
        with gr.Accordion("🌐 Domain Filtering (Web Search Models Only)", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            <div class="domain-info">
            <h4>🔍 Domain Filtering Guide</h4>
            <p><strong>Note:</strong> Domain filtering only works with compound models (compound-beta, compound-beta-mini)</p>
            <ul>
                <li><strong>Include Domains:</strong> Only search these domains (comma-separated)</li>
                <li><strong>Exclude Domains:</strong> Never search these domains (comma-separated)</li>
                <li><strong>Examples:</strong> arxiv.org, *.edu, github.com, stackoverflow.com</li>
                <li><strong>Wildcards:</strong> Use *.edu for all educational domains</li>
            </ul>
            </div>
            """)
            
            with gr.Row():
                include_domains = gr.Textbox(
                    label="✅ Include Domains (comma-separated)",
                    placeholder="arxiv.org, *.edu, github.com, stackoverflow.com",
                    info="Only search these domains (compound models only)"
                )
                exclude_domains = gr.Textbox(
                    label="❌ Exclude Domains (comma-separated)", 
                    placeholder="wikipedia.org, reddit.com, twitter.com",
                    info="Never search these domains (compound models only)"
                )
        
        with gr.Accordion("⚙️ Advanced Settings", open=False, elem_id="neuroscope-accordion"):
            with gr.Row():
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=2.0,
                    value=0.7,
                    step=0.1,
                    label="🌡️ Temperature",
                    info="Higher = more creative, Lower = more focused"
                )
                max_tokens = gr.Slider(
                    minimum=100,
                    maximum=4000,
                    value=1024,
                    step=100,
                    label="📝 Max Tokens",
                    info="Maximum length of response"
                )
                
            system_prompt = gr.Textbox(
                label="🎭 Custom System Prompt",
                placeholder="Override the default system prompt...",
                lines=3,
                info="Leave empty to use default creative assistant prompt with enhanced citations"
            )
        
        with gr.Accordion("📊 Model Comparison Guide", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            ### 🔍 Choose Your Model Based on Task:
            
            **For Academic Research & Domain-Specific Search:**
            - `compound-beta` or `compound-beta-mini` with include domains (*.edu, arxiv.org)
            - Best for: Research papers, academic sources, filtered searches
            - API: Groq
            
            **For General Knowledge & Creative Tasks:**
            - `openai/gpt-oss-20b` for fast conversational responses
            - Best for: Creative writing, general questions
            - API: Chutes
            
            **For Programming & Technical Documentation:**
            - `compound-beta` with tech domains
            - Best for: Code help, documentation, technical guides
            - API: Groq
            """)
        
        with gr.Accordion("🔗 Common Domain Examples", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            **Academic & Research:**
            - `arxiv.org`, `*.edu`, `scholar.google.com`, `researchgate.net`, `pubmed.ncbi.nlm.nih.gov`
            
            **Technology & Programming:**
            - `github.com`, `stackoverflow.com`, `docs.python.org`, `developer.mozilla.org`, `medium.com`
            
            **News & Media:**
            - `reuters.com`, `bbc.com`, `npr.org`, `apnews.com`, `cnn.com`, `nytimes.com`
            
            **Business & Finance:**
            - `bloomberg.com`, `wsj.com`, `nasdaq.com`, `sec.gov`, `investopedia.com`
            
            **Science & Medicine:**
            - `nature.com`, `science.org`, `pubmed.ncbi.nlm.nih.gov`, `who.int`, `cdc.gov`
            
            **Government & Official:**
            - `*.gov`, `*.org`, `un.org`, `worldbank.org`, `imf.org`
            """)
        
        with gr.Accordion("📖 How to Use This Enhanced App", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            ### 🚀 Getting Started
            1. **Enter your API Keys** - Groq from [console.groq.com](https://console.groq.com/), Chutes for openai/gpt-oss-20b
            2. **Select a model** - Choose based on your needs:
               - **Compound models** (Groq): For web search with domain filtering
               - **openai/gpt-oss-20b** (Chutes): For general conversational tasks
            3. **Configure search settings** - Choose search type and options (Groq models only)
            4. **Click Connect** - Validate your keys and connect to the AI
            5. **Start chatting!** - Type your message and get intelligent responses with citations
            
            ### 🎯 Key Features
            - **Dual APIs**: Web search (Groq) + Basic chat (Chutes)
            - **Smart Citations**: Automatic source linking and citation formatting (Groq models)
            - **Domain Filtering**: Control which websites the AI searches (Groq models)
            - **Model Flexibility**: Choose the right model and API for your task
            - **Enhanced Tool Visibility**: See search tools used (Groq models)
            
            ### 💡 Tips for Best Results
            
            **For Research Tasks:**
            - Use compound models with domain filtering
            - Include academic domains (*.edu, arxiv.org) for scholarly sources
            - Use "Force Search" for the most current information
            
            **For Creative Tasks:**
            - Use openai/gpt-oss-20b (Chutes) or any model
            - Set search type to "none" for purely creative responses
            - Use higher temperature (0.8-1.0) for more creativity
            """)
        
        with gr.Accordion("🎯 Sample Examples to Test Enhanced Search", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            <div class="example-box">
            <h4>🔬 Research & Analysis</h4>
            
            **Compound Model + Domain Filtering (Groq):**
            - Query: "What are the latest breakthroughs in quantum computing?"
            - Model: compound-beta
            - Include domains: "arxiv.org, *.edu, nature.com"
            - Search type: web_search
            
            <h4>💬 General Knowledge (Chutes):**
            - Query: "Tell me about quantum computing"
            - Model: openai/gpt-oss-20b
            - Search type: none
            
            <h4>💻 Programming & Tech</h4>
            
            **Technical Documentation (Groq):**
            - Query: "How to implement OAuth 2.0 in Python Flask?"
            - Model: compound-beta
            - Include domains: "github.com, docs.python.org, stackoverflow.com"
            - Search type: web_search
            
            **Code Help (Chutes):**
            - Same query with openai/gpt-oss-20b
            - Search type: none
            
            <h4>🎨 Creative Tasks</h4>
            - Query: "Write a short story about AI and humans working together"
            - Any model with search_type: "none"
            - Higher temperature (0.8-1.0)
            
            <h4>📊 Business Analysis</h4>
            
            **Business Analysis (Filtered, Groq):**
            - Query: "Cryptocurrency adoption in enterprise"
            - Model: compound-beta
            - Include domains: "bloomberg.com, wsj.com, harvard.edu"
            - Search type: web_search
            </div>
            """)
        
        send_btn.click(
            fn=chat_with_ai,
            inputs=[msg, include_domains, exclude_domains, system_prompt, temperature, max_tokens, search_type, force_search, chatbot],
            outputs=[chatbot, msg]
        )
        
        msg.submit(
            fn=chat_with_ai,
            inputs=[msg, include_domains, exclude_domains, system_prompt, temperature, max_tokens, search_type, force_search, chatbot],
            outputs=[chatbot, msg]
        )
        
        clear_btn.click(
            fn=clear_chat_history,
            outputs=[chatbot]
        )
        
        with gr.Accordion("🚀 About This Enhanced NeuroScope AI", open=False, elem_id="neuroscope-accordion"):
            gr.Markdown("""
            **Enhanced Creative Agentic AI Chat Tool** with dual API support:
            
            ### 🆕 **New in This Version:**
            - 💬 **Chutes API Integration**: For openai/gpt-oss-20b model
            - 🌐 **Dual API System**: Web search (Groq) + Basic chat (Chutes)
            - 🎯 **Model Flexibility**: Multiple models across two APIs
            - ⚡ **Force Search Option**: Make AI search for Groq models  
            - 🔧 **Enhanced Tool Visibility**: See search tools used (Groq models)
            - 📊 **Model Comparison Guide**: Choose the right model and API
            
            ### 🏆 **Core Features:**
            - 🔗 **Automatic Source Citations**: Clickable links to sources (Groq models)
            - 📚 **Sources Used Section**: Dedicated section for websites (Groq models)
            - 🌐 **Smart Domain Filtering**: Control search scope (Groq models)
            - 💬 **Conversational Memory**: Maintains context throughout the session
            - ⚙️ **Full Customization**: Adjust all parameters and prompts
            - 🎨 **Creative & Analytical**: Optimized for both creative and research tasks
            
            ### 🛠️ **Technical Details:**
            - **Compound Models (Groq)**: compound-beta, compound-beta-mini (web search + domain filtering)
            - **Chat Model (Chutes)**: openai/gpt-oss-20b (basic conversational capabilities)
            - **Automatic Search Type Detection**: AI chooses best search method (Groq models)
            - **Enhanced Error Handling**: Robust error management and user feedback
            - **Real-time Status Updates**: Live feedback on model capabilities and search settings
            """)
    
    return app

# Main execution
if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(
        share=True
    )