"""
EWAAST: Equitable Wound Assessment for All Skin Tones

Main Gradio application for the MedGemma Impact Challenge.
Uses MedGemma 1.5 4B with Monk Skin Tone (MST) awareness.

Features:
- Wound Assessment Mode: Upload images for AI-powered staging
- Training Mode: Interactive quiz to improve nurse confidence on diverse skin tones
"""

import os
import gradio as gr
from PIL import Image
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Check for demo mode
DEMO_MODE = os.environ.get("DEMO_MODE", "false").lower() == "true"

# Import the wound assessment agent
from src.agent.reasoning import WoundAssessmentAgent, WoundStage
from src.agent.classifier import MSTCategory

# Import training engine
from src.training.engine import TrainingEngine, TrainingCase

# Import Purpose-T form
from src.forms.purpose_t import (
    PurposeTForm, PatientInfo, SkinAssessment, VisualAssessment,
    PalpationAssessment, RiskFactors, AIAssessment,
    AssessmentLocation, SuggestedStage, get_mst_specific_questions, get_care_recommendations
)

# Initialize agent (lazy loading)
_agent = None
_training_engine = None
_current_case = None


def get_agent() -> WoundAssessmentAgent:
    """Lazy-load the wound assessment agent."""
    global _agent
    if _agent is None:
        _agent = WoundAssessmentAgent()
    return _agent


def get_training_engine() -> TrainingEngine:
    """Lazy-load the training engine."""
    global _training_engine
    if _training_engine is None:
        _training_engine = TrainingEngine(weight_deep_tones=True)
    return _training_engine


def format_assessment(assessment) -> str:
    """Format the WoundAssessment into a readable markdown string."""
    
    # Urgency badge colors
    urgency_badges = {
        "immediate": "🔴 **IMMEDIATE**",
        "urgent": "🟠 **URGENT**",
        "standard": "🟡 **STANDARD**",
        "routine": "🟢 **ROUTINE**"
    }
    
    # Stage descriptions
    stage_icons = {
        WoundStage.STAGE_1: "1️⃣",
        WoundStage.STAGE_2: "2️⃣",
        WoundStage.STAGE_3: "3️⃣",
        WoundStage.STAGE_4: "4️⃣",
        WoundStage.UNSTAGEABLE: "❓",
        WoundStage.DEEP_TISSUE_INJURY: "🟣",
        WoundStage.NOT_A_PRESSURE_ULCER: "✅"
    }
    
    # MST category colors
    mst_colors = {
        MSTCategory.LIGHT: "🏻",
        MSTCategory.MEDIUM: "🏽",
        MSTCategory.DEEP: "🏿"
    }
    
    mst = assessment.mst_result
    stage_icon = stage_icons.get(assessment.stage, "❓")
    mst_color = mst_colors.get(mst.category, "🏽")
    urgency_badge = urgency_badges.get(assessment.urgency, "🟡 **STANDARD**")
    
    # Format care plan with bullet points
    care_plan_formatted = assessment.care_plan.replace("\\n", "\n")
    
    output = f"""
## 🩹 Wound Assessment Report

---

### Skin Tone Analysis {mst_color}

| Attribute | Value |
|-----------|-------|
| **Monk Scale (MST)** | {mst.value}/10 |
| **Category** | {mst.category.value.title()} |
| **Confidence** | {mst.confidence:.0%} |

**Visual Guidance Applied:**
> {mst.visual_guidance}

---

### Clinical Assessment

| | |
|-|-|
| **Stage** | {stage_icon} **{assessment.stage.value}** |
| **Urgency** | {urgency_badge} |
| **Confidence** | {assessment.confidence:.0%} |

**Clinical Rationale:**
{assessment.rationale}

---

### 📋 Recommended Care Plan

{care_plan_formatted}

---

*⚠️ This is an AI-assisted assessment using MST-aware criteria. Always consult a qualified healthcare professional for clinical decisions.*
"""
    
    if DEMO_MODE:
        output += "\n\n*🎭 Running in DEMO MODE - using simulated responses*"
    
    return output


