import gradio as gr
import base64
import os
from pathlib import Path
import tempfile
import time
from groq import Groq
import requests  # We need to keep requests for Kimi K2

# --- Configuration ---
GROQ_CLIENT = Groq(api_key=os.environ.get('groq'))
HYPERBOLIC_KEY = os.environ.get('h_api_key')  # For Kimi K2

class MediClearBackend:
    """Handles the API logic to keep the UI code clean."""
    
    @staticmethod
    def encode_image(image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')

    @staticmethod
    def analyze_medical_image(image_path):
        """Step 1: Extract technical info from a single image"""
        try:
            base64_image = MediClearBackend.encode_image(image_path)
            
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "You are a highly experienced medical doctor. Extract every piece of information from this medical image. Provide detailed observations about any findings, measurements, abnormalities, or notable features."},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                    ]
                }
            ]

            completion = GROQ_CLIENT.chat.completions.create(
                model="meta-llama/llama-4-scout-17b-16e-instruct",
                messages=messages,
                temperature=0.1,
                max_completion_tokens=4096,
                top_p=0.9,
                stream=True,
                stop=None
            )

            technical_data = ""
            for chunk in completion:
                if chunk.choices[0].delta.content:
                    technical_data += chunk.choices[0].delta.content

            return technical_data
        except Exception as e:
            raise Exception(f"Image Analysis Error: {str(e)}")

    @staticmethod
    def extract_text_from_multiple_images(image_files):
        """Extract and combine text from multiple medical images"""
        combined_text = "COMBINED MEDICAL REPORTS ANALYSIS:\n\n"
        
        for index, image_file in enumerate(image_files, 1):
            try:
                filename = Path(image_file.name).name
                combined_text += f"=== REPORT {index}: {filename} ===\n"
                
                # Extract text from this image
                image_text = MediClearBackend.analyze_medical_image(image_file.name)
                combined_text += f"{image_text}\n\n"
                combined_text += "=" * 50 + "\n\n"
                
            except Exception as e:
                combined_text += f"❌ Failed to process {filename}: {str(e)}\n\n"
                combined_text += "=" * 50 + "\n\n"
        
        return combined_text

    @staticmethod
    def summarize_combined_reports(combined_text):
        """Step 2: Convert combined technical info to patient-friendly summary"""
        try:
            if not HYPERBOLIC_KEY:
                raise Exception("Hyperbolic API key not configured")
                
            url = "https://api.hyperbolic.xyz/v1/chat/completions"
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {HYPERBOLIC_KEY}"
            }
            
            system_prompt = """You are a senior medical doctor with 50 years of experience. 
            Your role is to analyze multiple medical reports and provide a comprehensive, unified analysis to patients in simple, reassuring, everyday language.
            Never mention the doctors name whatsoever and never ask any counter questions.
            
            You will receive multiple medical reports combined together. Analyze them as a complete patient case and provide:
            
            Structure:
            1. Warm greeting acknowledging this is a comprehensive analysis of multiple reports.
            2. Overall unified assessment of the patient's condition across all reports.
            3. Key findings from all reports, highlighting patterns, consistencies, or important observations.
            4. Comprehensive recommendations based on the complete picture.
            5. Reassuring closing.
            
            Speak directly to the patient and provide a holistic view of their medical situation."""

            data = {
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": combined_text}
                ],
                "model": "moonshotai/Kimi-K2-Instruct",  # This is the Kimi K2 model
                "max_tokens": 4096,
                "temperature": 0.1,
                "top_p": 0.9
            }

            response = requests.post(url, headers=headers, json=data)
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except Exception as e:
            raise Exception(f"Summary Error (Kimi K2): {str(e)}")

def process_single_image(image_file):
    """Process a single medical image and return the analysis"""
    backend = MediClearBackend()
    
    if image_file is None:
        return "❌ Please upload a medical image first."
    
    try:
        # Step 1: Technical analysis with Groq (Llama 4 Scout)
        technical_data = backend.analyze_medical_image(image_file.name)
        
        # Step 2: Patient-friendly summary with Kimi K2
        final_report = backend.summarize_combined_reports(technical_data)
        
        # Format the final report as Markdown
        filename = Path(image_file.name).name
        formatted_report = f"""
# 🩺 MEDICAL ANALYSIS REPORT: {filename}
---
## 🔬 Technical Analysis
*Medical image processing complete*
## 📝 Patient Report Summary
{final_report}
---
> ✅ **Analysis Complete** - This tool provides AI-powered insights and is not a substitute for professional medical diagnosis.
"""
        return formatted_report
        
    except Exception as e:
        return f"## ❌ Error\n**Error processing image:** {str(e)}"

