import gradio as gr
import os
import subprocess
import tempfile
import shutil
from zipfile import ZipFile
import logging
import json
import threading
import psutil
from flask import Flask, request, jsonify

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

# Initialize Flask app
app = Flask(__name__)

connected_cpus = {}

# Endpoint to donate CPU resources
@app.route('/donate_cpu', methods=['POST'])
def donate_cpu_handler():
    data = request.get_json()
    host = data['host']
    cpu_count = data['cpu_count']
    connected_cpus[host] = {"cpu_count": cpu_count, "usage": 0.0}
    logger.info(f"CPU donated by {host} with {cpu_count} CPUs.")
    return jsonify({"status": "success", "message": f"CPU donated by {host}"})

# Endpoint to update CPU usage
@app.route('/update_cpu_usage', methods=['POST'])
def update_cpu_usage_handler():
    data = request.get_json()
    host = data['host']
    usage = data['usage']
    if host in connected_cpus:
        connected_cpus[host]['usage'] = usage
        logger.info(f"Updated CPU usage for {host}: {usage}%")
    return jsonify({"status": "success"})

# Function to run the provided Python script using MPI
def run_script(script_name, folder_path):
    output_log = tempfile.TemporaryFile(mode='w+t')
    try:
        # Collect all available CPUs
        total_cpus = sum(cpu['cpu_count'] for cpu in connected_cpus.values())
        
        # Run the script using MPI
        result = subprocess.run(['mpiexec', '-n', str(total_cpus), 'python', script_name], cwd=folder_path, stdout=output_log, stderr=subprocess.STDOUT)
        output_log.seek(0)
        log_output = output_log.read()
    except Exception as e:
        log_output = str(e)
    finally:
        output_log.close()
    
    return log_output

# Function to handle file uploads and script execution
def handle_upload(folder, script_name):
    # Create a temporary directory to store uploaded files
    temp_dir = tempfile.mkdtemp()
    
    # Save the uploaded folder contents to the temporary directory
    folder_path = os.path.join(temp_dir, 'uploaded_folder')
    os.makedirs(folder_path, exist_ok=True)
    for file_name, file_obj in folder.items():
        with open(os.path.join(folder_path, file_name), 'wb') as f:
            f.write(file_obj.read())

    # Run the script
    log_output = run_script(script_name, folder_path)
    
    # Create a zip file of the entire folder (including any new files created by the script)
    zip_path = os.path.join(temp_dir, 'output_folder.zip')
    with ZipFile(zip_path, 'w') as zipf:
        for root, _, files in os.walk(folder_path):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
    
    return log_output, zip_path

# Function to get connected CPUs information
def get_cpu_info():
    info = []
    for host, data in connected_cpus.items():
        info.append(f"{host}: {data['cpu_count']} CPUs, {data['usage']}% usage")
    return "\n".join(info)

# Gradio interface
def gradio_interface():
    interface_inputs = [
        gr.File(label="Upload Folder", file_count="multiple", file_types=['file']),
        gr.Textbox(label="Python Script Name")
    ]
    interface_outputs = [
        gr.Textbox(label="Log Output", interactive=False),
        gr.File(label="Download Output Folder"),
        gr.Textbox(label="Connected CPUs Info", interactive=False)
    ]
    
    iface = gr.Interface(
        fn=handle_upload,
        inputs=interface_inputs,
        outputs=interface_outputs,
        live=True,
        theme="light"  # Specify a theme that works, "light" is an example
    )
    
    iface.launch()

# Launch the Gradio interface using Flask's run method
if __name__ == "__main__":
    gradio_interface()
    # Uncomment the line below if using Flask's run method for local testing
    # app.run(host='0.0.0.0', port=7860)