#!/usr/bin/env python3
"""
Flask API for Memory Module Detection
This API processes motherboard images and detects memory modules using YOLOv8.
"""

import os
import io
import base64
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
from flask_restx import Api, Resource, fields, reqparse
from PIL import Image
import numpy as np
from werkzeug.utils import secure_filename
from werkzeug.datastructures import FileStorage
import tempfile
import logging
from datetime import datetime
from inference_utils import MemoryModuleDetector

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

# Initialize Flask app
app = Flask(__name__)
CORS(app)

# Initialize Flask-RESTX API with custom configuration
api = Api(
    app,
    version='1.0',
    title='Memory Module Detection API',
    description='AI-powered memory module detection system for motherboard images using YOLOv8',
    doc='/docs/',
    prefix='/api/v1'
)

# Create namespaces
ns_health = api.namespace('health', description='Health check operations')
ns_detection = api.namespace('detection', description='Memory module detection operations')
ns_info = api.namespace('info', description='API information')

# Define API models for documentation
detection_result = api.model('DetectionResult', {
    'success': fields.Boolean(required=True, description='Whether detection was successful'),
    'detections': fields.List(fields.Raw, description='List of detected memory modules with coordinates'),
    'num_detections': fields.Integer(description='Number of memory modules detected'),
    'annotated_image': fields.String(description='Base64 encoded annotated image'),
    'confidence_threshold': fields.Float(description='Confidence threshold used for detection'),
    'test_image_path': fields.String(description='Path to the test image (for hardcoded tests)')
})

error_response = api.model('ErrorResponse', {
    'error': fields.String(required=True, description='Error message'),
    'success': fields.Boolean(required=True, description='Always false for errors')
})

health_response = api.model('HealthResponse', {
    'status': fields.String(required=True, description='Health status'),
    'model_loaded': fields.Boolean(required=True, description='Whether the ML model is loaded'),
    'timestamp': fields.String(required=True, description='Current timestamp')
})

api_info_response = api.model('ApiInfoResponse', {
    'name': fields.String(required=True, description='API name'),
    'version': fields.String(required=True, description='API version'),
    'description': fields.String(required=True, description='API description'),
    'model_info': fields.Raw(description='Information about the ML model'),
    'endpoints': fields.List(fields.String, description='Available endpoints')
})

# Configuration
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}

# Create upload folder if it doesn't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Initialize detector
MODEL_PATH = 'runs/detect/memory_module_detection/weights/best.pt'
detector = MemoryModuleDetector(MODEL_PATH)

# Hardcoded test image path
HARDCODED_IMAGE_PATH = 'training/memory/out1.png'

def allowed_file(filename):
    """Check if file extension is allowed."""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def image_to_base64(image):
    """Convert PIL Image to base64 string."""
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    img_str = base64.b64encode(buffer.getvalue()).decode()
    return img_str

def base64_to_image(base64_string):
    """Convert base64 string to PIL Image."""
    img_data = base64.b64decode(base64_string)
    image = Image.open(io.BytesIO(img_data))
    return image

@app.route('/', methods=['GET'])
def home():
    """Home endpoint - serve frontend or API information based on Accept header."""
    # Check if request is from a browser (wants HTML)
    if 'text/html' in request.headers.get('Accept', ''):
        return render_template('index.html')

    # Otherwise return JSON API information
    return jsonify({
        'message': 'Memory Module Detection API',
        'version': '1.0.0',
        'endpoints': {
            '/': 'GET - Frontend interface or API information',
            '/api': 'GET - API information (JSON)',
            '/detect': 'POST - Upload image for memory module detection',
            '/detect/hardcoded': 'GET - Process hardcoded test image',
            '/detect/base64': 'POST - Process base64 encoded image',
            '/health': 'GET - Health check'
        },
        'model_loaded': detector.model is not None,
        'supported_formats': list(ALLOWED_EXTENSIONS)
    })

@app.route('/api', methods=['GET'])
def api_info():
    """API information endpoint (always returns JSON)."""
    return jsonify({
        'message': 'Memory Module Detection API',
        'version': '1.0.0',
        'endpoints': {
            '/': 'GET - Frontend interface or API information',
            '/api': 'GET - API information (JSON)',
            '/detect': 'POST - Upload image for memory module detection',
            '/detect/hardcoded': 'GET - Process hardcoded test image',
            '/detect/base64': 'POST - Process base64 encoded image',
            '/health': 'GET - Health check'
        },
        'model_loaded': detector.model is not None,
        'supported_formats': list(ALLOWED_EXTENSIONS)
    })