def process_multiple_images(image_files):
    """Process multiple medical images and return a SINGLE combined analysis"""
    backend = MediClearBackend()
    
    if not image_files:
        return "❌ Please upload at least one medical image."
    
    try:
        # Step 1: Extract and combine text from all images
        combined_text = backend.extract_text_from_multiple_images(image_files)
        
        # Step 2: Generate a single comprehensive summary from combined text
        comprehensive_summary = backend.summarize_combined_reports(combined_text)
        
        # Format as a single comprehensive report
        full_report = f"""# 🩺 COMPREHENSIVE MEDICAL ANALYSIS

## 📊 Analysis of {len(image_files)} Medical Reports

**Combined insights from all uploaded medical images:**

---

## 📝 Comprehensive Patient Summary

{comprehensive_summary}

---

## 📋 Reports Analyzed:
"""
        # List all analyzed files
        for index, image_file in enumerate(image_files, 1):
            filename = Path(image_file.name).name
            full_report += f"- **Report {index}:** {filename}\n"
        
        full_report += """
---
> 🏁 **Comprehensive Analysis Complete** - This unified analysis provides insights across all your medical reports. Remember: This tool provides AI-powered insights and is not a substitute for professional medical diagnosis.
"""
        return full_report
        
    except Exception as e:
        return f"## ❌ Processing Error\n**Failed to process multiple images:** {str(e)}"

