app.py

#!/usr/bin/env python3
"""
Flask API for Mi Band 4 connector
Provides REST API endpoints for all Mi Band 4 functionalities
"""

import asyncio
import json
import logging
import struct
import threading
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import os
import uuid
from functools import wraps

from flask import Flask, request, jsonify
from flask_cors import CORS

# Import our Mi Band 4 library
from mib4 import MiBand4, AlertType, MusicState, find_miband_devices, HeartRateData

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# Flask app setup
app = Flask(__name__)
app.secret_key = os.getenv('FLASK_SECRET_KEY', 'miband4-api-secret-key-change-me')
CORS(app, supports_credentials=True)

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global event loop for async operations
_event_loop = None
_loop_thread = None


def get_event_loop():
    """Get or create the global event loop"""
    global _event_loop, _loop_thread
    
    if _event_loop is None or _event_loop.is_closed():
        def run_loop():
            global _event_loop
            _event_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(_event_loop)
            _event_loop.run_forever()
        
        _loop_thread = threading.Thread(target=run_loop, daemon=True)
        _loop_thread.start()
        
        # Wait for loop to be created
        import time
        while _event_loop is None:
            time.sleep(0.01)
    
    return _event_loop


class DeviceManager:
    """Manages Mi Band 4 device connections and sessions"""
    
    def __init__(self):
        self.devices = {}  # session_id -> device_config
        self.connections = {}  # session_id -> MiBand4 instance
    
    def register_device(self, mac_address: str, auth_key: Optional[str] = None) -> str:
        """Register a device and return session ID"""
        session_id = str(uuid.uuid4())
        self.devices[session_id] = {
            'mac_address': mac_address.upper(),
            'auth_key': auth_key,
            'registered_at': datetime.now(),
            'last_used': datetime.now()
        }
        logger.info(f"Device registered: {mac_address} with session {session_id}")
        return session_id
    
    def get_device_config(self, session_id: str) -> Optional[Dict]:
        """Get device configuration by session ID"""
        config = self.devices.get(session_id)
        if config:
            config['last_used'] = datetime.now()
        return config
    
    async def get_connection(self, session_id: str) -> Optional[MiBand4]:
        """Get or create Mi Band 4 connection"""
        if session_id in self.connections:
            band = self.connections[session_id]
            if band.connected:
                return band
        
        config = self.get_device_config(session_id)
        if not config:
            return None
        
        # Create new connection
        band = MiBand4(
            config['mac_address'], 
            config['auth_key'], 
            timeout=15.0
        )
        
        try:
            success = await band.connect()
            if success:
                self.connections[session_id] = band
                
                # Try authentication if auth key is available
                if config['auth_key']:
                    await band.authenticate()
                
                return band
        except Exception as e:
            logger.error(f"Connection failed for {session_id}: {e}")
        
        return None
    
    async def disconnect(self, session_id: str):
        """Disconnect device"""
        if session_id in self.connections:
            try:
                await self.connections[session_id].disconnect()
            except Exception as e:
                logger.debug(f"Disconnect error (ignored): {e}")
            del self.connections[session_id]
    
    def remove_device(self, session_id: str):
        """Remove device registration"""
        if session_id in self.devices:
            del self.devices[session_id]
        if session_id in self.connections:
            del self.connections[session_id]


# Global device manager instance
dm = DeviceManager()


def async_route(f):
    """Decorator to run async functions in Flask routes using the global event loop"""
    @wraps(f)
    def wrapper(*args, **kwargs):
        loop = get_event_loop()
        
        # Run the coroutine in the global event loop
        future = asyncio.run_coroutine_threadsafe(f(*args, **kwargs), loop)
        
        try:
            # Wait for result with timeout
            result = future.result(timeout=30)  # 30 second timeout
            return result
        except Exception as e:
            logger.error(f"Async route error: {e}")
            return jsonify({'error': str(e), 'code': 'ASYNC_ERROR'}), 500
    
    return wrapper


