import os
import streamlit as st
from PIL import Image
import google.generativeai as genai
import json
from datetime import datetime

# Configure Gemini API key (ideally load from environment variable in production)
# Using st.secrets for better security in Streamlit
api_key = "AIzaSyDWm-Z7X4H-b41r2ZRN61UfABdv81D2Gxo"  # In production, use st.secrets["GEMINI_API_KEY"]
genai.configure(api_key=api_key)

# --- Generation Configuration and Chat Session Setup ---
generation_config = {
    "temperature": 0.8,  # Slightly reduced for more consistent medical responses
    "top_p": 0.95,
    "top_k": 40,
    "max_output_tokens": 8192,
    "response_mime_type": "application/json",  # Changed to JSON for structured parsing
}

model = genai.GenerativeModel(
    model_name="gemini-1.5-flash",
    generation_config=generation_config,
)

# Define the system prompt for structured JSON responses
SYSTEM_PROMPT = """
As a highly skilled medical practitioner specializing in image analysis, 
you are tasked with examining medical images for a renowned hospital. 
Your expertise is crucial in identifying any anomalies, diseases, or health issues 
that may be present in the images.

Your responsibilities:
1. **Detailed Analysis**: Thoroughly examine the image for abnormalities.
2. **Analysis Report**: Document findings clearly.
3. **Recommendations**: Suggest necessary tests or treatments.
4. **Treatments**: Provide possible treatments for better recovery.
5. **Risk Level**: Indicate the severity level as "Low", "Medium", "High", or "Critical".

**Important Notes**:
- Only respond if the image is related to human health.
- If the image is unclear, state: "Unable to determine based on the uploaded image."
- Include a disclaimer about consulting with a doctor.

Format the response as a JSON object with these fields:
{
  "image_type": "String describing the type of medical image",
  "detailed_analysis": "Thorough analysis of visible features",
  "analysis_report": "Summary of findings",
  "recommendations": "Suggested tests or follow-ups",
  "treatments": "Possible treatment options",
  "risk_level": "Low/Medium/High/Critical",
  "confidence_score": 75,
  "areas_of_concern": ["List of specific areas or features of concern"]
}
"""


# --- Functions for UI Components ---

def display_sidebar():
    with st.sidebar:
        st.title("🏥 Cure Connect")
        st.subheader("Your AI Medical Assistant")

        st.markdown("---")
        st.markdown("### About This Tool")
        st.markdown("""
        Cure Connect uses advanced AI to analyze medical images and provide potential insights.
        **IMPORTANT:** This tool is for educational purposes only and should not replace professional medical advice.
        """)

        st.markdown("---")
        st.markdown("### Supported Image Types")
        st.markdown("""
        - X-rays
        - MRI scans
        - CT scans
        - Ultrasound images
        - Dermatological photos
        """)

        st.markdown("---")
        st.markdown("### Usage Tips")
        st.markdown("""
        1. Upload a clear medical image
        2. Click "Analyze Image"
        3. Review the AI analysis
        4. Share results with your healthcare provider
        """)


def create_analysis_card(title, content, color):
    st.markdown(
        f"""
        <div style="
            background-color: {color};
            padding: 20px;
            border-radius: 10px;
            margin-bottom: 10px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        ">
            <h3 style="color: white;">{title}</h3>
            <p style="color: white;">{content}</p>
        </div>
        """,
        unsafe_allow_html=True
    )


def display_risk_gauge(risk_level, confidence):
    # Map risk levels to numerical values
    risk_values = {"Low": 1, "Medium": 2, "High": 3, "Critical": 4}
    risk_value = risk_values.get(risk_level, 0)

    # Colors for different risk levels
    colors = {
        "Low": "#4CAF50",  # Green
        "Medium": "#FFC107",  # Amber
        "High": "#FF5722",  # Deep Orange
        "Critical": "#F44336"  # Red
    }

    col1, col2 = st.columns(2)

    with col1:
        st.markdown("### Risk Assessment")
        st.markdown(
            f"""
            <div style="
                background-color: {colors.get(risk_level, "#607D8B")};
                color: white;
                text-align: center;
                padding: 20px;
                border-radius: 10px;
                font-size: 24px;
                font-weight: bold;
            ">
                {risk_level}
            </div>
            """,
            unsafe_allow_html=True
        )

    with col2:
        st.markdown("### AI Confidence")
        # Ensure confidence is an integer before using it
        if isinstance(confidence, str):
            try:
                confidence = int(float(confidence))
            except (ValueError, TypeError):
                confidence = 75  # Default value if conversion fails

        st.progress(confidence / 100)
        st.markdown(f"<h1 style='text-align: center;'>{confidence}%</h1>", unsafe_allow_html=True)