def process_wound_image(image: Image.Image, description: str) -> str:
    """Main processing function for the Gradio interface."""
    if image is None:
        return "❌ **Please upload an image of the wound.**\n\nYou can upload a photo or use your webcam."
    
    try:
        agent = get_agent()
        assessment = agent.assess(image, description or "")
        return format_assessment(assessment)
        
    except Exception as e:
        return f"""
## ❌ Assessment Error

An error occurred during assessment:

```
{str(e)}
```

Please try again with a different image, or check that MedGemma is properly configured.
"""


# ===== TRAINING MODE FUNCTIONS =====

def get_new_training_case():
    """Get a new training case and return display elements."""
    global _current_case
    
    engine = get_training_engine()
    _current_case = engine.get_random_case()
    
    # Create placeholder image based on skin tone
    img = Image.new('RGB', (300, 200), _current_case.image_placeholder_color)
    
    # Format case display
    case_display = f"""
## 📋 Case #{_current_case.case_id[-4:]}

**Clinical Context:**
> {_current_case.clinical_context}

**Visual Description:**
> {_current_case.visual_description}

---

*Difficulty: {_current_case.difficulty.upper()}*
"""
    
    # Get session stats
    stats = engine.get_session_stats()
    stats_display = f"Cases: {stats['cases_completed']} | Avg Score: {stats['average_score']:.0f}%"
    
    return (
        img,  # placeholder image
        case_display,  # case markdown
        5,  # reset MST slider to middle
        "Stage 1",  # reset stage dropdown
        50,  # reset confidence to 50%
        "*Submit your answer to see feedback...*",  # clear feedback
        stats_display  # update stats
    )


def submit_training_answer(mst_guess: int, stage_guess: str, confidence: int):
    """Evaluate the user's answer and provide feedback."""
    global _current_case
    
    if _current_case is None:
        return "❌ No active case. Click 'Next Case' to start.", ""
    
    engine = get_training_engine()
    
    feedback = engine.evaluate_submission(
        case=_current_case,
        user_stage=stage_guess,
        user_mst=mst_guess,
        user_confidence=confidence
    )
    
    # Build feedback display
    result_display = f"""
## 📊 Results

### Score: **{feedback.score}/100** {'🎉' if feedback.score >= 80 else '📈' if feedback.score >= 50 else '📚'}

---

### Stage Assessment
{feedback.stage_feedback}

### Skin Tone Assessment  
{feedback.mst_feedback}

---

### 📚 Visual Cue Lesson

{feedback.visual_cue_lesson}

---

### 🧠 Confidence Analysis

{feedback.confidence_analysis}
"""
    
    # Update stats
    stats = engine.get_session_stats()
    stats_display = f"Cases: {stats['cases_completed']} | Avg Score: {stats['average_score']:.0f}%"
    
    return result_display, stats_display


def reset_training_session():
    """Reset the training session."""
    engine = get_training_engine()
    engine.reset_session()
    return "*Session reset. Click 'Next Case' to begin.*", "Cases: 0 | Avg Score: 0%"


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

TITLE = "🏥 EWAAST: Equitable Wound Assessment for All Skin Tones"

DESCRIPTION = """
**EWAAST** uses MedGemma 1.5 combined with the **Monk Skin Tone (MST) Scale** 
to provide equitable wound assessment across all skin tones.
"""