@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({
        'status': 'healthy',
        'model_loaded': detector.model is not None,
        'model_path': MODEL_PATH
    })

@app.route('/detect', methods=['POST'])
def detect_memory_modules():
    """
    Detect memory modules in uploaded image.
    
    Expected input:
    - File upload with key 'image'
    - Optional: confidence threshold as form data
    
    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500
        
        # Check if file is present
        if 'image' not in request.files:
            return jsonify({
                'error': 'No image file provided',
                'success': False
            }), 400
        
        file = request.files['image']
        
        # Check if file is selected
        if file.filename == '':
            return jsonify({
                'error': 'No file selected',
                'success': False
            }), 400
        
        # Check file extension
        if not allowed_file(file.filename):
            return jsonify({
                'error': f'File type not allowed. Supported formats: {ALLOWED_EXTENSIONS}',
                'success': False
            }), 400
        
        # Get confidence threshold from form data (default 80%)
        conf_threshold = float(request.form.get('confidence', 0.8))
        
        # Save uploaded file temporarily
        filename = secure_filename(file.filename)
        temp_path = os.path.join(UPLOAD_FOLDER, filename)
        file.save(temp_path)
        
        try:
            # Run detection
            detections, annotated_image = detector.detect(
                temp_path, 
                conf_threshold=conf_threshold
            )
            
            # Convert annotated image to base64
            annotated_base64 = image_to_base64(annotated_image)
            
            # Prepare response
            response_data = {
                'success': True,
                'detections': detections,
                'num_detections': len(detections),
                'annotated_image': annotated_base64,
                'confidence_threshold': conf_threshold,
                'original_filename': filename
            }
            
            logger.info(f"Processed {filename}: found {len(detections)} memory modules")
            
            return jsonify(response_data)
        
        finally:
            # Clean up temporary file
            if os.path.exists(temp_path):
                os.remove(temp_path)
    
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        return jsonify({
            'error': f'Error processing image: {str(e)}',
            'success': False
        }), 500

@app.route('/detect/hardcoded', methods=['GET'])
def detect_hardcoded_image():
    """
    Process hardcoded test image for memory module detection.

    Optional query parameters:
    - confidence: confidence threshold (default: 0.8)

    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500

        # Check if hardcoded image exists
        if not os.path.exists(HARDCODED_IMAGE_PATH):
            return jsonify({
                'error': f'Hardcoded test image not found at {HARDCODED_IMAGE_PATH}',
                'success': False
            }), 404

        # Get confidence threshold from query parameters (default 80%)
        conf_threshold = float(request.args.get('confidence', 0.8))

        # Run detection
        detections, annotated_image = detector.detect(
            HARDCODED_IMAGE_PATH,
            conf_threshold=conf_threshold
        )

        # Convert annotated image to base64
        annotated_base64 = image_to_base64(annotated_image)

        # Prepare response
        response_data = {
            'success': True,
            'detections': detections,
            'num_detections': len(detections),
            'annotated_image': annotated_base64,
            'confidence_threshold': conf_threshold,
            'test_image_path': HARDCODED_IMAGE_PATH
        }

        logger.info(f"Processed hardcoded image: found {len(detections)} memory modules")

        return jsonify(response_data)
    
    except Exception as e:
        logger.error(f"Error processing hardcoded image: {str(e)}")
        return jsonify({
            'error': f'Error processing hardcoded image: {str(e)}',
            'success': False
        }), 500