def generate_pdf_report(analysis_data, image):
    # In a real app, you would implement PDF generation here
    st.download_button(
        label="📄 Download PDF Report",
        data=f"Sample report for {analysis_data['image_type']} - {datetime.now().strftime('%Y-%m-%d')}",
        file_name="medical_report.txt",
        mime="text/plain"
    )


# --- Main Application UI ---
st.set_page_config(page_title='Cure Connect - Medical Image Analytics',
                   page_icon='🏥',
                   layout='wide',
                   initial_sidebar_state='expanded')

# Apply custom CSS for better styling
st.markdown("""
<style>
    .main {
        padding: 2rem;
        background-color: #f8f9fa;
    }
    .stButton button {
        background-color: #4285F4;
        color: white;
        border-radius: 10px;
        padding: 0.5rem 1rem;
        font-weight: bold;
    }
    h1, h2, h3 {
        color: #01579B;
    }
    .stProgress > div > div {
        background-color: #4285F4;
    }
    .disclaimer {
        background-color: #FFECB3;
        padding: 10px;
        border-radius: 5px;
        border-left: 5px solid #FFC107;
    }
</style>
""", unsafe_allow_html=True)

# Display sidebar
display_sidebar()

# Main content area
st.title("🩺 Cure Connect - Medical Image Analysis")
st.subheader("AI-assisted diagnostic support platform")

st.markdown("""
This platform uses advanced AI to analyze medical images and provide potential insights.
Upload your medical scan below to get started.
""")

# Create tabs for different sections
tab1, tab2, tab3 = st.tabs(["📊 Analysis", "ℹ️ How It Works", "❓ FAQ"])

with tab1:
    # File upload with enhanced UI
    col1, col2 = st.columns([1, 1])

    with col1:
        st.markdown("### Upload Medical Image")
        uploaded_file = st.file_uploader('Select an image file',
                                         type=['png', 'jpg', 'jpeg'],
                                         help="Upload a clear, high-resolution medical image")

        analyze_col, clear_col = st.columns([1, 1])
        with analyze_col:
            analyze_button = st.button('🔍 Analyze Image', use_container_width=True)
        with clear_col:
            clear_button = st.button('🗑️ Clear Results', use_container_width=True)

    # Analysis section
    if analyze_button and uploaded_file:
        try:
            # Open and display the uploaded image
            img = Image.open(uploaded_file)

            with col2:
                st.markdown("### Uploaded Medical Image")
                st.image(img, use_column_width=True)
                st.caption("The AI model will analyze this image to identify potential areas of concern.")

            # Show processing indicator
            with st.spinner('Analyzing image... Please wait'):
                # In a production app, you would handle the JSON response properly
                # For this prototype, we'll simulate a structured response
                try:
                    # Generate analysis using the system prompt and the uploaded image
                    response = model.generate_content([SYSTEM_PROMPT, img])

                    # Try to parse as JSON first
                    try:
                        # Parse JSON response
                        analysis_data = json.loads(response.text)
                    except json.JSONDecodeError:
                        # If not valid JSON, create a structured dictionary from text
                        analysis_data = {
                            "image_type": "Medical scan",
                            "detailed_analysis": response.text[:500],
                            "analysis_report": "See detailed analysis",
                            "recommendations": "Consult with a healthcare professional",
                            "treatments": "To be determined by your doctor",
                            "risk_level": "Medium",
                            "confidence_score": 75,
                            "areas_of_concern": ["Unable to parse specific areas from text response"]
                        }

                    # Ensure confidence_score is an integer
                    if isinstance(analysis_data["confidence_score"], str):
                        try:
                            analysis_data["confidence_score"] = int(float(analysis_data["confidence_score"]))
                        except (ValueError, TypeError):
                            analysis_data["confidence_score"] = 75  # Default if conversion fails

                    # Display results in a structured format
                    st.markdown("## Analysis Results")
                    st.markdown(f"### Image Type: {analysis_data['image_type']}")

                    # Display risk gauge
                    display_risk_gauge(analysis_data['risk_level'], analysis_data['confidence_score'])

                    # Display color-coded analysis sections
                    st.markdown("### Key Findings")
                    create_analysis_card("Detailed Analysis", analysis_data['detailed_analysis'], "#0277BD")
                    create_analysis_card("Analysis Report", analysis_data['analysis_report'], "#00897B")
                    create_analysis_card("Recommendations", analysis_data['recommendations'], "#7B1FA2")
                    create_analysis_card("Treatments", analysis_data['treatments'], "#C2185B")

                    # Areas of concern
                    st.markdown("### Areas of Concern")
                    for i, area in enumerate(analysis_data['areas_of_concern']):
                        st.markdown(f"- **Area {i + 1}:** {area}")

                    # Disclaimer
                    st.markdown("""
                    <div class="disclaimer">
                        <strong>⚠️ IMPORTANT DISCLAIMER:</strong> This analysis is for informational purposes only 
                        and should not be considered medical advice. Always consult with a qualified healthcare 
                        professional for proper diagnosis and treatment options.
                    </div>
                    """, unsafe_allow_html=True)

                    # Generate PDF report option
                    st.markdown("### Export Options")
                    generate_pdf_report(analysis_data, img)

                except Exception as e:
                    st.error(f"Error processing response: {str(e)}")

        except Exception as e:
            st.error(f"Analysis failed: {str(e)}")
            st.error("Please ensure you're using a valid medical image format (JPEG/PNG)")