def require_session(f):
    """Decorator to require valid session"""
    @wraps(f)
    def wrapper(*args, **kwargs):
        session_id = request.headers.get('X-Session-ID')
        
        if not session_id:
            # Try to get from JSON body as fallback
            try:
                data = request.get_json() if request.is_json else {}
                session_id = data.get('session_id') if data else None
            except:
                pass
        
        if not session_id:
            return jsonify({'error': 'Session ID required', 'code': 'NO_SESSION'}), 400
        
        config = dm.get_device_config(session_id)
        if not config:
            return jsonify({'error': 'Invalid session ID', 'code': 'INVALID_SESSION'}), 401
        
        return f(session_id, *args, **kwargs)
    return wrapper


def require_auth(f):
    """Decorator to require authentication"""
    @wraps(f)
    def wrapper(session_id, *args, **kwargs):
        config = dm.get_device_config(session_id)
        if not config or not config.get('auth_key'):
            return jsonify({
                'error': 'Authentication required for this operation',
                'code': 'AUTH_REQUIRED'
            }), 403
        
        return f(session_id, *args, **kwargs)
    return wrapper


# API Routes

@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'healthy',
        'timestamp': datetime.now().isoformat(),
        'version': '1.0.0',
        'active_sessions': len(dm.devices),
        'active_connections': len(dm.connections),
        'event_loop_running': _event_loop is not None and not _event_loop.is_closed()
    })


@app.route('/api/discover', methods=['POST'])
@async_route
async def discover_devices():
    """Discover nearby Mi Band devices"""
    try:
        devices = await find_miband_devices()
        return jsonify({
            'success': True,
            'devices': devices,
            'count': len(devices)
        })
    except Exception as e:
        logger.error(f"Discovery error: {e}")
        return jsonify({'error': str(e), 'code': 'DISCOVERY_FAILED'}), 500


@app.route('/api/register', methods=['POST'])
def register_device():
    """Register a device and get session ID"""
    data = request.get_json()
    
    if not data or 'mac_address' not in data:
        return jsonify({'error': 'MAC address required', 'code': 'MISSING_MAC'}), 400
    
    mac_address = data['mac_address'].strip().upper()
    auth_key = data.get('auth_key', '').strip() or None
    
    # Validate MAC address
    if len(mac_address) != 17 or mac_address.count(':') != 5:
        return jsonify({'error': 'Invalid MAC address format', 'code': 'INVALID_MAC'}), 400
    
    # Validate auth key if provided
    if auth_key and (len(auth_key) != 32 or not all(c in '0123456789ABCDEFabcdef' for c in auth_key)):
        return jsonify({'error': 'Invalid auth key format', 'code': 'INVALID_AUTH_KEY'}), 400
    
    session_id = dm.register_device(mac_address, auth_key)
    
    return jsonify({
        'success': True,
        'session_id': session_id,
        'mac_address': mac_address,
        'auth_enabled': bool(auth_key),
        'message': 'Device registered successfully'
    })


@app.route('/api/connect', methods=['POST'])
@require_session
@async_route
async def connect_device(session_id):
    """Connect to registered device"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Connection failed', 'code': 'CONNECTION_FAILED'}), 500
        
        config = dm.get_device_config(session_id)
        return jsonify({
            'success': True,
            'connected': True,
            'authenticated': band._is_authenticated(),
            'mac_address': config['mac_address'],
            'auth_enabled': bool(config['auth_key'])
        })
    except Exception as e:
        logger.error(f"Connection error: {e}")
        return jsonify({'error': str(e), 'code': 'CONNECTION_ERROR'}), 500


@app.route('/api/disconnect', methods=['POST'])
@require_session
@async_route
async def disconnect_device(session_id):
    """Disconnect from device"""
    try:
        await dm.disconnect(session_id)
        return jsonify({'success': True, 'message': 'Disconnected successfully'})
    except Exception as e:
        logger.error(f"Disconnect error: {e}")
        return jsonify({'error': str(e), 'code': 'DISCONNECT_ERROR'}), 500


@app.route('/api/device/info', methods=['GET'])
@require_session
@async_route
async def get_device_info(session_id):
    """Get device information"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        info = await band.get_device_info()
        return jsonify({
            'success': True,
            'device_info': {
                'firmware_version': info.firmware_version,
                'hardware_version': info.hardware_version,
                'serial_number': info.serial_number
            }
        })
    except Exception as e:
        logger.error(f"Get device info error: {e}")
        return jsonify({'error': str(e), 'code': 'DEVICE_INFO_ERROR'}), 500