# Build Gradio interface with tabs
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
    ),
    title=TITLE,
    css="""
    .gradio-container {
        max-width: 1200px !important;
    }
    .markdown-text h2 {
        border-bottom: 2px solid #4A90D9;
        padding-bottom: 8px;
    }
    .training-case {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 15px;
        border-radius: 10px;
        color: white;
    }
    """
) as demo:
    
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)
    
    with gr.Tabs():
        # ===== TAB 1: WOUND ASSESSMENT =====
        with gr.Tab("🔍 Wound Assessment"):
            gr.Markdown("""
### How to Use

1. **Upload** an image of the wound (including some healthy surrounding skin)
2. **Describe** any symptoms (optional: pain, warmth, duration)
3. **Receive** an MST-aware clinical assessment

> ⚠️ **Disclaimer**: This is a research demonstration. Not for clinical use without professional oversight.
            """)
            
            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    gr.Markdown("### 📸 Upload Wound Image")
                    image_input = gr.Image(
                        label="Wound Image",
                        type="pil",
                        sources=["upload", "webcam"],
                        height=350
                    )
                    
                    description_input = gr.Textbox(
                        label="🗣️ Patient Description (optional)",
                        placeholder="e.g., 'Painful area on heel, noticed warmth for 2 days'",
                        lines=3,
                        max_lines=5
                    )
                    
                    with gr.Row():
                        clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                        submit_btn = gr.Button("🔍 Assess Wound", variant="primary", size="lg")
                
                with gr.Column(scale=1):
                    gr.Markdown("### 📋 Assessment Result")
                    output = gr.Markdown(
                        value="*Upload an image to begin assessment...*",
                        label="Assessment"
                    )
            
            # Event handlers for assessment tab
            submit_btn.click(
                fn=process_wound_image,
                inputs=[image_input, description_input],
                outputs=output
            )
            
            clear_btn.click(
                fn=lambda: (None, "", "*Upload an image to begin assessment...*"),
                outputs=[image_input, description_input, output]
            )
        
        # ===== TAB 2: TRAINING MODE =====
        with gr.Tab("🎓 Training Mode"):
            gr.Markdown("""
### 🎯 React to Deep Tones - Nurse Training

This interactive training module addresses the **confidence gap** identified in NHS research:
nurses report lower confidence when assessing wounds on darker skin tones.

**How it works:**
1. Study the clinical scenario
2. Identify the **Monk Skin Tone (MST)** value
3. Determine the **Pressure Ulcer Stage**
4. Rate your **confidence** level
5. Receive immediate feedback with educational lessons

> 💡 Cases are **weighted towards MST 7-10** to specifically address the confidence gap.
            """)
            
            with gr.Row():
                # Session stats
                session_stats = gr.Textbox(
                    value="Cases: 0 | Avg Score: 0%",
                    label="📊 Session Progress",
                    interactive=False
                )
                reset_btn = gr.Button("🔄 Reset Session", variant="secondary", size="sm")
            
            with gr.Row(equal_height=True):
                # LEFT: Case presentation
                with gr.Column(scale=1):
                    gr.Markdown("### 📋 Clinical Case")
                    
                    case_image = gr.Image(
                        label="Wound Visualization",
                        height=200,
                        interactive=False
                    )
                    
                    case_display = gr.Markdown(
                        value="*Click 'Next Case' to begin training...*"
                    )
                    
                    next_case_btn = gr.Button("➡️ Next Case", variant="primary", size="lg")
                
                # RIGHT: Answer input
                with gr.Column(scale=1):
                    gr.Markdown("### ✏️ Your Assessment")
                    
                    mst_slider = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="🎨 Monk Skin Tone (MST) - What skin tone do you observe?",
                        info="1 = Very Light, 10 = Very Deep"
                    )
                    
                    stage_dropdown = gr.Dropdown(
                        choices=[
                            "Stage 1",
                            "Stage 2", 
                            "Stage 3",
                            "Stage 4",
                            "Deep Tissue Injury",
                            "Unstageable"
                        ],
                        value="Stage 1",
                        label="📊 Pressure Ulcer Stage"
                    )
                    
                    confidence_slider = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=50,
                        step=5,
                        label="🧠 Your Confidence Level",
                        info="How confident are you in your assessment?"
                    )
                    
                    submit_answer_btn = gr.Button("✅ Submit Answer", variant="primary", size="lg")
            
            # Feedback section
            gr.Markdown("---")
            gr.Markdown("### 📝 Feedback")
            feedback_display = gr.Markdown(
                value="*Submit your answer to see feedback...*"
            )
            
            # Event handlers for training tab
            next_case_btn.click(
                fn=get_new_training_case,
                outputs=[
                    case_image,
                    case_display,
                    mst_slider,
                    stage_dropdown,
                    confidence_slider,
                    feedback_display,
                    session_stats
                ]
            )
            
            submit_answer_btn.click(
                fn=submit_training_answer,
                inputs=[mst_slider, stage_dropdown, confidence_slider],
                outputs=[feedback_display, session_stats]
            )
            
            reset_btn.click(
                fn=reset_training_session,
                outputs=[feedback_display, session_stats]
            )
        
        # ===== TAB 3: PURPOSE-T FORM =====
        with gr.Tab("📋 Purpose-T Form"):
            gr.Markdown("""
### 📋 Digital Purpose-T Assessment

This **inclusive digital form** replaces traditional Waterlow/Braden paper forms with **MST-aware** assessment fields.

**Key Features:**
- ✅ Monk Skin Tone (MST) is **REQUIRED** - forces conscious assessment
- ✅ Dynamic questions based on skin tone
- ✅ Palpation assessment (temperature, texture, induration)
- ✅ Integrated risk scoring
- ✅ FHIR-compatible export
            """)
            
            with gr.Accordion("📝 Section 1: Patient & Skin Information", open=True):
                with gr.Row():
                    patient_id = gr.Textbox(label="Patient ID", placeholder="Enter patient ID")
                    assessor_name = gr.Textbox(label="Assessor Name", placeholder="Your name")
                
                with gr.Row():
                    mst_slider_form = gr.Slider(
                        minimum=1, maximum=10, value=5, step=1,
                        label="🎨 Monk Skin Tone (MST) - REQUIRED",
                        info="Observe healthy skin near affected area. 1=Very Light, 10=Very Deep"
                    )
                    location_dropdown = gr.Dropdown(
                        choices=[loc.value for loc in AssessmentLocation],
                        value=AssessmentLocation.SACRUM.value,
                        label="📍 Assessment Location"
                    )
                
                mst_guidance_display = gr.Markdown(
                    value="*Select MST to see visual guidance...*"
                )
            
            with gr.Accordion("👁️ Section 2: Visual Assessment", open=True):
                visual_guidance_banner = gr.Markdown(value="")
                
                with gr.Row():
                    discoloration_present = gr.Checkbox(label="Discoloration present?", value=False)
                    skin_intact = gr.Checkbox(label="Skin intact?", value=True)
                
                discoloration_desc = gr.Dropdown(
                    choices=["Red/Pink", "Darkening from baseline", "Purple/violet tinge", 
                             "Deep purple/maroon", "Ashen/greyish", "Blue-black"],
                    label="Describe discoloration color",
                    visible=False
                )
                
                wound_description = gr.Textbox(
                    label="Wound Description (if skin not intact)",
                    placeholder="Describe wound characteristics...",
                    lines=2,
                    visible=False
                )
            
            with gr.Accordion("🖐️ Section 3: Palpation Assessment", open=True):
                gr.Markdown("**⚠️ Palpation is CRITICAL for accurate assessment on darker skin tones**")
                
                with gr.Row():
                    temp_assessed = gr.Checkbox(label="Temperature assessed?", value=False)
                    temp_difference = gr.Radio(
                        choices=["Warmer", "Cooler", "Same"],
                        label="Temperature difference?",
                        visible=False
                    )
                
                with gr.Row():
                    induration_present = gr.Checkbox(label="Induration (firmness) present?", value=False)
                    induration_severity = gr.Radio(
                        choices=["Mild", "Moderate", "Severe"],
                        label="Induration severity",
                        visible=False
                    )
                
                with gr.Row():
                    boggy_texture = gr.Checkbox(label="Boggy/spongy texture?", value=False)
                    pain_present = gr.Checkbox(label="Pain on palpation?", value=False)
                
                pain_score = gr.Slider(
                    minimum=0, maximum=10, value=0, step=1,
                    label="Pain score (0-10)",
                    visible=False
                )
            
            with gr.Accordion("⚖️ Section 4: Risk Factors (Braden Subscales)", open=False):
                with gr.Row():
                    sensory = gr.Slider(1, 4, 4, step=1, label="Sensory Perception (1-4)")
                    moisture = gr.Slider(1, 4, 4, step=1, label="Moisture (1-4)")
                    activity = gr.Slider(1, 4, 4, step=1, label="Activity (1-4)")
                
                with gr.Row():
                    mobility = gr.Slider(1, 4, 4, step=1, label="Mobility (1-4)")
                    nutrition = gr.Slider(1, 4, 4, step=1, label="Nutrition (1-4)")
                    friction = gr.Slider(1, 3, 3, step=1, label="Friction/Shear (1-3)")
                
                with gr.Row():
                    diabetes = gr.Checkbox(label="Diabetes", value=False)
                    vascular = gr.Checkbox(label="Vascular Disease", value=False)
                    previous_pi = gr.Checkbox(label="Previous Pressure Injury", value=False)
            
            with gr.Accordion("🤖 Section 5: AI-Assisted Staging (Optional)", open=False):
                ai_image = gr.Image(label="Upload wound image for AI staging", type="pil")
                ai_stage_btn = gr.Button("🔍 Get AI Staging Suggestion", variant="secondary")
                ai_result = gr.Markdown(value="*Upload image and click to get AI suggestion...*")
            
            gr.Markdown("---")
            
            with gr.Row():
                final_stage = gr.Dropdown(
                    choices=[s.value for s in SuggestedStage],
                    label="📊 Final Stage Determination",
                    value=SuggestedStage.INTACT_AT_RISK.value
                )
                generate_btn = gr.Button("📄 Generate Report", variant="primary", size="lg")
            
            report_output = gr.Markdown(value="*Complete the form and click Generate Report...*")
            
            with gr.Row():
                download_pdf_btn = gr.Button("⬇️ Download PDF", variant="secondary")
                download_fhir_btn = gr.Button("⬇️ Download FHIR JSON", variant="secondary")
            
            # === PURPOSE-T EVENT HANDLERS ===
            
            def update_mst_guidance(mst_val):
                """Update guidance when MST changes."""
                questions = get_mst_specific_questions(mst_val)
                guidance = f"""
### MST {mst_val} Visual Guidance

{questions['banner'] if questions['banner'] else ''}

**{questions['primary_question']}**

*Color options to look for: {', '.join([c.value for c in questions['color_options']])}*
"""
                return guidance
            
            def toggle_discoloration(is_present):
                return gr.update(visible=is_present)
            
            def toggle_wound_desc(is_intact):
                return gr.update(visible=not is_intact)
            
            def toggle_temp(is_assessed):
                return gr.update(visible=is_assessed)
            
            def toggle_induration(is_present):
                return gr.update(visible=is_present)
            
            def toggle_pain(has_pain):
                return gr.update(visible=has_pain)
            
            def get_ai_staging(image):
                if image is None:
                    return "❌ Please upload an image first."
                try:
                    agent = get_agent()
                    assessment = agent.assess(image, "")
                    return f"""
### AI Staging Suggestion

**Stage:** {assessment.stage.value}  
**Confidence:** {assessment.confidence:.0%}

**Rationale:**
{assessment.rationale}

> ⚠️ This is an AI suggestion. Clinician should verify and may override.
"""
                except Exception as e:
                    return f"❌ AI staging error: {str(e)}"
            
            def generate_report(
                pid, assessor, mst_val, location, 
                discol_present, intact, discol_desc, wound_desc,
                temp_assessed, temp_diff, indur_present, indur_sev, boggy, pain, pain_sc,
                sens, moist, act, mob, nutr, fric, diab, vasc, prev,
                stage
            ):
                from datetime import datetime
                
                # Build form
                form = PurposeTForm(
                    patient=PatientInfo(
                        patient_id=pid or "UNKNOWN",
                        assessor_name=assessor or "Not specified",
                        assessment_date=datetime.now()
                    ),
                    skin=SkinAssessment(
                        mst_value=mst_val,
                        assessment_location=AssessmentLocation(location)
                    ),
                    visual=VisualAssessment(
                        discoloration_present=discol_present,
                        skin_intact=intact,
                        wound_present=not intact,
                        wound_description=wound_desc or ""
                    ),
                    palpation=PalpationAssessment(
                        temperature_assessed=temp_assessed,
                        temperature_difference=temp_diff.lower() if temp_diff else None,
                        induration_present=indur_present,
                        induration_severity=indur_sev.lower() if indur_sev else None,
                        boggy_texture=boggy,
                        pain_on_palpation=pain,
                        pain_score=pain_sc if pain else None
                    ),
                    risk=RiskFactors(
                        sensory_perception=sens,
                        moisture=moist,
                        activity=act,
                        mobility=mob,
                        nutrition=nutr,
                        friction_shear=fric,
                        diabetes=diab,
                        vascular_disease=vasc,
                        previous_pressure_injury=prev
                    ),
                    ai=AIAssessment(),
                    final_stage=SuggestedStage(stage),
                    care_plan=get_care_recommendations(form) if 'form' in dir() else ""
                )
                
                # Validate
                errors = form.validate()
                if errors:
                    return "## ❌ Validation Errors\n\n" + "\n".join(f"• {e}" for e in errors)
                
                # Add care plan
                form.care_plan = get_care_recommendations(form)
                
                return form.generate_report()
            
            # Wire up events
            mst_slider_form.change(update_mst_guidance, inputs=[mst_slider_form], outputs=[mst_guidance_display])
            discoloration_present.change(toggle_discoloration, inputs=[discoloration_present], outputs=[discoloration_desc])
            skin_intact.change(toggle_wound_desc, inputs=[skin_intact], outputs=[wound_description])
            temp_assessed.change(toggle_temp, inputs=[temp_assessed], outputs=[temp_difference])
            induration_present.change(toggle_induration, inputs=[induration_present], outputs=[induration_severity])
            pain_present.change(toggle_pain, inputs=[pain_present], outputs=[pain_score])
            
            ai_stage_btn.click(get_ai_staging, inputs=[ai_image], outputs=[ai_result])
            
            generate_btn.click(
                generate_report,
                inputs=[
                    patient_id, assessor_name, mst_slider_form, location_dropdown,
                    discoloration_present, skin_intact, discoloration_desc, wound_description,
                    temp_assessed, temp_difference, induration_present, induration_severity, 
                    boggy_texture, pain_present, pain_score,
                    sensory, moisture, activity, mobility, nutrition, friction,
                    diabetes, vascular, previous_pi,
                    final_stage
                ],
                outputs=[report_output]
            )
    
    # ===== FOOTER =====
    gr.Markdown("---")
    
    with gr.Accordion("ℹ️ About EWAAST", open=False):
        gr.Markdown("""
### The Problem

Current wound assessment AI tools perform poorly on darker skin tones because:
- **Training data bias**: Most dermatology datasets overrepresent lighter skin (MST 1-3)
- **Feature bias**: "Redness" as an indicator fails for deeper skin tones where inflammation presents as **purple/blue discoloration**

### The Solution

EWAAST fixes this by explicitly incorporating **Monk Skin Tone detection** into the assessment pipeline:

1. **MST Classifier** → Analyzes healthy skin to determine tone (1-10)
2. **Visual Guidance Injection** → Applies tone-specific feature detection
3. **Equitable Staging** → Uses EPUAP/NPUAP criteria adapted for all skin tones

### Monk Skin Tone Scale

| MST Range | Description | Wound Feature Guidance |
|-----------|-------------|------------------------|
| 1-3 | Light | Erythema (redness) visible |
| 4-7 | Medium | Look for warmth, subtle color change |
| 8-10 | Deep | Discoloration (purple/blue), induration, heat |

### References

- [Monk Skin Tone Scale](https://skintone.google/) - Dr. Ellis Monk & Google
- [EPUAP/NPUAP Pressure Ulcer Guidelines](https://www.epuap.org/)
- [MedGemma](https://developers.google.com/health-ai-developer-foundations) - Google Health AI
        """)
    
    with gr.Row():
        gr.Markdown(
            "**Built for the [MedGemma Impact Challenge 2026](https://www.kaggle.com/competitions/med-gemma-impact-challenge)** 🏆 | "
            "[GitHub](https://github.com/ClinyQAi/EWAAST) | "
            "[Monk Scale](https://skintone.google/)"
        )


if __name__ == "__main__":
    # Check for demo mode and notify
    if DEMO_MODE:
        print("🎭 DEMO MODE: Using simulated MST-aware responses")
    else:
        print("🚀 PRODUCTION MODE: Using MedGemma for inference")
    
    demo.launch(
        share=False,
        server_name="127.0.0.1",
        server_port=8080
    )