import gradio as gr
import requests
import time
import re
import threading
import uvicorn
import logging
import os
import signal
import sys
from typing import Dict, List, Optional, Tuple
from collections import defaultdict

# Import your backend modules
from service_v2 import app as fastapi_app
from chapter_retrieval_system_v2 import MultiCollectionChapterRetrieval

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

class ICD10SearchInterface:
    def __init__(self, api_base_url: str = "http://127.0.0.1:8000"):
        """Initialize the interface with API base URL"""
        self.api_base_url = api_base_url.rstrip('/')
        self.server_ready = False
        self.max_retries = 30  # Increased for Spaces startup time
        
        # ICD-10 code to chapter mapping
        self.code_to_chapter = self._build_code_to_chapter_mapping()
        
    def _build_code_to_chapter_mapping(self) -> Dict[str, Dict[str, str]]:
        """Build mapping from ICD-10 code ranges to chapters"""
        return {
            # Chapter I: Certain infectious and parasitic diseases (A00-B99)
            "chapter_1_I": {
                "title": "Certain infectious and parasitic diseases",
                "code_ranges": ["A", "B"],
                "description": "Infectious diseases, parasitic diseases, and related conditions"
            },
            
            # Chapter II: Neoplasms (C00-D49)
            "chapter_2_II": {
                "title": "Neoplasms", 
                "code_ranges": ["C", "D"],
                "description": "Malignant neoplasms, benign neoplasms, and neoplasms of uncertain behavior"
            },
            
            # Chapter III: Diseases of blood and blood-forming organs (D50-D89)
            "chapter_3_III": {
                "title": "Diseases of the blood and blood-forming organs",
                "code_ranges": ["D5", "D6", "D7", "D8"],
                "description": "Anemias, coagulation defects, and other blood disorders"
            },
            
            # Chapter IV: Endocrine, nutritional and metabolic diseases (E00-E89)
            "chapter_4_IV": {
                "title": "Endocrine, nutritional and metabolic diseases",
                "code_ranges": ["E"],
                "description": "Diabetes, thyroid disorders, nutritional deficiencies, and metabolic disorders"
            },
            
            # Chapter V: Mental and behavioural disorders (F01-F99)
            "chapter_5_V": {
                "title": "Mental and behavioural disorders",
                "code_ranges": ["F"],
                "description": "Mental disorders, substance abuse, and behavioral conditions"
            },
            
            # Chapter VI: Diseases of the nervous system (G00-G99)
            "chapter_6_VI": {
                "title": "Diseases of the nervous system",
                "code_ranges": ["G"],
                "description": "Neurological disorders, epilepsy, migraines, and nervous system diseases"
            },
            
            # Chapter VII: Diseases of the eye and adnexa (H00-H59)
            "chapter_7_VII": {
                "title": "Diseases of the eye and adnexa",
                "code_ranges": ["H0", "H1", "H2", "H3", "H4", "H5"],
                "description": "Eye diseases, visual disorders, and related conditions"
            },
            
            # Chapter VIII: Diseases of the ear and mastoid process (H60-H95)
            "chapter_8_VIII": {
                "title": "Diseases of the ear and mastoid process",
                "code_ranges": ["H6", "H7", "H8", "H9"],
                "description": "Hearing disorders, ear infections, and mastoid conditions"
            },
            
            # Chapter IX: Diseases of the circulatory system (I00-I99)
            "chapter_9_IX": {
                "title": "Diseases of the circulatory system",
                "code_ranges": ["I"],
                "description": "Heart disease, hypertension, stroke, and vascular disorders"
            },
            
            # Chapter X: Diseases of the respiratory system (J00-J99)
            "chapter_10_X": {
                "title": "Diseases of the respiratory system",
                "code_ranges": ["J"],
                "description": "Pneumonia, asthma, COPD, and other respiratory conditions"
            },
            
            # Chapter XI: Diseases of the digestive system (K00-K95)
            "chapter_11_XI": {
                "title": "Diseases of the digestive system",
                "code_ranges": ["K"],
                "description": "Gastrointestinal disorders, liver disease, and digestive conditions"
            },
            
            # Chapter XII: Diseases of the skin and subcutaneous tissue (L00-L99)
            "chapter_12_XII": {
                "title": "Diseases of the skin and subcutaneous tissue",
                "code_ranges": ["L"],
                "description": "Skin infections, dermatitis, and subcutaneous tissue disorders"
            },
            
            # Chapter XIII: Diseases of the musculoskeletal system (M00-M99)
            "chapter_13_XIII": {
                "title": "Diseases of the musculoskeletal system and connective tissue",
                "code_ranges": ["M"],
                "description": "Arthritis, bone disorders, muscle diseases, and connective tissue conditions"
            },
            
            # Chapter XIV: Diseases of the genitourinary system (N00-N99)
            "chapter_14_XIV": {
                "title": "Diseases of the genitourinary system",
                "code_ranges": ["N"],
                "description": "Kidney disease, urinary disorders, and reproductive system conditions"
            },
            
            # Chapter XV: Pregnancy, childbirth and the puerperium (O00-O9A)
            "chapter_15_XV": {
                "title": "Pregnancy, childbirth and the puerperium",
                "code_ranges": ["O"],
                "description": "Pregnancy complications, delivery issues, and postpartum conditions"
            },
            
            # Chapter XVI: Certain conditions originating in the perinatal period (P00-P96)
            "chapter_16_XVI": {
                "title": "Certain conditions originating in the perinatal period",
                "code_ranges": ["P"],
                "description": "Newborn conditions and perinatal complications"
            },
            
            # Chapter XVII: Congenital malformations (Q00-Q99)
            "chapter_17_XVII": {
                "title": "Congenital malformations, deformations and chromosomal abnormalities",
                "code_ranges": ["Q"],
                "description": "Birth defects and chromosomal disorders"
            },
            
            # Chapter XVIII: Symptoms, signs and abnormal findings (R00-R99)
            "chapter_18_XVIII": {
                "title": "Symptoms, signs and abnormal clinical and laboratory findings",
                "code_ranges": ["R"],
                "description": "Symptoms and signs not elsewhere classified"
            },
            
            # Chapter XIX: Injury, poisoning and external causes (S00-T88)
            "chapter_19_XIX": {
                "title": "Injury, poisoning and certain other consequences of external causes",
                "code_ranges": ["S", "T"],
                "description": "Injuries, poisoning, and external cause consequences"
            },
            
            # Chapter XX: External causes of morbidity (V01-Y99)
            "chapter_20_XX": {
                "title": "External causes of morbidity",
                "code_ranges": ["V", "W", "X", "Y"],
                "description": "External causes of injury and poisoning"
            },
            
            # Chapter XXI: Factors influencing health status (Z00-Z99)
            "chapter_21_XXI": {
                "title": "Factors influencing health status and contact with health services",
                "code_ranges": ["Z"],
                "description": "Health maintenance, screening, and healthcare encounters"
            }
        }
    
    def wait_for_server(self, max_wait_time=60):
        """Wait for FastAPI server to be ready with enhanced logging"""
        logger.info(f"Waiting for FastAPI server at {self.api_base_url}")
        start_time = time.time()
        attempt = 0
        
        while time.time() - start_time < max_wait_time:
            attempt += 1
            try:
                response = requests.get(f"{self.api_base_url}/health", timeout=10)
                if response.status_code == 200:
                    self.server_ready = True
                    logger.info(f"FastAPI server ready after {attempt} attempts ({time.time() - start_time:.1f}s)")
                    return True
                else:
                    logger.warning(f"Server returned status {response.status_code}, attempt {attempt}")
            except requests.exceptions.RequestException as e:
                if attempt % 10 == 0:  # Log every 10 attempts
                    logger.info(f"Waiting for server... attempt {attempt} ({time.time() - start_time:.1f}s)")
                time.sleep(2)
                continue
        
        logger.error(f"FastAPI server failed to start within {max_wait_time} seconds")
        return False
    
    def get_server_status(self) -> Tuple[bool, str]:
        """Get current server status for UI display"""
        if not self.server_ready:
            return False, "Server starting up..."
        
        try:
            response = requests.get(f"{self.api_base_url}/health", timeout=5)
            if response.status_code == 200:
                return True, "Server Ready"
            else:
                return False, f"Server Error (Status: {response.status_code})"
        except requests.exceptions.RequestException as e:
            return False, f"Connection Error: {str(e)}"
    
    def test_connection(self) -> Tuple[bool, str]:
        """Test if the API is accessible"""
        return self.get_server_status()
    
    # Keep all your existing methods (copy from original code)
    def extract_category_code(self, icd_code: str) -> str:
        """Extract the main category code from ICD-10 code (e.g., I21.0 -> I21)"""
        if not icd_code:
            return ""
        
        code = icd_code.strip().upper()
        match = re.match(r'^([A-Z]\d{2,3})', code)
        if match:
            return match.group(1)
        
        return code
    
    def group_codes_by_category(self, results: List[Dict]) -> Dict[str, List[Dict]]:
        """Group ICD-10 codes by their main category"""
        categories = defaultdict(list)
        
        for result in results:
            code = result.get('code', '')
            category = self.extract_category_code(code)
            if category:
                categories[category].append(result)
        
        return dict(categories)
    
    def get_category_info(self, category_code: str, codes_in_category: List[Dict]) -> Dict:
        """Get information about a category from its codes"""
        category_result = None
        max_score = 0
        
        for code_info in codes_in_category:
            if code_info['code'] == category_code:
                category_result = code_info
                break
            if code_info['score'] > max_score:
                max_score = code_info['score']
                category_result = code_info
        
        return category_result or codes_in_category[0]
    
    def get_chapter_info_for_code(self, icd_code: str) -> Optional[Dict[str, str]]:
        """Get chapter information for a given ICD-10 code"""
        if not icd_code:
            return None
            
        code = icd_code.strip().upper()
        
        # Check each chapter's code ranges
        for chapter_id, chapter_data in self.code_to_chapter.items():
            for code_prefix in chapter_data["code_ranges"]:
                if code.startswith(code_prefix):
                    return {
                        "chapter_id": chapter_id,
                        "title": chapter_data["title"],
                        "description": chapter_data["description"]
                    }
        
        return None
    
    def search_icd10(
        self, 
        query: str,
        limit: int = 10,
        score_threshold: float = 0.3,
        search_mode: str = "smart",
        target_chapters: str = "",
        detailed_analysis: bool = False,
        chapters_per_sentence: int = 2
    ) -> str:
        """Search ICD-10 codes using the API with enhanced error handling for Spaces"""
        if not query or not query.strip():
            return "Please enter a diagnostic query."
        
        if not self.server_ready:
            return """
            <div style='text-align: center; padding: 20px; background: #ffeaa7; border-radius: 8px; margin: 20px 0;'>
                <h3>Server Starting Up</h3>
                <p>The FastAPI server is still initializing. Please wait a moment and try again.</p>
                <p><em>This usually takes 10-30 seconds on first load.</em></p>
            </div>
            """
        
        is_connected, connection_msg = self.test_connection()
        if not is_connected:
            return f"""
            <div style='text-align: center; padding: 20px; background: #fab1a0; border-radius: 8px; margin: 20px 0;'>
                <h3>Connection Error</h3>
                <p>{connection_msg}</p>
                <p><em>Please refresh the page and try again.</em></p>
            </div>
            """
        
        try:
            params = {
                "q": query.strip(),
                "limit": limit * 2,
                "score_threshold": score_threshold,
                "search_mode": search_mode or "smart",
                "detailed_analysis": detailed_analysis,
                "chapters_per_sentence": chapters_per_sentence
            }
            
            if target_chapters and target_chapters.strip():
                params["target_chapters"] = target_chapters.strip()
            
            start_time = time.time()
            response = requests.get(f"{self.api_base_url}/api/search", params=params, timeout=120)
            request_time = time.time() - start_time
            
            if response.status_code != 200:
                error_data = response.json() if response.headers.get('content-type', '').startswith('application/json') else {"detail": response.text}
                return f"""
                <div style='text-align: center; padding: 20px; background: #fab1a0; border-radius: 8px; margin: 20px 0;'>
                    <h3>API Error ({response.status_code})</h3>
                    <p>{error_data.get('detail', 'Unknown error')}</p>
                </div>
                """
            
            data = response.json()
            return self._format_sentence_results_with_enhanced_categories(data)
            
        except requests.exceptions.Timeout:
            return """
            <div style='text-align: center; padding: 20px; background: #fab1a0; border-radius: 8px; margin: 20px 0;'>
                <h3>Request Timeout</h3>
                <p>The search is taking too long. Try reducing the limit or increasing the score threshold.</p>
            </div>
            """
        except requests.exceptions.RequestException as e:
            logger.error(f"Request error: {e}")
            return f"""
            <div style='text-align: center; padding: 20px; background: #fab1a0; border-radius: 8px; margin: 20px 0;'>
                <h3>Request Error</h3>
                <p>{str(e)}</p>
            </div>
            """
        except Exception as e:
            logger.error(f"Unexpected error: {e}")
            return f"""
            <div style='text-align: center; padding: 20px; background: #fab1a0; border-radius: 8px; margin: 20px 0;'>
                <h3>Unexpected Error</h3>
                <p>{str(e)}</p>
            </div>
            """
    
    def _format_sentence_results_with_enhanced_categories(self, data: Dict) -> str:
        """Format sentence-based results with enhanced category and chapter information"""
        sentence_results = data.get('sentence_results', [])
        
        if not sentence_results:
            return "<div style='text-align: center; color: #666; padding: 20px;'>No sentence-based results available.</div>"

        html = """
        <div style='margin-bottom: 20px;'>
            <h3 style='color: #2c3e50; margin-bottom: 15px;'>Results by Sentence with Enhanced Category Information</h3>
            <p style='color: #666; margin-bottom: 20px;'>
                Results are organized by sentence and grouped by ICD-10 categories with chapter context. High-scoring codes are highlighted.
            </p>
        </div>
        """
        
        for i, sent_result in enumerate(sentence_results, 1):
            # Group results by category
            categories = self.group_codes_by_category(sent_result['results'])
            
            html += f"""
            <div style='margin-bottom: 30px; border: 2px solid #3498db; border-radius: 12px; overflow: hidden; box-shadow: 0 4px 6px rgba(0,0,0,0.1);'>
                <div style='background: linear-gradient(135deg, #3498db, #2980b9); color: white; padding: 15px;'>
                    <h4 style='margin: 0; font-size: 1.2em;'>
                        Sentence {i}: "{sent_result['sentence_text']}"
                    </h4>
                    <div style='margin-top: 8px; font-size: 0.9em; opacity: 0.9;'>
                        <span style='background-color: rgba(255,255,255,0.2); padding: 3px 8px; border-radius: 12px; margin-right: 10px;'>
                            {sent_result['total_results']} total results
                        </span>
                        <span style='background-color: rgba(255,255,255,0.2); padding: 3px 8px; border-radius: 12px;'>
                            Top 3 of {len(categories)} categories
                        </span>
                    </div>
                </div>
                <div style='padding: 20px;'>
            """
            
            # Sort categories by highest score and limit to top 3
            sorted_categories = sorted(
                categories.items(), 
                key=lambda x: max(code['score'] for code in x[1]), 
                reverse=True
            )[:3]
            
            for category_code, codes_in_category in sorted_categories:
                # Get category information
                category_info = self.get_category_info(category_code, codes_in_category)
                highest_score = max(code['score'] for code in codes_in_category)
                category_color = self._get_category_color(highest_score)
                
                # Get chapter information for this category
                sample_code = codes_in_category[0].get('code', category_code)
                chapter_info = self.get_chapter_info_for_code(sample_code)
                
                # Build enhanced category header
                category_title = category_info.get('title', 'Unknown Category')
                chapter_display = ""
                chapter_tooltip = ""
                
                if chapter_info:
                    chapter_display = f" • Chapter {chapter_info['chapter_id'].split('_')[1]} ({chapter_info['chapter_id'].split('_')[2]})"
                    chapter_tooltip = f"title='{chapter_info['description']}'"
                
                html += f"""
                <div style='margin-bottom: 20px; border: 1px solid {category_color}; border-radius: 8px; overflow: hidden;'>
                    <div style='background-color: {category_color}; color: white; padding: 12px 15px;'>
                        <div style='display: flex; justify-content: space-between; align-items: flex-start;'>
                            <div style='flex-grow: 1;'>
                                <h5 style='margin: 0; font-size: 1em; line-height: 1.3;'>
                                    <span style='display: block;'>
                                        Category {category_code}: {category_title}
                                    </span>
                                    {f'<span style="font-size: 0.85em; opacity: 0.9; display: block; margin-top: 4px;" {chapter_tooltip}>{chapter_display}</span>' if chapter_info else ''}
                                </h5>
                                {f'<div style="font-size: 0.8em; opacity: 0.8; margin-top: 6px; line-height: 1.2;">{chapter_info["description"]}</div>' if chapter_info else ''}
                            </div>
                            <div style='text-align: right; margin-left: 15px;'>
                                <span style='font-size: 0.8em; background-color: rgba(255,255,255,0.2); padding: 2px 6px; border-radius: 10px; display: block;'>
                                    Max: {highest_score:.3f}
                                </span>
                                <span style='font-size: 0.75em; opacity: 0.8; margin-top: 2px; display: block;'>
                                    {len(codes_in_category)} codes
                                </span>
                            </div>
                        </div>
                    </div>
                    <div style='padding: 12px;'>
                """
                
                # Sort codes within category by score
                sorted_codes = sorted(codes_in_category, key=lambda x: x['score'], reverse=True)
                
                # Filter out codes that are the same as the category code
                filtered_codes = [code for code in sorted_codes if code.get('code', '') != category_code]
                
                # If we filtered out all codes or have no codes, show a message
                if not filtered_codes:
                    html += f"""
                    <div style='margin-bottom: 8px; padding: 12px; background-color: #f8f9fa; border-radius: 6px; border-left: 4px solid #95a5a6;'>
                        <div style='color: #666; text-align: center; font-style: italic;'>
                            Category {category_code} represents the main code group. Specific subcodes available in detailed search.
                        </div>
                    </div>
                    """
                else:
                    for j, result in enumerate(filtered_codes, 1):
                        score_color = self._get_score_color(result['score'])
                        is_high_score = result['score'] >= 0.6
                        
                        # Add highlighting for high-scoring codes
                        highlight_style = ""
                        if is_high_score:
                            highlight_style = "box-shadow: 0 0 0 2px #f39c12; background: linear-gradient(135deg, #fff9e6, #ffffff);"
                        
                        html += f"""
                        <div style='margin-bottom: 8px; padding: 12px; background-color: #f8f9fa; border-radius: 6px; border-left: 4px solid {score_color}; {highlight_style}'>
                            <div style='display: flex; justify-content: space-between; align-items: center;'>
                                <div style='flex-grow: 1;'>
                                    <strong style='color: #2c3e50; font-size: 1em;'>
                                        {result['code']} - {result['title']}
                                        {' ⭐' if is_high_score else ''}
                                    </strong>
                                </div>
                                <span style='background-color: {score_color}; color: white; padding: 3px 8px; border-radius: 4px; font-size: 0.85em; font-weight: bold;'>
                                    {result['score']:.3f}
                                </span>
                            </div>
                            {f"<div style='font-size: 0.9em; color: #666; margin-top: 8px; line-height: 1.4;'>{result['description'][:250]}{'...' if len(result.get('description', '')) > 250 else ''}</div>" if result.get('description') else ""}
                        </div>
                        """
                
                html += "</div></div>"
            
            html += "</div></div>"
        
        # Enhanced legend with chapter info
        html += """
        <div style='background-color: #f8f9fa; border-radius: 8px; padding: 15px; margin-top: 20px;'>
            <h4 style='color: #2c3e50; margin-bottom: 15px;'>Enhanced Legend</h4>
            
            <div style='margin-bottom: 15px;'>
                <h5 style='color: #2c3e50; margin-bottom: 8px;'>Score Quality:</h5>
                <div style='display: flex; flex-wrap: wrap; gap: 15px; align-items: center;'>
                    <div style='display: flex; align-items: center;'>
                        <div style='width: 20px; height: 20px; background-color: #27ae60; border-radius: 3px; margin-right: 8px;'></div>
                        <span style='font-size: 0.9em;'>Excellent Match (≥0.8)</span>
                    </div>
                    <div style='display: flex; align-items: center;'>
                        <div style='width: 20px; height: 20px; background-color: #f39c12; border-radius: 3px; margin-right: 8px;'></div>
                        <span style='font-size: 0.9em;'>Good Match (≥0.6)</span>
                    </div>
                    <div style='display: flex; align-items: center;'>
                        <div style='width: 20px; height: 20px; background-color: #e67e22; border-radius: 3px; margin-right: 8px;'></div>
                        <span style='font-size: 0.9em;'>Fair Match (≥0.4)</span>
                    </div>
                    <div style='display: flex; align-items: center;'>
                        <div style='width: 20px; height: 20px; background-color: #e74c3c; border-radius: 3px; margin-right: 8px;'></div>
                        <span style='font-size: 0.9em;'>Low Match (<0.4)</span>
                    </div>
                </div>
            </div>
            
            <div>
                <h5 style='color: #2c3e50; margin-bottom: 8px;'>Features:</h5>
                <div style='display: flex; flex-wrap: wrap; gap: 20px; align-items: center; font-size: 0.9em;'>
                    <span>High-scoring codes (≥0.6)</span>
                    <span>Category grouping by ICD-10 structure</span>
                    <span>Chapter context and descriptions</span>
                    <span>Score-based category prioritization</span>
                    <span>Duplicate category codes filtered</span>
                </div>
            </div>
        </div>
        """
        
        return html
    
    def _get_score_color(self, score: float) -> str:
        """Get color based on similarity score"""
        if score >= 0.8:
            return "#27ae60"  # Green
        elif score >= 0.6:
            return "#f39c12"  # Orange
        elif score >= 0.4:
            return "#e67e22"  # Dark orange
        else:
            return "#e74c3c"  # Red
    
    def _get_category_color(self, max_score: float) -> str:
        """Get category header color based on highest score in category"""
        if max_score >= 0.8:
            return "#2ecc71"  # Bright green
        elif max_score >= 0.6:
            return "#3498db"  # Blue
        elif max_score >= 0.4:
            return "#9b59b6"  # Purple
        else:
            return "#95a5a6"  # Gray

