import os
import re
import tempfile
from typing import Dict, List, Tuple
import gradio as gr

# Optional imports with graceful fallback
try:
    import PyPDF2
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False

try:
    from pdf2image import convert_from_bytes
    PDF2IMAGE_AVAILABLE = True
except ImportError:
    PDF2IMAGE_AVAILABLE = False

try:
    import pytesseract
    from PIL import Image
    OCR_AVAILABLE = True
except ImportError:
    OCR_AVAILABLE = False

try:
    import docx
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

try:
    from transformers import pipeline
    import torch
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False

try:
    from groq import Groq
    GROQ_AVAILABLE = True
except ImportError:
    GROQ_AVAILABLE = False

# Configuration
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
ENABLE_GROQ = bool(GROQ_API_KEY) and GROQ_AVAILABLE

print(f"🚀 Medical Report Summarizer Initializing...")
print(f"📊 Configuration:")
print(f"  - GROQ Available: {GROQ_AVAILABLE}")
print(f"  - GROQ Enabled: {ENABLE_GROQ}")
print(f"  - Transformers: {TRANSFORMERS_AVAILABLE}")

# ==================== DOCUMENT PROCESSOR ====================

class DocumentProcessor:
    def __init__(self):
        self.reader = None
    
    def extract_text(self, file_path: str, file_type: str) -> Tuple[str, str]:
        """Extract text from various file formats"""
        if not os.path.exists(file_path):
            return "", "File not found"
        
        file_type = file_type.lower()
        
        try:
            # Text files
            if file_type == 'txt':
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        text = f.read()
                except:
                    with open(file_path, 'r', encoding='latin-1') as f:
                        text = f.read()
                return self._clean_text(text), ""
            
            # PDF files
            elif file_type == 'pdf':
                if not PDF_AVAILABLE:
                    return "", "PDF processing library not available"
                
                text = ""
                try:
                    with open(file_path, 'rb') as file:
                        pdf_reader = PyPDF2.PdfReader(file)
                        for page in pdf_reader.pages:
                            page_text = page.extract_text()
                            if page_text:
                                text += page_text + "\n"
                except Exception as e:
                    return "", f"PDF error: {str(e)}"
                
                if text.strip():
                    return self._clean_text(text), ""
                else:
                    return "", "No text extracted from PDF"
            
            # Image files
            elif file_type in ['jpg', 'jpeg', 'png', 'bmp']:
                if not OCR_AVAILABLE:
                    return "", "OCR library not available"
                
                try:
                    image = Image.open(file_path)
                    text = pytesseract.image_to_string(image)
                    if text.strip():
                        return self._clean_text(text), ""
                    else:
                        return "", "No text found in image"
                except Exception as e:
                    return "", f"Image processing error: {str(e)}"
            
            # Word documents
            elif file_type in ['docx']:
                if not DOCX_AVAILABLE:
                    return "", "Word document library not available"
                
                try:
                    doc = docx.Document(file_path)
                    text = ""
                    for paragraph in doc.paragraphs:
                        if paragraph.text.strip():
                            text += paragraph.text + "\n"
                    return self._clean_text(text), ""
                except Exception as e:
                    return "", f"Word document error: {str(e)}"
            
            else:
                return "", f"Unsupported file type: {file_type}"
                
        except Exception as e:
            return "", f"Processing error: {str(e)}"
    
    def _clean_text(self, text: str) -> str:
        """Clean and normalize text"""
        if not text:
            return ""
        # Remove excessive whitespace
        text = re.sub(r'\s+', ' ', text)
        # Normalize line breaks
        text = re.sub(r'\n+', '\n', text)
        return text.strip()

# ==================== SERIOUSNESS ANALYZER ====================