@app.route('/api/device/battery', methods=['GET'])
@require_session
@async_route
async def get_battery_info(session_id):
    """Get battery information"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        battery = await band.get_battery_info()
        return jsonify({
            'success': True,
            'battery': {
                'level': battery.level,
                'status': battery.status,
                'last_level': battery.last_level,
                'last_charge': battery.last_charge.isoformat() if battery.last_charge else None,
                'last_off': battery.last_off.isoformat() if battery.last_off else None
            }
        })
    except Exception as e:
        logger.error(f"Get battery error: {e}")
        return jsonify({'error': str(e), 'code': 'BATTERY_ERROR'}), 500


@app.route('/api/device/time', methods=['GET'])
@require_session
@async_route
async def get_device_time(session_id):
    """Get device time"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        device_time = await band.get_current_time()
        return jsonify({
            'success': True,
            'device_time': device_time.isoformat(),
            'system_time': datetime.now().isoformat()
        })
    except Exception as e:
        logger.error(f"Get time error: {e}")
        return jsonify({'error': str(e), 'code': 'TIME_ERROR'}), 500


@app.route('/api/device/time', methods=['POST'])
@require_session
@require_auth
@async_route
async def set_device_time(session_id):
    """Set device time"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        data = request.get_json() or {}
        
        # Use provided time or current system time
        if 'time' in data:
            set_time = datetime.fromisoformat(data['time'])
        else:
            set_time = datetime.now()
        
        await band.set_current_time(set_time)
        
        return jsonify({
            'success': True,
            'time_set': set_time.isoformat(),
            'message': 'Time set successfully'
        })
    except Exception as e:
        logger.error(f"Set time error: {e}")
        return jsonify({'error': str(e), 'code': 'SET_TIME_ERROR'}), 500


@app.route('/api/notification', methods=['POST'])
@require_session
@async_route
async def send_notification(session_id):
    """Send notification to device"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        data = request.get_json()
        if not data:
            return jsonify({'error': 'Request data required', 'code': 'MISSING_DATA'}), 400
        
        # Parse alert type
        alert_type_map = {
            'message': AlertType.MESSAGE,
            'mail': AlertType.MAIL,
            'call': AlertType.CALL,
            'missed_call': AlertType.MISSED_CALL
        }
        
        alert_type_str = data.get('type', 'message').lower()
        if alert_type_str not in alert_type_map:
            return jsonify({'error': 'Invalid notification type', 'code': 'INVALID_TYPE'}), 400
        
        alert_type = alert_type_map[alert_type_str]
        title = data.get('title', 'Notification')
        message = data.get('message', '')
        
        # Send notification
        await band.send_notification(alert_type, title, message)
        
        return jsonify({
            'success': True,
            'message': f'{alert_type.name} notification sent successfully',
            'notification': {
                'type': alert_type_str,
                'title': title,
                'message': message
            }
        })
    except Exception as e:
        logger.error(f"Send notification error: {e}")
        return jsonify({'error': str(e), 'code': 'NOTIFICATION_ERROR'}), 500