with tab2:
    st.markdown("## How Cure Connect Works")

    st.markdown("""
    ### 1. Image Analysis
    When you upload a medical image, our AI system processes it using advanced computer vision techniques
    to identify patterns, abnormalities, and potential areas of concern.

    ### 2. Medical Knowledge Base
    The analysis is informed by a vast medical knowledge base that includes information about various
    conditions, diseases, and their visual presentations in medical imaging.

    ### 3. Risk Assessment
    Based on the analysis, the system provides a risk assessment categorized as:
    - **Low Risk** (Green): Minor or no abnormalities detected
    - **Medium Risk** (Amber): Notable findings that should be discussed with a healthcare provider
    - **High Risk** (Orange): Significant findings that require prompt medical attention
    - **Critical Risk** (Red): Urgent findings that may require immediate medical intervention

    ### 4. Confidence Score
    The confidence score indicates how certain the AI is about its analysis based on image quality,
    clarity of findings, and comparison with known patterns.
    """)

    # Add a simple diagram
    st.markdown("""
    ```
    ┌─────────────┐    ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
    │  Image      │    │  AI         │    │  Analysis   │    │  Results    │
    │  Upload     │ -> │  Processing │ -> │  Generation │ -> │  Display    │
    └─────────────┘    └─────────────┘    └─────────────┘    └─────────────┘
    ```
    """)

    st.warning("""
    **Remember:** While our system uses advanced technology to analyze medical images,
    it is designed to be a supportive tool for healthcare professionals, not a replacement
    for proper medical consultation and diagnosis.
    """)

with tab3:
    st.markdown("## Frequently Asked Questions")

    faq_data = [
        {
            "question": "How accurate is the AI analysis?",
            "answer": "The AI analysis provides an estimated confidence score with each result. However, accuracy varies based on image quality, the type of medical condition, and other factors. Always consult with a healthcare professional for accurate diagnosis."
        },
        {
            "question": "Is my medical data secure?",
            "answer": "We take data privacy seriously. Your uploaded images are processed securely and not stored permanently unless explicitly requested. All processing is done in compliance with healthcare data protection standards."
        },
        {
            "question": "What types of medical images can I upload?",
            "answer": "The system can analyze various medical imaging types including X-rays, MRIs, CT scans, ultrasound images, and dermatological photos. The clearer and higher resolution the image, the better the analysis will be."
        },
        {
            "question": "How should I use the results?",
            "answer": "Consider the results as informational insights that you can discuss with your healthcare provider. The analysis is meant to supplement, not replace, professional medical advice."
        },
        {
            "question": "Can I share the analysis with my doctor?",
            "answer": "Yes! You can download a PDF report of the analysis to share with your healthcare provider. This can be a helpful starting point for discussions about your health concerns."
        }
    ]

    for i, faq in enumerate(faq_data):
        with st.expander(f"{i + 1}. {faq['question']}"):
            st.markdown(faq['answer'])

    st.markdown("""
    ### Have more questions?
    If you have any other questions about using Cure Connect, please contact our support team
    at support@cureconnect.example.com.
    """)