from flask import Flask, request, jsonify
from STOUT import translate_forward, translate_reverse
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
import time

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)

# Thread pool for parallel processing
executor = ThreadPoolExecutor(max_workers=4)

@app.route('/')
def home():
    return jsonify({
        "message": "STOUT V2 API - SMILES to IUPAC Translator",
        "endpoints": {
            "/smiles_to_iupac": "GET/POST - Convert SMILES to IUPAC name",
            "/iupac_to_smiles": "GET/POST - Convert IUPAC name to SMILES",
            "/batch/smiles_to_iupac": "POST - Convert multiple SMILES to IUPAC names",
            "/batch/iupac_to_smiles": "POST - Convert multiple IUPAC names to SMILES",
            "/health": "GET - Health check"
        }
    })

@app.route('/health')
def health():
    return jsonify({"status": "healthy"})

def process_single_smiles(smiles):
    """Process a single SMILES string"""
    try:
        iupac_name = translate_forward(smiles)
        return {"smiles": smiles, "iupac": iupac_name, "success": True}
    except Exception as e:
        return {"smiles": smiles, "error": str(e), "success": False}

def process_single_iupac(iupac):
    """Process a single IUPAC name"""
    try:
        smiles = translate_reverse(iupac)
        return {"iupac": iupac, "smiles": smiles, "success": True}
    except Exception as e:
        return {"iupac": iupac, "error": str(e), "success": False}

@app.route('/smiles_to_iupac', methods=['GET', 'POST'])
def smiles_to_iupac():
    try:
        if request.method == 'GET':
            smiles = request.args.get('smiles')
        else:  # POST
            data = request.get_json()
            smiles = data.get('smiles') if data else None
        
        if not smiles:
            return jsonify({"error": "Missing 'smiles' parameter"}), 400
        
        result = process_single_smiles(smiles)
        if result['success']:
            return jsonify(result)
        else:
            return jsonify(result), 500
            
    except Exception as e:
        app.logger.error(f"Error translating SMILES: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/batch/smiles_to_iupac', methods=['POST'])
def batch_smiles_to_iupac():
    """Process multiple SMILES strings in parallel"""
    try:
        start_time = time.time()
        data = request.get_json()
        
        if not data or 'smiles_list' not in data:
            return jsonify({"error": "Missing 'smiles_list' in request body"}), 400
        
        smiles_list = data['smiles_list']
        if not isinstance(smiles_list, list):
            return jsonify({"error": "'smiles_list' must be an array"}), 400
        
        # Limit batch size to prevent abuse
        if len(smiles_list) > 100:
            return jsonify({"error": "Maximum batch size is 100"}), 400
        
        # Process in parallel using thread pool
        futures = [executor.submit(process_single_smiles, smiles) for smiles in smiles_list]
        results = []
        
        for future in as_completed(futures):
            results.append(future.result())
        
        # Sort results to maintain input order
        results_dict = {r['smiles']: r for r in results}
        ordered_results = [results_dict.get(smiles, {"smiles": smiles, "error": "Not processed", "success": False}) 
                          for smiles in smiles_list]
        
        processing_time = time.time() - start_time
        
        return jsonify({
            "results": ordered_results,
            "total": len(smiles_list),
            "successful": sum(1 for r in ordered_results if r.get('success', False)),
            "failed": sum(1 for r in ordered_results if not r.get('success', False)),
            "processing_time_seconds": round(processing_time, 3)
        })
        
    except Exception as e:
        app.logger.error(f"Error in batch processing: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/batch/iupac_to_smiles', methods=['POST'])
def batch_iupac_to_smiles():
    """Process multiple IUPAC names in parallel"""
    try:
        start_time = time.time()
        data = request.get_json()
        
        if not data or 'iupac_list' not in data:
            return jsonify({"error": "Missing 'iupac_list' in request body"}), 400
        
        iupac_list = data['iupac_list']
        if not isinstance(iupac_list, list):
            return jsonify({"error": "'iupac_list' must be an array"}), 400
        
        # Limit batch size to prevent abuse
        if len(iupac_list) > 100:
            return jsonify({"error": "Maximum batch size is 100"}), 400
        
        # Process in parallel using thread pool
        futures = [executor.submit(process_single_iupac, iupac) for iupac in iupac_list]
        results = []
        
        for future in as_completed(futures):
            results.append(future.result())
        
        # Sort results to maintain input order
        results_dict = {r['iupac']: r for r in results}
        ordered_results = [results_dict.get(iupac, {"iupac": iupac, "error": "Not processed", "success": False}) 
                          for iupac in iupac_list]
        
        processing_time = time.time() - start_time
        
        return jsonify({
            "results": ordered_results,
            "total": len(iupac_list),
            "successful": sum(1 for r in ordered_results if r.get('success', False)),
            "failed": sum(1 for r in ordered_results if not r.get('success', False)),
            "processing_time_seconds": round(processing_time, 3)
        })
        
    except Exception as e:
        app.logger.error(f"Error in batch processing: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/iupac_to_smiles', methods=['GET', 'POST'])
def iupac_to_smiles():
    try:
        if request.method == 'GET':
            iupac = request.args.get('iupac')
        else:  # POST
            data = request.get_json()
            iupac = data.get('iupac') if data else None
        
        if not iupac:
            return jsonify({"error": "Missing 'iupac' parameter"}), 400
        
        result = process_single_iupac(iupac)
        if result['success']:
            return jsonify(result)
        else:
            return jsonify(result), 500
            
    except Exception as e:
        app.logger.error(f"Error translating IUPAC: {str(e)}")
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)