class SeriousnessAnalyzer:
    def __init__(self):
        self.critical_terms = {
            "high": [
                "cancer", "malignant", "metastasis", "tumor",
                "heart attack", "myocardial infarction", "stroke",
                "sepsis", "organ failure", "critical condition", "emergency",
                "life-threatening", "rupture", "internal bleeding"
            ],
            "medium": [
                "infection", "inflammation", "hypertension", "diabetes",
                "arthritis", "pneumonia", "bronchitis", "fracture",
                "ulcer", "kidney disease", "liver disease", "moderate",
                "chronic", "worsening"
            ],
            "low": [
                "mild", "slight", "minor", "stable", "improving",
                "benign", "routine", "checkup", "follow-up"
            ]
        }
    
    def analyze(self, text: str) -> Dict:
        """Analyze seriousness of medical findings"""
        if not text:
            return {"level": "Unknown", "score": 0, "recommendation": "No text to analyze"}
        
        text_lower = text.lower()
        severity_scores = {"high": 0, "medium": 0, "low": 0}
        
        for severity, terms in self.critical_terms.items():
            for term in terms:
                if term in text_lower:
                    severity_scores[severity] += text_lower.count(term)
        
        overall_score = (
            severity_scores["high"] * 3 +
            severity_scores["medium"] * 2 +
            severity_scores["low"]
        )
        
        if overall_score >= 5 or severity_scores["high"] >= 2:
            return {
                "level": "High",
                "score": overall_score,
                "recommendation": "🔴 URGENT: Consult healthcare provider immediately.",
                "details": f"Found {severity_scores['high']} high-risk terms"
            }
        elif overall_score >= 3:
            return {
                "level": "Medium",
                "score": overall_score,
                "recommendation": "🟡 MODERATE: Schedule follow-up with your doctor.",
                "details": f"Found {severity_scores['medium']} medium-risk terms"
            }
        else:
            return {
                "level": "Low",
                "score": overall_score,
                "recommendation": "🟢 ROUTINE: Discuss at your next appointment.",
                "details": f"Found {severity_scores['low']} routine terms"
            }

# ==================== CONSULTANT SEARCH ====================

class ConsultantSearch:
    def __init__(self):
        self.doctors = [
            {
                "id": 1,
                "name": "Dr. Ahmed Raza",
                "specialty": "Cardiology",
                "hospital": "Aga Khan University Hospital",
                "address": "Stadium Road, Karachi",
                "phone": "+92-21-111-111-111",
                "rating": 4.8,
                "experience": "15 years",
                "fee": "₨ 3,000",
                "online": True,
                "city": "Karachi"
            },
            {
                "id": 2,
                "name": "Dr. Saima Khan",
                "specialty": "Endocrinology",
                "hospital": "Shaukat Khanum Memorial Hospital",
                "address": "Johar Town, Lahore",
                "phone": "+92-42-111-111-111",
                "rating": 4.7,
                "experience": "12 years",
                "fee": "₨ 2,500",
                "online": True,
                "city": "Lahore"
            },
            {
                "id": 3,
                "name": "Dr. Usman Ali",
                "specialty": "General Physician",
                "hospital": "Telemedicine Pakistan",
                "address": "Online - Nationwide",
                "phone": "0300-123-4567",
                "rating": 4.4,
                "experience": "8 years",
                "fee": "₨ 1,500",
                "online": True,
                "city": "Online"
            },
            {
                "id": 4,
                "name": "Dr. Fatima Shah",
                "specialty": "Pediatrics",
                "hospital": "Children Hospital Lahore",
                "address": "Ferozepur Road, Lahore",
                "phone": "+92-42-111-111-112",
                "rating": 4.9,
                "experience": "10 years",
                "fee": "₨ 2,000",
                "online": True,
                "city": "Lahore"
            }
        ]
        
        self.condition_specialty_map = {
            "heart": "Cardiology",
            "diabetes": "Endocrinology",
            "blood pressure": "Cardiology",
            "cancer": "Oncology",
            "child": "Pediatrics",
            "fracture": "Orthopedics",
            "skin": "Dermatology",
            "pregnancy": "Gynecology"
        }
    
    def suggest_specialties(self, medical_terms: List[str]) -> List[str]:
        """Suggest specialties based on medical terms"""
        specialties = set()
        all_text = " ".join(medical_terms).lower()
        
        for condition, specialty in self.condition_specialty_map.items():
            if condition in all_text:
                specialties.add(specialty)
        
        if not specialties:
            specialties = {"General Physician"}
        
        return list(specialties)[:3]
    
    def search(self, specialty: str, city: str = "Any") -> List[Dict]:
        """Search for consultants"""
        results = []
        for doctor in self.doctors:
            # Check specialty match
            if specialty.lower() not in doctor["specialty"].lower():
                continue
            
            # Check city match
            if city.lower() != "any" and city.lower() not in doctor["city"].lower():
                continue
            
            results.append(doctor)
        
        # Sort by rating
        results.sort(key=lambda x: x["rating"], reverse=True)
        return results
    
    def format_doctor(self, doctor: Dict) -> str:
        """Format doctor info for display"""
        online_icon = "💻" if doctor["online"] else "🏥"
        return f"""
**{online_icon} {doctor['name']}** ⭐{doctor['rating']}
📍 **Specialty**: {doctor['specialty']}
🏥 **Hospital**: {doctor['hospital']}
🌆 **Location**: {doctor['city']}
💰 **Fee**: {doctor['fee']}
📞 **Phone**: `{doctor['phone']}`
👨‍⚕️ **Experience**: {doctor['experience']}
---
"""
    
    def format_results(self, doctors: List[Dict]) -> str:
        """Format search results"""
        if not doctors:
            return "❌ No doctors found matching your criteria. Try a different search."
        
        result = f"## 👨‍⚕️ Found {len(doctors)} Doctors\n\n"
        for i, doctor in enumerate(doctors, 1):
            result += f"**{i}.** {self.format_doctor(doctor)}\n"
        
        result += "\n💡 **Tip**: Call during business hours (9 AM - 5 PM) for appointments."
        return result