@app.route('/detect/base64', methods=['POST'])
def detect_base64_image():
    """
    Detect memory modules in base64 encoded image.
    
    Expected JSON input:
    {
        "image": "base64_encoded_image_string",
        "confidence": 0.5  // optional
    }
    
    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500
        
        # Get JSON data
        data = request.get_json()
        
        if not data or 'image' not in data:
            return jsonify({
                'error': 'No base64 image data provided',
                'success': False
            }), 400
        
        # Get confidence threshold (default 80%)
        conf_threshold = float(data.get('confidence', 0.8))
        
        # Decode base64 image
        try:
            image = base64_to_image(data['image'])
        except Exception as e:
            return jsonify({
                'error': f'Invalid base64 image data: {str(e)}',
                'success': False
            }), 400
        
        # Run detection
        detections, annotated_image = detector.detect_from_array(
            np.array(image), 
            conf_threshold=conf_threshold
        )
        
        # Convert annotated image to base64
        annotated_base64 = image_to_base64(annotated_image)
        
        # Prepare response
        response_data = {
            'success': True,
            'detections': detections,
            'num_detections': len(detections),
            'annotated_image': annotated_base64,
            'confidence_threshold': conf_threshold
        }
        
        logger.info(f"Processed base64 image: found {len(detections)} memory modules")
        
        return jsonify(response_data)
    
    except Exception as e:
        logger.error(f"Error processing base64 image: {str(e)}")
        return jsonify({
            'error': f'Error processing base64 image: {str(e)}',
            'success': False
        }), 500

@app.errorhandler(413)
def too_large(e):
    """Handle file too large error."""
    return jsonify({
        'error': 'File too large. Maximum size is 16MB.',
        'success': False
    }), 413

@app.errorhandler(404)
def not_found(e):
    """Handle 404 errors."""
    return jsonify({
        'error': 'Endpoint not found',
        'success': False
    }), 404

@app.errorhandler(500)
def internal_error(e):
    """Handle internal server errors."""
    return jsonify({
        'error': 'Internal server error',
        'success': False
    }), 500


# ============================================================================
# SWAGGER API RESOURCES
# ============================================================================

@ns_health.route('')
class HealthCheck(Resource):
    @ns_health.doc('health_check')
    @ns_health.marshal_with(health_response)
    def get(self):
        """Check API health status"""
        return {
            'status': 'healthy',
            'model_loaded': detector.model is not None,
            'timestamp': datetime.now().isoformat()
        }

@ns_info.route('')
class ApiInfo(Resource):
    @ns_info.doc('api_info')
    @ns_info.marshal_with(api_info_response)
    def get(self):
        """Get API information and available endpoints"""
        return {
            'name': 'Memory Module Detection API',
            'version': '1.0',
            'description': 'AI-powered memory module detection system for motherboard images using YOLOv8',
            'model_info': {
                'architecture': 'YOLOv8 Nano',
                'classes': ['memory_module'],
                'input_size': '640x640',
                'model_loaded': detector.model is not None
            },
            'endpoints': [
                '/api/v1/health',
                '/api/v1/info',
                '/api/v1/detection/upload',
                '/api/v1/detection/hardcoded',
                '/api/v1/detection/base64'
            ]
        }

# File upload parser
upload_parser = reqparse.RequestParser()
upload_parser.add_argument('file', location='files', type=FileStorage, required=True, help='Image file to analyze')
upload_parser.add_argument('confidence', type=float, default=0.8, help='Confidence threshold (0.0-1.0)')

@ns_detection.route('/upload')
class DetectionUpload(Resource):
    @ns_detection.doc('upload_detection')
    @ns_detection.expect(upload_parser)
    @ns_detection.marshal_with(detection_result, code=200)
    @ns_detection.marshal_with(error_response, code=400)
    @ns_detection.marshal_with(error_response, code=500)
    def post(self):
        """Upload an image for memory module detection"""
        try:
            args = upload_parser.parse_args()
            file = args['file']
            confidence = args.get('confidence', 0.8)

            if not file or file.filename == '':
                return {'error': 'No file provided', 'success': False}, 400

            if not allowed_file(file.filename):
                return {'error': 'Invalid file type. Allowed: PNG, JPG, JPEG, GIF, BMP', 'success': False}, 400

            # Save uploaded file temporarily
            filename = secure_filename(file.filename)
            temp_path = os.path.join(UPLOAD_FOLDER, filename)
            file.save(temp_path)

            # Run detection
            detections, annotated_image = detector.detect(temp_path, conf_threshold=confidence)

            # Convert annotated image to base64
            annotated_base64 = image_to_base64(annotated_image)

            return {
                'success': True,
                'detections': detections,
                'num_detections': len(detections),
                'annotated_image': annotated_base64,
                'confidence_threshold': confidence
            }

        except Exception as e:
            return {'error': f'Error processing image: {str(e)}', 'success': False}, 500

# Hardcoded test parser
hardcoded_parser = reqparse.RequestParser()
hardcoded_parser.add_argument('confidence', type=float, default=0.8, help='Confidence threshold (0.0-1.0)')

@ns_detection.route('/hardcoded')
class DetectionHardcoded(Resource):
    @ns_detection.doc('hardcoded_detection')
    @ns_detection.expect(hardcoded_parser)
    @ns_detection.marshal_with(detection_result, code=200)
    @ns_detection.marshal_with(error_response, code=404)
    @ns_detection.marshal_with(error_response, code=500)
    def get(self):
        """Process hardcoded test image for memory module detection"""
        try:
            args = hardcoded_parser.parse_args()
            confidence = args.get('confidence', 0.8)

            if detector.model is None:
                return {'error': 'Model not loaded. Please train the model first.', 'success': False}, 500

            if not os.path.exists(HARDCODED_IMAGE_PATH):
                return {'error': f'Hardcoded test image not found at {HARDCODED_IMAGE_PATH}', 'success': False}, 404

            # Run detection
            detections, annotated_image = detector.detect(HARDCODED_IMAGE_PATH, conf_threshold=confidence)

            # Convert annotated image to base64
            annotated_base64 = image_to_base64(annotated_image)

            return {
                'success': True,
                'detections': detections,
                'num_detections': len(detections),
                'annotated_image': annotated_base64,
                'confidence_threshold': confidence,
                'test_image_path': HARDCODED_IMAGE_PATH
            }

        except Exception as e:
            return {'error': f'Error processing hardcoded image: {str(e)}', 'success': False}, 500

# Base64 detection parser
base64_parser = reqparse.RequestParser()
base64_parser.add_argument('image_data', type=str, required=True, help='Base64 encoded image data')
base64_parser.add_argument('confidence', type=float, default=0.8, help='Confidence threshold (0.0-1.0)')

@ns_detection.route('/base64')
class DetectionBase64(Resource):
    @ns_detection.doc('base64_detection')
    @ns_detection.expect(base64_parser)
    @ns_detection.marshal_with(detection_result, code=200)
    @ns_detection.marshal_with(error_response, code=400)
    @ns_detection.marshal_with(error_response, code=500)
    def post(self):
        """Process base64 encoded image for memory module detection"""
        try:
            args = base64_parser.parse_args()
            image_data = args['image_data']
            confidence = args.get('confidence', 0.8)

            if detector.model is None:
                return {'error': 'Model not loaded. Please train the model first.', 'success': False}, 500

            # Decode base64 image
            try:
                image_bytes = base64.b64decode(image_data)
                image = Image.open(io.BytesIO(image_bytes))
            except Exception as e:
                return {'error': f'Invalid base64 image data: {str(e)}', 'success': False}, 400

            # Save temporarily for processing
            temp_path = os.path.join(UPLOAD_FOLDER, 'temp_base64.png')
            image.save(temp_path)

            # Run detection
            detections, annotated_image = detector.detect(temp_path, conf_threshold=confidence)

            # Convert annotated image to base64
            annotated_base64 = image_to_base64(annotated_image)

            return {
                'success': True,
                'detections': detections,
                'num_detections': len(detections),
                'annotated_image': annotated_base64,
                'confidence_threshold': confidence
            }

        except Exception as e:
            return {'error': f'Error processing base64 image: {str(e)}', 'success': False}, 500


if __name__ == '__main__':
    # Check if model exists
    if not os.path.exists(MODEL_PATH):
        print(f"Warning: Model not found at {MODEL_PATH}")
        print("Please train the model first using: python3 train.py")
        print("The API will still start but detection endpoints will return errors.")
    
    # Start the Flask app
    print("🚀 Starting Memory Module Detection API...")
    print(f"📊 Model path: {MODEL_PATH}")
    print(f"🤖 Model loaded: {detector.model is not None}")
    print(f"🖼️  Hardcoded test image: {HARDCODED_IMAGE_PATH}")
    print("")
    print("🌐 Web Interface: http://localhost:5002")
    print("📚 API Documentation: http://localhost:5002/docs/")
    print("🔧 Swagger UI (Professional): Run 'python3 swagger_app.py' for port 5003")
    print("")

    app.run(host='0.0.0.0', port=5002, debug=True)