flask_app.py

"""
EWAAST: Premium Flask Application

A high-fidelity web application for equitable wound assessment,
matching the official Google MedGemma aesthetic standards.

Based on the architecture of:
- google/ehr-navigator-agent-with-medgemma
- google/appoint-ready
"""

import os
import json
import base64
from io import BytesIO
from flask import Flask, render_template, request, jsonify
from PIL import Image
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Import our existing agent modules
from src.agent.classifier import MSTClassifier
from src.agent.reasoning import WoundAssessmentAgent

app = Flask(__name__, 
            template_folder='templates',
            static_folder='static')

# Initialize agents (lazy loading for HF Spaces)
mst_classifier = None
wound_agent = None


def get_classifier():
    """Lazy-load the MST classifier."""
    global mst_classifier
    if mst_classifier is None:
        mst_classifier = MSTClassifier()
    return mst_classifier


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

# ===== STATIC ASSETS (for React /assets/ path) =====

@app.route('/assets/<path:filename>')
def serve_assets(filename):
    """Serve files from the assets folder (patient images, JSON data)."""
    return app.send_static_file(f'assets/{filename}')


# ===== PAGES =====

@app.route('/')
def index():
    """Landing page with hero section."""
    return render_template('index.html')


@app.route('/assess')
def assess_page():
    """Main assessment interface."""
    return render_template('assess.html')


# ===== API ENDPOINTS =====

@app.route('/api/classify', methods=['POST'])
def api_classify():
    """
    Classify skin tone from uploaded image.
    
    Returns:
        JSON with MST value (1-10) and category
    """
    try:
        # Get image from request
        if 'image' not in request.files:
            return jsonify({'error': 'No image provided'}), 400
        
        file = request.files['image']
        image = Image.open(file.stream)
        
        # Classify skin tone
        classifier = get_classifier()
        result = classifier.classify(image)
        
        return jsonify({
            'success': True,
            'mst_value': result.value,
            'mst_category': result.category.value,
            'confidence': result.confidence,
            'visual_guidance': result.visual_guidance
        })
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/assess', methods=['POST'])
def api_assess():
    """
    Perform wound assessment with MST context.
    
    Expects:
        - image: Wound image file
        - context: Patient clinical context (optional)
        - mst_value: Pre-classified MST value (optional)
    
    Returns:
        JSON with stage, rationale, care plan, and visual guidance
    """
    try:
        # Get image
        if 'image' not in request.files:
            return jsonify({'error': 'No image provided'}), 400
        
        file = request.files['image']
        image = Image.open(file.stream)
        
        # Get optional context
        context = request.form.get('context', '')
        
        # Perform assessment
        agent = get_agent()
        assessment = agent.assess(image, context)
        
        return jsonify({
            'success': True,
            'stage': assessment.stage.value,
            'mst_value': assessment.mst_result.value,
            'mst_category': assessment.mst_result.category.value,
            'visual_guidance': assessment.mst_result.visual_guidance,
            'rationale': assessment.rationale,
            'care_plan': assessment.care_plan,
            'urgency': assessment.urgency,
            'confidence': assessment.confidence
        })
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/health')
def health_check():
    """Health check endpoint for HF Spaces."""
    return jsonify({
        'status': 'healthy',
        'version': '1.0.0',
        'model': 'MedGemma 1.5 4B (EWAAST Fine-tuned)'
    })


if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=True)