# ==================== AI ASSISTANT ====================

class AIAssistant:
    def __init__(self):
        self.ner_model = None
        self.groq_client = None
        self.seriousness_analyzer = SeriousnessAnalyzer()
        
        # Initialize NER model (lightweight)
        if TRANSFORMERS_AVAILABLE:
            try:
                print("🤖 Loading medical NER model...")
                self.ner_model = pipeline(
                    "ner",
                    model="samrawal/bert-base-uncased_clinical-ner",
                    aggregation_strategy="simple"
                )
                print("✅ NER model loaded successfully")
            except Exception as e:
                print(f"⚠️ Could not load NER model: {e}")
                self.ner_model = None
        
        # Initialize Groq client
        if ENABLE_GROQ:
            try:
                print("🤖 Initializing Groq client...")
                self.groq_client = Groq(api_key=GROQ_API_KEY)
                print("✅ Groq client initialized successfully")
            except Exception as e:
                print(f"⚠️ Could not initialize Groq: {e}")
                self.groq_client = None
        else:
            print("ℹ️ Groq not enabled (no API key or library)")
    
    def extract_medical_terms(self, text: str) -> Dict[str, List[str]]:
        """Extract medical terms from text"""
        if not text or not self.ner_model:
            return {"conditions": [], "medications": [], "symptoms": []}
        
        try:
            # Limit text length for performance
            processed_text = text[:2000]
            entities = self.ner_model(processed_text)
            
            terms = {
                "conditions": [],
                "medications": [],
                "symptoms": []
            }
            
            for entity in entities:
                if entity['score'] > 0.7:
                    category = entity['entity_group']
                    term = entity['word'].strip()
                    
                    if category == "DISEASE" and term not in terms["conditions"]:
                        terms["conditions"].append(term)
                    elif category == "MEDICATION" and term not in terms["medications"]:
                        terms["medications"].append(term)
                    elif category in ["SYMPTOM", "PROBLEM"] and term not in terms["symptoms"]:
                        terms["symptoms"].append(term)
            
            return terms
            
        except Exception as e:
            print(f"⚠️ Error extracting medical terms: {e}")
            return {"conditions": [], "medications": [], "symptoms": []}
    
    def generate_summary(self, text: str, medical_terms: Dict) -> str:
        """Generate patient-friendly summary"""
        # If Groq is available, use it
        if self.groq_client and ENABLE_GROQ:
            try:
                print("🤖 Generating AI summary with Groq...")
                
                # Prepare medical terms string
                terms_str = ""
                if medical_terms["conditions"]:
                    terms_str += f"Conditions: {', '.join(medical_terms['conditions'][:3])}\n"
                if medical_terms["medications"]:
                    terms_str += f"Medications: {', '.join(medical_terms['medications'][:3])}\n"
                
                # CORRECT: Use proper message format
                messages = [
                    {
                        "role": "system",
                        "content": "You are a helpful medical assistant that explains medical reports in simple, patient-friendly language. Be compassionate and clear."
                    },
                    {
                        "role": "user",
                        "content": f"""Please summarize this medical report in simple, patient-friendly language:

REPORT TEXT:
{text[:1500]}

MEDICAL TERMS IDENTIFIED:
{terms_str}

Please provide:
1. A simple overview of what the report is about
2. Key findings in everyday language
3. What the patient should do next
4. Any important warnings or next steps

Use bullet points and avoid medical jargon. Be empathetic and clear."""
                    }
                ]
                
                response = self.groq_client.chat.completions.create(
                    messages=messages,
                    model="llama-3.1-8b-instant",
                    temperature=0.3,
                    max_tokens=500
                )
                
                summary = response.choices[0].message.content
                print("✅ AI summary generated successfully")
                return summary
                
            except Exception as e:
                print(f"⚠️ AI summary failed, using fallback: {e}")
                return self._generate_fallback_summary(text, medical_terms)
        
        # Fallback summary
        return self._generate_fallback_summary(text, medical_terms)
    
    def _generate_fallback_summary(self, text: str, medical_terms: Dict) -> str:
        """Generate fallback summary without AI"""
        summary = ["## 🏥 Medical Report Summary", ""]
        
        # Add key findings
        if any(medical_terms.values()):
            summary.append("### 🔍 Key Findings:")
            
            if medical_terms["conditions"]:
                conditions = medical_terms["conditions"][:3]
                summary.append(f"- **Conditions identified**: {', '.join(conditions)}")
            
            if medical_terms["medications"]:
                medications = medical_terms["medications"][:3]
                summary.append(f"- **Medications mentioned**: {', '.join(medications)}")
            
            if medical_terms["symptoms"]:
                symptoms = medical_terms["symptoms"][:3]
                summary.append(f"- **Symptoms reported**: {', '.join(symptoms)}")
        else:
            summary.append("No specific medical terms were identified in the report.")
        
        summary.append("\n### 💡 What This Means:")
        summary.append("- This is a summary of your medical report findings")
        summary.append("- These results should be discussed with your healthcare provider")
        summary.append("- Your doctor can provide personalized interpretation")
        
        summary.append("\n### 📝 Next Steps:")
        summary.append("1. **Schedule** an appointment with your doctor")
        summary.append("2. **Bring** this report to your appointment")
        summary.append("3. **Ask questions** about anything you don't understand")
        summary.append("4. **Follow** your doctor's recommendations")
        
        summary.append("\n---")
        summary.append("**⚠️ Important**: This is an AI-generated summary for educational purposes only.")
        summary.append("Always consult qualified healthcare professionals for medical advice.")
        
        return "\n".join(summary)
    
    def chat_about_report(self, question: str, report_text: str, medical_terms: Dict) -> str:
        """Chat about the medical report - FIXED VERSION"""
        print(f"💬 Chat question: {question[:50]}...")
        
        # If no report processed yet
        if not report_text:
            return "Please upload and process a medical report first."
        
        # If Groq is available, use it
        if self.groq_client and ENABLE_GROQ:
            try:
                # Prepare medical terms string
                terms_str = ""
                if medical_terms["conditions"]:
                    terms_str += f"Conditions: {', '.join(medical_terms['conditions'][:3])}\n"
                if medical_terms["medications"]:
                    terms_str += f"Medications: {', '.join(medical_terms['medications'][:3])}\n"
                
                # CORRECT: Use proper message format with system and user roles
                messages = [
                    {
                        "role": "system",
                        "content": """You are a compassionate medical assistant that helps patients understand their medical reports. 
                        Always be clear, accurate, and supportive. 
                        Base your answers only on the information provided in the medical report.
                        If the information isn't in the report, politely say so and suggest asking their doctor."""
                    },
                    {
                        "role": "user",
                        "content": f"""I have a medical report and need your help understanding it.

MEDICAL REPORT CONTENT (first 1000 characters):
{report_text[:1000]}

MEDICAL TERMS IDENTIFIED:
{terms_str}

My question is: {question}

Please provide a helpful response that:
1. Directly answers my question based on the medical report
2. Uses simple, easy-to-understand language
3. Explains any medical terms mentioned
4. Is compassionate and supportive
5. Encourages me to discuss with my healthcare provider
6. Stays within the information available in the report

If the information isn't in the report, politely say so and suggest asking my doctor."""
                    }
                ]
                
                response = self.groq_client.chat.completions.create(
                    messages=messages,
                    model="llama-3.1-8b-instant",
                    temperature=0.4,
                    max_tokens=300
                )
                
                answer = response.choices[0].message.content
                print("✅ Chat response generated successfully")
                return answer
                
            except Exception as e:
                print(f"⚠️ Chat AI failed: {e}")
                return self._simple_chat_response(question, medical_terms)
        
        # Simple rule-based response
        return self._simple_chat_response(question, medical_terms)
    
    def _simple_chat_response(self, question: str, medical_terms: Dict) -> str:
        """Simple rule-based chat response"""
        question_lower = question.lower()
        
        # Check for specific types of questions
        if "what does" in question_lower or "mean" in question_lower:
            # Look for medical terms in the question
            all_terms = []
            for term_list in medical_terms.values():
                all_terms.extend(term_list)
            
            for term in all_terms:
                if term.lower() in question_lower:
                    return f"**{term}** is a medical term from your report. For its specific meaning in your context, please discuss it with your healthcare provider during your appointment."
        
        if "serious" in question_lower or "urgent" in question_lower or "emergency" in question_lower:
            return "The seriousness of your condition should be assessed by a healthcare professional. If you're experiencing severe symptoms, please seek immediate medical attention."
        
        if "doctor" in question_lower or "specialist" in question_lower or "appointment" in question_lower:
            return "Based on your report, I recommend discussing your findings with a healthcare provider. They can recommend the appropriate specialist if needed."
        
        # Default response
        return "Thank you for your question. I recommend discussing this with your healthcare provider who can give you personalized medical advice based on your complete health history and this report."