@app.route('/api/music', methods=['POST'])
@require_session
@async_route
async def set_music_info(session_id):
    """Set music information"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        data = request.get_json() or {}
        
        state = MusicState.PLAYING if data.get('state', 'playing').lower() == 'playing' else MusicState.PAUSED
        artist = data.get('artist', 'Unknown Artist')
        album = data.get('album', 'Unknown Album')
        track = data.get('track', 'Unknown Track')
        volume = int(data.get('volume', 75))
        position = int(data.get('position', 0))
        duration = int(data.get('duration', 240))
        
        await band.set_music_info(
            state=state,
            artist=artist,
            album=album,
            track=track,
            volume=volume,
            position=position,
            duration=duration
        )
        
        return jsonify({
            'success': True,
            'music_info': {
                'state': state.name.lower(),
                'artist': artist,
                'album': album,
                'track': track,
                'volume': volume,
                'position': position,
                'duration': duration
            }
        })
    except Exception as e:
        logger.error(f"Set music error: {e}")
        return jsonify({'error': str(e), 'code': 'MUSIC_ERROR'}), 500


@app.route('/api/steps', methods=['GET'])
@require_session
@require_auth
@async_route
async def get_steps_info(session_id):
    """Get steps information (requires auth)"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        steps = await band.get_steps()
        return jsonify({
            'success': True,
            'steps': {
                'steps': steps.steps,
                'meters': steps.meters,
                'calories': steps.calories,
                'fat_burned': steps.fat_burned
            }
        })
    except Exception as e:
        logger.error(f"Get steps error: {e}")
        return jsonify({'error': str(e), 'code': 'STEPS_ERROR'}), 500