def start_fastapi_server():
    """Start FastAPI server with enhanced error handling for Spaces"""
    try:
        logger.info("Starting FastAPI server...")
        # Use environment variable for port if available
        port = int(os.environ.get("FASTAPI_PORT", "8000"))
        
        # Enhanced server configuration for Spaces
        uvicorn.run(
            fastapi_app, 
            host="127.0.0.1", 
            port=port, 
            log_level="info",
            access_log=False,  # Reduce log noise
            workers=1,  # Single worker for Spaces
            timeout_keep_alive=30
        )
    except Exception as e:
        logger.error(f"FastAPI server failed to start: {e}")
        # Don't raise - let Gradio continue with error messages

def create_gradio_interface():
    """Create the Gradio interface with server status monitoring"""
    search_interface = ICD10SearchInterface()
    
    css = """
    .gradio-container {
        max-width: 1400px !important;
        margin: auto !important;
    }
    
    .server-status {
        transition: all 0.3s ease;
    }
    """
    
    with gr.Blocks(css=css, title="ICD-10 Smart Search", theme=gr.themes.Soft()) as demo:
        gr.HTML("""
        <div style='text-align: center; margin-bottom: 30px; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);'>
            <h1 style='color: white; margin: 0; font-size: 2.5em;'>ICD-10 Smart Search</h1>
            <p style='color: #f1f2f6; margin: 15px 0 0 0; font-size: 1.2em;'>Advanced diagnostic code search with AI-powered sentence analysis</p>
        </div>
        """)
        
        # Server status indicator
        def get_server_status():
            is_ready, msg = search_interface.get_server_status()
            if is_ready:
                return "<div class='server-status' style='text-align: center; padding: 10px; background: #00b894; color: white; border-radius: 5px; margin-bottom: 20px;'>🟢 Server Ready</div>"
            else:
                return f"<div class='server-status' style='text-align: center; padding: 10px; background: #e17055; color: white; border-radius: 5px; margin-bottom: 20px;'>🔴 {msg}</div>"
        
        server_status = gr.HTML(value=get_server_status())
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML("<h3>Search Parameters</h3>")
                
                query_input = gr.Textbox(
                    label="Diagnostic Query",
                    placeholder="Enter diagnostic description (e.g., 'chest pain with shortness of breath')",
                    lines=3,
                    value=""
                )
                
                with gr.Accordion("Advanced Options", open=False):
                    with gr.Row():
                        limit_input = gr.Slider(
                            label="Maximum Results per Sentence",
                            minimum=5,
                            maximum=50,
                            value=15,
                            step=5,
                            info="Higher values show more codes per category"
                        )
                        
                        score_threshold_input = gr.Slider(
                            label="Score Threshold",
                            minimum=0.1,
                            maximum=0.9,
                            value=0.2,
                            step=0.05,
                            info="Lower values include more potential matches"
                        )
                    
                    search_mode_input = gr.Dropdown(
                        label="Search Mode",
                        choices=["smart", "all_chapters", "specific_chapters"],
                        value="smart"
                    )
                    
                    target_chapters_input = gr.Textbox(
                        label="Target Chapters (comma-separated)",
                        placeholder="e.g., chapter_9_IX, chapter_10_X",
                        visible=False
                    )
                    
                    with gr.Row():
                        detailed_analysis_input = gr.Checkbox(
                            label="Include Detailed Analysis",
                            value=True
                        )
                        
                        chapters_per_sentence_input = gr.Slider(
                            label="Chapters per Sentence",
                            minimum=1,
                            maximum=5,
                            value=3,
                            step=1
                        )
                
                search_button = gr.Button("Search ICD-10 Codes", variant="primary", size="lg")
                
                def update_target_chapters_visibility(search_mode):
                    return gr.update(visible=(search_mode == "specific_chapters"))
                
                search_mode_input.change(
                    update_target_chapters_visibility,
                    inputs=search_mode_input,
                    outputs=target_chapters_input
                )
            
            with gr.Column(scale=2):
                gr.HTML("<h3>Enhanced Category-Grouped Results</h3>")
                sentence_results_output = gr.HTML(
                    value="<div style='text-align: center; color: #666; padding: 40px;'>Enter a diagnostic query and click search to see categorized results with chapter context.</div>"
                )
        
        # Example queries
        gr.HTML("<h3>Example Queries</h3>")
        
        example_queries = [
            "acute myocardial infarction with chest pain",
            "type 2 diabetes with diabetic nephropathy",
            "major depressive disorder with anxiety", 
            "fracture of distal radius from fall",
            "acute appendicitis with peritonitis",
            "gestational diabetes in pregnancy",
            "chronic kidney disease stage 3",
            "essential hypertension with heart disease"
        ]
        
        with gr.Row():
            for i in range(0, len(example_queries), 2):
                with gr.Column():
                    for j in range(2):
                        if i + j < len(example_queries):
                            example_btn = gr.Button(
                                example_queries[i + j],
                                variant="secondary",
                                size="sm"
                            )
                            example_btn.click(
                                lambda x=example_queries[i + j]: x,
                                outputs=query_input
                            )
        
        # Search functionality
        search_button.click(
            fn=search_interface.search_icd10,
            inputs=[
                query_input,
                limit_input,
                score_threshold_input,
                search_mode_input,
                target_chapters_input,
                detailed_analysis_input,
                chapters_per_sentence_input
            ],
            outputs=sentence_results_output
        )
        
        # Enhanced footer
        gr.HTML("""
        <div style='text-align: center; margin-top: 30px; padding: 20px; background-color: #f8f9fa; border-radius: 12px; border: 1px solid #e9ecef;'>
            <p style='margin: 0; color: #666; line-height: 1.6;'>
                Powered by advanced semantic search and AI-driven sentence analysis<br>
                <strong>Features:</strong> Chapter context • Category descriptions • Score-based prioritization<br>
                <strong>Note:</strong> This tool is for research purposes only and should not replace professional medical diagnosis
            </p>
        </div>
        """)
        
        # Auto-refresh server status every 10 seconds
        demo.load(get_server_status, outputs=server_status, every=10)
    
    return demo