# ==================== MAIN APPLICATION ====================

class MedicalApp:
    def __init__(self):
        self.doc_processor = DocumentProcessor()
        self.ai_assistant = AIAssistant()
        self.consultant_search = ConsultantSearch()
        
        # State
        self.current_report_text = ""
        self.current_medical_terms = {"conditions": [], "medications": [], "symptoms": []}
        self.current_seriousness = None
    
    def process_uploaded_file(self, file):
        """Process uploaded file and return all outputs"""
        print(f"📄 Processing uploaded file...")
        
        if file is None:
            return self._get_placeholder_outputs()
        
        try:
            # Get file info
            file_path = file.name
            file_name = os.path.basename(file_path)
            file_ext = file_name.split('.')[-1].lower() if '.' in file_name else 'txt'
            
            print(f"📁 File: {file_name}, Type: {file_ext}")
            
            # Extract text
            extracted_text, error = self.doc_processor.extract_text(file_path, file_ext)
            
            if error:
                return [
                    f"❌ Error: {error}",
                    "Unable to assess without text.",
                    "No medical terms extracted.",
                    "",  # Clear chat input
                    []   # Clear chat history
                ]
            
            if len(extracted_text) < 20:
                return [
                    "❌ Not enough text extracted. The file may be empty or unreadable.",
                    "Unable to assess.",
                    "No medical terms found.",
                    "",
                    []
                ]
            
            print(f"✅ Extracted {len(extracted_text)} characters")
            
            # Store current report
            self.current_report_text = extracted_text
            
            # Extract medical terms
            self.current_medical_terms = self.ai_assistant.extract_medical_terms(extracted_text)
            print(f"🔍 Found medical terms: {sum(len(v) for v in self.current_medical_terms.values())}")
            
            # Generate summary
            summary = self.ai_assistant.generate_summary(extracted_text, self.current_medical_terms)
            
            # Analyze seriousness
            self.current_seriousness = self.ai_assistant.seriousness_analyzer.analyze(extracted_text)
            seriousness_text = self._format_seriousness(self.current_seriousness)
            
            # Format medical terms for display
            terms_text = self._format_medical_terms(self.current_medical_terms)
            
            print("✅ Processing complete!")
            
            return [
                summary,
                seriousness_text,
                terms_text,
                "",  # Clear chat input
                []   # Clear chat history
            ]
            
        except Exception as e:
            print(f"❌ Error in process_uploaded_file: {e}")
            import traceback
            traceback.print_exc()
            return [
                f"❌ Processing error: {str(e)[:100]}",
                "Assessment failed",
                "Term extraction failed",
                "",
                []
            ]
    
    def handle_chat_message(self, message: str, chat_history: List[Tuple[str, str]]):
        """Handle chat message from user"""
        print(f"💬 Received chat message: {message}")
        
        if not self.current_report_text:
            response = "Please upload a medical report first to ask questions about it."
        else:
            response = self.ai_assistant.chat_about_report(
                message,
                self.current_report_text,
                self.current_medical_terms
            )
        
        # Add to chat history
        chat_history.append((message, response))
        
        # Return empty string for input (clears it) and updated history
        return "", chat_history
    
    def search_consultants(self, city: str, specialty: str):
        """Search for consultants"""
        print(f"🔍 Searching consultants: {specialty} in {city}")
        
        if not specialty:
            return "Please select a specialty."
        
        doctors = self.consultant_search.search(specialty, city)
        return self.consultant_search.format_results(doctors)
    
    def get_specialty_suggestions(self):
        """Get specialty suggestions based on current report"""
        if not self.current_medical_terms:
            return []
        
        # Flatten all terms
        all_terms = []
        for term_list in self.current_medical_terms.values():
            all_terms.extend(term_list)
        
        return self.consultant_search.suggest_specialties(all_terms)
    
    def _format_seriousness(self, seriousness_data: Dict) -> str:
        """Format seriousness data for display"""
        if not seriousness_data:
            return "No seriousness assessment available."
        
        icon_map = {"High": "🔴", "Medium": "🟡", "Low": "🟢"}
        icon = icon_map.get(seriousness_data["level"], "⚪")
        
        return f"""
{icon} **Seriousness Level**: {seriousness_data['level']}

**Risk Score**: {seriousness_data['score']}/10

**Recommendation**:
{seriousness_data['recommendation']}

{seriousness_data.get('details', '')}
"""
    
    def _format_medical_terms(self, medical_terms: Dict) -> str:
        """Format medical terms for display"""
        if not any(medical_terms.values()):
            return "No specific medical terms were identified in the document."
        
        text = "### 🔍 Medical Terms Identified:\n\n"
        
        if medical_terms["conditions"]:
            text += "**Medical Conditions:**\n"
            for term in medical_terms["conditions"][:5]:
                text += f"- {term}\n"
            text += "\n"
        
        if medical_terms["medications"]:
            text += "**Medications:**\n"
            for term in medical_terms["medications"][:5]:
                text += f"- {term}\n"
            text += "\n"
        
        if medical_terms["symptoms"]:
            text += "**Symptoms:**\n"
            for term in medical_terms["symptoms"][:5]:
                text += f"- {term}\n"
        
        return text
    
    def _get_placeholder_outputs(self):
        return [
            "## 🏥 Medical Report Summarizer\n\nPlease upload a medical report to begin analysis.",
            "Upload a report to see seriousness assessment.",
            "Medical terms will be extracted here.",
            "",  # Empty chat input
            []   # Empty chat history
        ]