@app.route('/api/heart_rate', methods=['GET'])
@require_session
@require_auth
@async_route
async def get_heart_rate(session_id):
    """Get single heart rate measurement (requires auth)"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        heart_rate = await band.get_heart_rate_one_time()
        return jsonify({
            'success': True,
            'heart_rate': {
                'bpm': heart_rate,
                'timestamp': datetime.now().isoformat()
            }
        })
    except Exception as e:
        logger.error(f"Get heart rate error: {e}")
        return jsonify({'error': str(e), 'code': 'HEART_RATE_ERROR'}), 500


@app.route('/api/alarm', methods=['POST'])
@require_session
@require_auth
@async_route
async def set_alarm(session_id):
    """Set alarm (requires auth)"""
    try:
        band = await dm.get_connection(session_id)
        if not band:
            return jsonify({'error': 'Not connected', 'code': 'NOT_CONNECTED'}), 400
        
        data = request.get_json() or {}
        
        hour = int(data.get('hour', 9))
        minute = int(data.get('minute', 0))
        days = tuple(int(d) for d in data.get('days', [1, 2, 3, 4, 5]) if 1 <= int(d) <= 7)
        enabled = bool(data.get('enabled', True))
        snooze = bool(data.get('snooze', True))
        alarm_id = int(data.get('alarm_id', 0))
        
        await band.set_alarm(
            hour=hour,
            minute=minute,
            days=days,
            enabled=enabled,
            snooze=snooze,
            alarm_id=alarm_id
        )
        
        day_names = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
        day_str = ", ".join(day_names[d-1] for d in days) if days else "No days"
        
        return jsonify({
            'success': True,
            'alarm': {
                'hour': hour,
                'minute': minute,
                'days': list(days),
                'days_text': day_str,
                'enabled': enabled,
                'snooze': snooze,
                'alarm_id': alarm_id,
                'time_text': f"{hour:02d}:{minute:02d}"
            }
        })
    except Exception as e:
        logger.error(f"Set alarm error: {e}")
        return jsonify({'error': str(e), 'code': 'ALARM_ERROR'}), 500


@app.route('/api/session/<session_id>', methods=['GET'])
def get_session_info(session_id):
    """Get session information"""
    config = dm.get_device_config(session_id)
    if not config:
        return jsonify({'error': 'Invalid session ID', 'code': 'INVALID_SESSION'}), 404
    
    is_connected = session_id in dm.connections and dm.connections[session_id].connected
    
    return jsonify({
        'success': True,
        'session': {
            'session_id': session_id,
            'mac_address': config['mac_address'],
            'auth_enabled': bool(config['auth_key']),
            'registered_at': config['registered_at'].isoformat(),
            'last_used': config['last_used'].isoformat(),
            'connected': is_connected
        }
    })


@app.route('/api/session/<session_id>', methods=['DELETE'])
@async_route
async def delete_session(session_id):
    """Delete session and disconnect device"""
    try:
        await dm.disconnect(session_id)
        dm.remove_device(session_id)
        return jsonify({'success': True, 'message': 'Session deleted successfully'})
    except Exception as e:
        logger.error(f"Delete session error: {e}")
        return jsonify({'error': str(e), 'code': 'DELETE_SESSION_ERROR'}), 500


# Error handlers
@app.errorhandler(404)
def not_found(error):
    return jsonify({'error': 'Endpoint not found', 'code': 'NOT_FOUND'}), 404


@app.errorhandler(405)
def method_not_allowed(error):
    return jsonify({'error': 'Method not allowed', 'code': 'METHOD_NOT_ALLOWED'}), 405


@app.errorhandler(500)
def internal_error(error):
    return jsonify({'error': 'Internal server error', 'code': 'INTERNAL_ERROR'}), 500


# API documentation endpoint
@app.route('/api/docs', methods=['GET'])
def api_docs():
    """API documentation"""
    docs = {
        'title': 'Mi Band 4 API',
        'version': '1.0.0',
        'description': 'REST API for Mi Band 4 device control',
        'endpoints': {
            'Device Discovery': {
                'POST /api/discover': 'Discover nearby Mi Band devices'
            },
            'Session Management': {
                'POST /api/register': 'Register device and get session ID',
                'GET /api/session/{session_id}': 'Get session information',
                'DELETE /api/session/{session_id}': 'Delete session and disconnect'
            },
            'Connection': {
                'POST /api/connect': 'Connect to registered device',
                'POST /api/disconnect': 'Disconnect from device'
            },
            'Device Information': {
                'GET /api/device/info': 'Get device information',
                'GET /api/device/battery': 'Get battery information',
                'GET /api/device/time': 'Get device time',
                'POST /api/device/time': 'Set device time (requires auth)'
            },
            'Notifications': {
                'POST /api/notification': 'Send notification (types: message, mail, call, missed_call)'
            },
            'Music Control': {
                'POST /api/music': 'Set music information'
            },
            'Health & Fitness (requires auth)': {
                'GET /api/steps': 'Get steps information',
                'GET /api/heart_rate': 'Get single heart rate measurement'
            },
            'Alarms (requires auth)': {
                'POST /api/alarm': 'Set alarm'
            },
            'Utility': {
                'GET /api/health': 'API health check',
                'GET /api/docs': 'This documentation'
            }
        },
        'authentication': {
            'description': 'Some endpoints require authentication (auth key)',
            'session_header': 'X-Session-ID',
            'required_for': ['steps', 'heart_rate', 'alarms', 'set_time']
        },
        'session_management': {
            'description': 'Register device first, then use session ID for all operations',
            'steps': [
                '1. POST /api/register with mac_address and optional auth_key',
                '2. Use returned session_id in X-Session-ID header for all requests',
                '3. POST /api/connect to establish connection',
                '4. Use other endpoints as needed'
            ]
        }
    }
    
    return jsonify(docs)


if __name__ == '__main__':
    print("Starting Mi Band 4 Flask API...")
    print("API Documentation: http://localhost:5000/api/docs")
    print("Health Check: http://localhost:5000/api/health")
    
    # Load default configuration
    default_mac = os.getenv('MAC_ADDRESS')
    default_auth = os.getenv('AUTH_KEY')
    
    if default_mac:
        print(f"Default device configured: {default_mac}")
        if default_auth:
            print("Default auth key configured")
    
    app.run(
        host=os.getenv('FLASK_HOST', '0.0.0.0'),
        port=int(os.getenv('FLASK_PORT', 5000)),
        debug=os.getenv('FLASK_DEBUG', 'False').lower() == 'true'
    )