# Global variable to track server thread
server_thread = None

def graceful_shutdown():
    """Handle graceful shutdown"""
    logger.info("Shutting down application...")
    # Add any cleanup code here if needed

# Signal handlers for graceful shutdown
signal.signal(signal.SIGTERM, lambda signum, frame: graceful_shutdown())
signal.signal(signal.SIGINT, lambda signum, frame: graceful_shutdown())

# Main application entry point for Hugging Face Spaces
if __name__ == "__main__":
    logger.info("Starting ICD-10 Search Application for Hugging Face Spaces...")
    
    try:
        # Start FastAPI server in background thread
        logger.info("Initializing FastAPI server thread...")
        server_thread = threading.Thread(target=start_fastapi_server, daemon=True)
        server_thread.start()
        logger.info("FastAPI server thread started")
        
        # Give server time to start (increased for Spaces)
        logger.info("Waiting for FastAPI server initialization...")
        time.sleep(8)  # Increased wait time for Spaces
        
        # Create and launch Gradio interface
        logger.info("Creating Gradio interface...")
        demo = create_gradio_interface()
        
        # Launch for Spaces environment
        logger.info("Launching Gradio interface for Hugging Face Spaces...")
        demo.launch(
            share=False,  # Don't create public link
            show_error=True,  # Show errors for debugging
            show_tips=False,  # Don't show Gradio tips
            quiet=False,  # Show startup info
            server_name="0.0.0.0",  # Listen on all interfaces for Spaces
            server_port=7860,  # Default Gradio port for Spaces
            prevent_thread_lock=False,
            root_path=os.environ.get("GRADIO_ROOT_PATH", "")  # Support for Spaces routing
        )
        
    except Exception as e:
        logger.error(f"Application failed to start: {e}")
        sys.exit(1)