# ==================== GRADIO INTERFACE ====================

def create_gradio_interface():
    """Create the Gradio interface"""
    
    app = MedicalApp()
    
    # Use simple Blocks without css parameter for older Gradio versions
    with gr.Blocks() as demo:
        
        # Header with embedded CSS
        gr.Markdown("""
        <style>
        .gradio-container {
            max-width: 1200px !important;
            margin: auto !important;
        }
        .chatbot {
            min-height: 300px !important;
        }
        .summary-box {
            border: 1px solid #e0e0e0;
            border-radius: 10px;
            padding: 15px;
            background: #f9f9f9;
        }
        .seriousness-high {
            background-color: #ffebee;
            padding: 10px;
            border-radius: 5px;
            border-left: 4px solid #f44336;
        }
        .seriousness-medium {
            background-color: #fff3e0;
            padding: 10px;
            border-radius: 5px;
            border-left: 4px solid #ff9800;
        }
        .seriousness-low {
            background-color: #e8f5e9;
            padding: 10px;
            border-radius: 5px;
            border-left: 4px solid #4caf50;
        }
        </style>
        
        <h1 style="text-align: center; color: #2c3e50;">🏥 AI Medical Report Summarizer - Pakistan</h1>
        <p style="text-align: center; color: #7f8c8d;">Upload medical reports, get AI-powered summaries, and connect with healthcare providers</p>
        """)
        
        with gr.Tabs():
            # ===== TAB 1: REPORT ANALYSIS =====
            with gr.Tab("📄 Analyze Medical Report"):
                with gr.Row():
                    # Left column: File upload
                    with gr.Column(scale=1):
                        gr.Markdown("### 📤 Upload Medical Report")
                        file_input = gr.File(
                            label="Choose a file",
                            file_types=[
                                ".pdf", ".txt", ".docx", 
                                ".jpg", ".jpeg", ".png", ".bmp"
                            ],
                            type="filepath"
                        )
                        
                        gr.Markdown("""
                        ### 📋 Supported Formats:
                        - **PDF documents**
                        - **Text files** (.txt)
                        - **Word documents** (.docx)
                        - **Images** (.jpg, .png, .bmp)
                        
                        *Note: Processing may take a few moments*
                        """)
                    
                    # Right column: Results
                    with gr.Column(scale=2):
                        gr.Markdown("### 📝 AI-Powered Summary")
                        summary_output = gr.Markdown(
                            value="Upload a document to see the summary..."
                        )
                
                # Seriousness and Medical Terms in a row below
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 🔍 Seriousness Assessment")
                        seriousness_output = gr.Markdown(
                            value="Assessment will appear here..."
                        )
                    
                    with gr.Column():
                        gr.Markdown("### 💊 Medical Terms Found")
                        terms_output = gr.Markdown(
                            value="Medical terms will appear here..."
                        )
                
                # Divider
                gr.Markdown("---")
                
                # Chat interface
                gr.Markdown("### 💬 Ask Questions About Your Report")
                
                chatbot = gr.Chatbot(
                    label="Medical Assistant",
                    height=350,
                    elem_id="chatbot"
                )
                
                with gr.Row():
                    chat_input = gr.Textbox(
                        label="Type your question here...",
                        placeholder="Example: What does this finding mean? Is this serious?",
                        scale=4,
                        interactive=True
                    )
                    send_btn = gr.Button("Send", variant="primary", scale=1)
            
            # ===== TAB 2: FIND DOCTORS =====
            with gr.Tab("🩺 Find Pakistani Doctors"):
                with gr.Row():
                    # Left column: Search filters
                    with gr.Column(scale=1):
                        gr.Markdown("### 🔍 Search Filters")
                        
                        # Auto-suggest button
                        suggest_btn = gr.Button(
                            "💡 Get Suggestions from Report",
                            variant="secondary"
                        )
                        
                        # Search inputs
                        city_input = gr.Dropdown(
                            label="📍 City",
                            choices=["Any", "Karachi", "Lahore", "Islamabad", "Online"],
                            value="Any"
                        )
                        
                        specialty_input = gr.Dropdown(
                            label="🎯 Specialty",
                            choices=[
                                "General Physician", "Cardiology", "Endocrinology",
                                "Pediatrics", "Gynecology", "Dermatology"
                            ],
                            value="General Physician"
                        )
                        
                        search_btn = gr.Button(
                            "🔍 Search Doctors",
                            variant="primary"
                        )
                        
                        gr.Markdown("""
                        ### 💡 Tips:
                        - Search by city or select "Online" for telemedicine
                        - Use suggestions based on your medical report
                        - Call during business hours (9 AM - 5 PM)
                        """)
                    
                    # Right column: Results
                    with gr.Column(scale=2):
                        gr.Markdown("### 👨‍⚕️ Available Doctors")
                        doctor_results = gr.Markdown(
                            value="Enter search criteria to find doctors..."
                        )
            
            # ===== TAB 3: HEALTHCARE INFO =====
            with gr.Tab("ℹ️ Healthcare Information"):
                gr.Markdown("""
                <div style="padding: 20px;">
                <h1 style="color: #2c3e50;">🇵🇰 Pakistan Healthcare Resources</h1>
                
                <h2 style="color: #3498db;">🏥 Major Hospital Networks</h2>
                <ul>
                <li><strong>Aga Khan University Hospital</strong> (Karachi)</li>
                <li><strong>Shaukat Khanum Memorial Hospital</strong> (Lahore, Karachi)</li>
                <li><strong>Pakistan Institute of Medical Sciences (PIMS)</strong> (Islamabad)</li>
                <li><strong>Civil Hospital</strong> (Major cities)</li>
                <li><strong>Services Hospital</strong> (Lahore)</li>
                </ul>
                
                <h2 style="color: #3498db;">📞 Emergency Services</h2>
                <ul>
                <li><strong>Rescue 1122</strong>: Nationwide emergency service</li>
                <li><strong>Edhi Foundation</strong>: 115-123-321</li>
                <li><strong>Chhipa Ambulance</strong>: 1020 (in some areas)</li>
                </ul>
                
                <h2 style="color: #3498db;">💻 Telemedicine Services</h2>
                <ul>
                <li><strong>Sehat Kahani</strong>: Online doctor consultations</li>
                <li><strong>Marham.pk</strong>: Doctor appointments & reviews</li>
                <li><strong>Oladdoctor</strong>: Video consultations</li>
                <li><strong>DoctHERs</strong>: Female doctors for women</li>
                </ul>
                
                <h2 style="color: #3498db;">💰 Typical Consultation Fees</h2>
                <ul>
                <li>General Physician: ₨ 1,000 - ₨ 2,000</li>
                <li>Specialists: ₨ 2,000 - ₨ 4,000</li>
                <li>Senior Consultants: ₨ 3,000 - ₨ 5,000</li>
                <li>Online Consultation: ₨ 1,000 - ₨ 2,500</li>
                </ul>
                
                <h2 style="color: #3498db;">🆘 Important Notes</h2>
                <ol>
                <li>Always verify doctor credentials with PMDC</li>
                <li>Keep copies of all medical reports</li>
                <li>Ask about payment plans if needed</li>
                <li>For emergencies, go directly to the nearest hospital</li>
                <li>This tool is for educational purposes only</li>
                </ol>
                
                <div style="background-color: #fff3cd; padding: 15px; border-radius: 5px; border-left: 4px solid #ffc107;">
                <h3 style="color: #856404;">⚠️ Medical Disclaimer</h3>
                <p>This AI-powered tool provides educational summaries only. It does NOT provide medical advice, diagnosis, or treatment. Always consult qualified healthcare professionals for medical decisions.</p>
                </div>
                </div>
                """)
        
        # ===== EVENT HANDLERS =====
        
        # Process file upload
        file_input.change(
            fn=app.process_uploaded_file,
            inputs=[file_input],
            outputs=[summary_output, seriousness_output, terms_output, chat_input, chatbot]
        )
        
        # Chat functionality
        def process_chat(message, history):
            return app.handle_chat_message(message, history)
        
        # Submit on Enter key
        chat_input.submit(
            fn=process_chat,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )
        
        # Submit on button click
        send_btn.click(
            fn=process_chat,
            inputs=[chat_input, chatbot],
            outputs=[chat_input, chatbot]
        )
        
        # Consultant search
        search_btn.click(
            fn=app.search_consultants,
            inputs=[city_input, specialty_input],
            outputs=[doctor_results]
        )
        
        # Auto-suggest specialties
        def update_suggestions():
            suggestions = app.get_specialty_suggestions()
            if suggestions:
                return gr.update(choices=suggestions, value=suggestions[0])
            return gr.update(choices=["General Physician"], value="General Physician")
        
        suggest_btn.click(
            fn=update_suggestions,
            outputs=[specialty_input]
        )
    
    return demo

# ==================== MAIN ENTRY POINT ====================

print("🚀 Initializing Medical Report Summarizer Application...")

# Create the interface
demo = create_gradio_interface()

# For Hugging Face Spaces
if __name__ == "__main__":
    # Launch with appropriate settings
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_error=True
    )