def create_gradio_interface():
    """Create the Gradio interface"""
    
    # Custom CSS for better styling
    custom_css = """
    .gradio-container {
        background: linear-gradient(135deg, #1A1A1A 0%, #2D2D2D 100%);
        font-family: 'Segoe UI', system-ui, sans-serif;
    }
    .dark {
        background: transparent !important;
    }
    .panel {
        background: #252525 !important;
        border-radius: 12px !important;
        border: 1px solid #404040 !important;
        padding: 20px !important;
    }
    .header {
        background: linear-gradient(135deg, #1A1A1A 0%, #252525 100%) !important;
        border-radius: 12px !important;
        padding: 20px !important;
        margin-bottom: 20px !important;
        border: 1px solid #404040 !important;
    }
    .warning {
        background: #FF980020 !important;
        border: 1px solid #FF9800 !important;
        border-radius: 8px !important;
        padding: 15px !important;
        margin: 10px 0 !important;
    }
    .success {
        background: #4CAF5020 !important;
        border: 1px solid #4CAF50 !important;
        border-radius: 8px !important;
        padding: 15px !important;
        margin: 10x 0 !important;
    }
    .processing {
        background: #FFD60A20 !important;
        border: 1px solid #FFD60A !important;
        border-radius: 8px !important;
        padding: 15px !important;
        margin: 10px 0 !important;
    }
    .markdown-output {
        background: #1a1a1a !important;
        border-radius: 8px !important;
        padding: 20px !important;
    }
    .markdown-output h1, .markdown-output h2, .markdown-output h3 {
        color: #FFD60A !important;
        margin-top: 1em !important;
        margin-bottom: 0.5em !important;
    }
    .markdown-output p {
        color: #E8E8E8 !important;
        line-height: 1.6 !important;
    }
    .markdown-output ul, .markdown-output ol {
        color: #E8E8E8 !important;
        margin-left: 20px !important;
    }
    .markdown-output blockquote {
        border-left: 4px solid #FFD60A !important;
        margin: 1em 0 !important;
        padding-left: 1em !important;
        color: #B0B0B0 !important;
        font-style: italic !important;
    }
    .markdown-output code {
        background: #404040 !important;
        padding: 2px 6px !important;
        border-radius: 4px !important;
        color: #E8E8E8 !important;
    }
    .markdown-output pre {
        background: #404040 !important;
        padding: 15px !important;
        border-radius: 8px !important;
        overflow-x: auto !important;
    }
    .markdown-output hr {
        border: none !important;
        border-top: 2px solid #404040 !important;
        margin: 2em 0 !important;
    }
    """
    
    with gr.Blocks(
        title="SehatScan - Koshur AI", 
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="yellow",
            neutral_hue="slate"
        )
    ) as demo:
        
        # Header Section
        with gr.Column(elem_classes="header"):
            gr.Markdown(
                """
                # 🩺 SehatScan - Koshur AI
                ### Professional Medical Image Analysis
                
                """
            )
        
        # Main Content
        with gr.Row():
            # Left Panel - Image Input
            with gr.Column(scale=1, min_width=400):
                with gr.Group(elem_classes="panel"):
                    gr.Markdown("### 📁 Upload Medical Images")
                    
                    # Single image upload
                    single_image = gr.File(
                        label="Single Image Analysis",
                        file_types=[".jpg", ".jpeg", ".png", ".bmp"],
                        file_count="single",
                        type="filepath"
                    )
                    
                    # Multiple image upload
                    multiple_images = gr.File(
                        label="Batch Image Analysis (Multiple) - Combined Report",
                        file_types=[".jpg", ".jpeg", ".png", ".bmp"],
                        file_count="multiple",
                        type="filepath"
                    )
                    
                    # Analysis buttons
                    with gr.Row():
                        analyze_single_btn = gr.Button(
                            "🚀 Analyze Single Image",
                            variant="primary",
                            size="lg"
                        )
                        analyze_batch_btn = gr.Button(
                            "📊 Analyze Multiple Images", 
                            variant="secondary",
                            size="lg"
                        )
                    
                    # Clear button
                    clear_btn = gr.Button("🗑️ Clear All", variant="stop")
                
                # Processing Info
                with gr.Group(elem_classes="processing"):
                    gr.Markdown(
                        """
                        ### 🔄 Processing Information
                        **Single Image:** Individual analysis
                        **Multiple Images:** Combined comprehensive analysis from all reports
                        """
                    )
                
                # Disclaimer
                with gr.Group(elem_classes="warning"):
                    gr.Markdown(
                        """
                        ### ⚠️ Medical Disclaimer
                        This tool provides AI-powered insights and is not a substitute for professional medical diagnosis, 
                        treatment, or medical advice. Always consult with qualified healthcare providers for medical decisions.
                        """
                    )
            
            # Right Panel - Output
            with gr.Column(scale=2, min_width=600):
                with gr.Group(elem_classes="panel"):
                    gr.Markdown("### 📋 Doctor's Analysis Report")
                    
                    # Changed from Textbox to Markdown for pure markdown rendering
                    output_markdown = gr.Markdown(
                        label="Analysis Results",
                        value="*Upload a medical image and click analyze to see the results here...*",
                        elem_classes="markdown-output"
                    )
                    
                    # Output actions
                    with gr.Row():
                        copy_btn = gr.Button("📋 Copy to Clipboard", size="sm")
                        export_btn = gr.Button("💾 Export Report", size="sm")
        
        # Footer
        with gr.Column(elem_classes="panel"):
            with gr.Row():
                gr.Markdown(
                    """
                    **❤️ Built by Koshur AI • An initiative for Kashmir**  
                    [Visit our LinkedIn](https://www.linkedin.com/in/koshurai/)
                    """
                )
        
        # Event handlers - updated to use Markdown output
        analyze_single_btn.click(
            fn=process_single_image,
            inputs=[single_image],
            outputs=[output_markdown],
            api_name="analyze_single"
        )
        
        analyze_batch_btn.click(
            fn=process_multiple_images,
            inputs=[multiple_images],
            outputs=[output_markdown],
            api_name="analyze_batch"
        )
        
        # Clear functionality
        def clear_all():
            return None, None, "*Upload a medical image and click analyze to see the results here...*"
        
        clear_btn.click(
            fn=clear_all,
            outputs=[single_image, multiple_images, output_markdown]
        )
        
        # Copy functionality - for Markdown we'll use a different approach
        copy_btn.click(
            fn=lambda: gr.Info("Use the copy button in the Markdown output or select and copy the text directly"),
            outputs=None
        )
        
        # Export functionality
        def export_report():
            return gr.Info("To export, please select and copy the Markdown content above, then paste it into a text file.")
        
        export_btn.click(
            fn=export_report,
            outputs=None
        )
        
        # Instructions
        gr.Markdown("### 🎯 How to use:")
        gr.Markdown("""
        1. **Upload** a medical report image 
        2. **Click Analyze** to get AI-powered medical insights  
        3. **For multiple images:** Get a single combined analysis of all reports
        4. **Always consult** with healthcare professionals for medical decisions
        """)
    
    return demo

# Create and launch the interface
if __name__ == "__main__":
    demo = create_gradio_interface()
    
    # For Hugging Face Spaces, use this launch method
    demo.launch(
        server_name="0.0.0.0" if os.getenv("SPACE_ID") else None,
        share=False,
        show_error=True,
        debug=False
    )