import os
import sys
import pty
import select
import signal
import subprocess
import threading
import time
from flask import Flask, jsonify, request, render_template, send_from_directory, Response

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

# Global process variables
current_process = None
master_fd = None
log_buffer = ""
running_script = None
log_lock = threading.Lock()

# Define script workflow categories
SCRIPT_CATEGORIES = {
    "workspace": {
        "title": "Workspace & Setup",
        "icon": "folder",
        "scripts": [
            {"file": "1_clear_workspace.sh", "name": "Clear Workspace", "desc": "Deletes the workspace directory and all intermediate files to start clean."}
        ]
    },
    "src_video": {
        "title": "Step 1: Source Video (data_src)",
        "icon": "video",
        "scripts": [
            {"file": "2_extract_image_from_data_src.sh", "name": "Extract Images", "desc": "Extracts video frames from workspace/data_src.mp4 to PNG files."},
            {"file": "4_data_src_extract_faces_S3FD.sh", "name": "Extract Faces (S3FD)", "desc": "Uses S3FD detector to extract aligned face crops from source frames."},
            {"file": "4_data_src_extract_faces_MANUAL.sh", "name": "Extract Faces (Manual)", "desc": "Interactively extract or fix missing face detections manually."},
            {"file": "4.2_data_src_sort.sh", "name": "Sort Source Faces", "desc": "Sorts source faces by similarity, luminance, blur, etc., to prune trash."}
        ]
    },
    "dst_video": {
        "title": "Step 2: Destination Video (data_dst)",
        "icon": "video",
        "scripts": [
            {"file": "3_extract_image_from_data_dst.sh", "name": "Extract Images", "desc": "Extracts video frames from workspace/data_dst.mp4 to PNG files."},
            {"file": "3.1_denoise_data_dst_images.sh", "name": "Denoise Images", "desc": "Denoises the extracted destination frames for cleaner final merge."},
            {"file": "5_data_dst_extract_faces_S3FD.sh", "name": "Extract Faces (S3FD)", "desc": "Uses S3FD detector to extract aligned faces from destination frames."},
            {"file": "5_data_dst_extract_faces_S3FD_+_manual_fix.sh", "name": "Extract Faces (S3FD + Manual)", "desc": "Extract faces using S3FD with manual correction fallback."},
            {"file": "5_data_dst_extract_faces_MANUAL.sh", "name": "Extract Faces (Manual)", "desc": "Interactively extract or fix destination faces manually."}
        ]
    },
    "xseg": {
        "title": "Step 3: XSeg Masking (Optional)",
        "icon": "scissors",
        "scripts": [
            {"file": "5_XSeg_train.sh", "name": "Train XSeg Mask", "desc": "Trains a custom XSeg masking model on labeled faces."},
            {"file": "5_XSeg_data_src_mask_edit.sh", "name": "Edit Src Masks", "desc": "Label/edit XSeg masks on your source face images."},
            {"file": "5_XSeg_data_src_mask_apply.sh", "name": "Apply Src Masks", "desc": "Applies a trained XSeg model to mask source face images."},
            {"file": "5_XSeg_data_dst_mask_edit.sh", "name": "Edit Dst Masks", "desc": "Label/edit XSeg masks on your destination face images."},
            {"file": "5_XSeg_data_dst_mask_apply.sh", "name": "Apply Dst Masks", "desc": "Applies a trained XSeg model to mask destination face images."}
        ]
    },
    "training": {
        "title": "Step 4: Model Training",
        "icon": "cpu",
        "scripts": [
            {"file": "6_train_Quick96_no_preview.sh", "name": "Train Quick96 (Headless)", "desc": "Trains Quick96 model in the background (no local window)."},
            {"file": "6_train_Quick96.sh", "name": "Train Quick96 (Windowed)", "desc": "Trains Quick96 model. Requires local X11 visual window."},
            {"file": "6_train_SAEHD_no_preview.sh", "name": "Train SAEHD (Headless)", "desc": "Trains high-quality SAEHD model in the background."},
            {"file": "6_train_SAEHD.sh", "name": "Train SAEHD (Windowed)", "desc": "Trains SAEHD model. Requires local X11 visual window."}
        ]
    },
    "merging": {
        "title": "Step 5: Face Merging",
        "icon": "merge",
        "scripts": [
            {"file": "7_merge_Quick96.sh", "name": "Merge Quick96", "desc": "Merges Quick96 face swaps onto destination frames."},
            {"file": "7_merge_SAEHD.sh", "name": "Merge SAEHD", "desc": "Merges SAEHD face swaps onto destination frames."}
        ]
    },
    "export": {
        "title": "Step 6: Export & Render",
        "icon": "download",
        "scripts": [
            {"file": "8_merged_to_mp4.sh", "name": "Export to MP4", "desc": "Combines merged frames into an MP4 file (h264)."},
            {"file": "8_merged_to_mp4_lossless.sh", "name": "Export to MP4 (Lossless)", "desc": "Combines merged frames into a lossless MP4 file."},
            {"file": "8_merged_to_avi.sh", "name": "Export to AVI", "desc": "Combines merged frames into an uncompressed AVI file."}
        ]
    }
}

# Add pre-trained download links to categories
SCRIPT_CATEGORIES["pretrain"] = {
    "title": "Download Pre-trained Weights",
    "icon": "cloud-download",
    "scripts": [
        {"file": "4.1_download_Quick96.sh", "name": "Download Quick96 pretrain", "desc": "Downloads Quick96 pretraining zip file."},
        {"file": "4.1_download_CelebA.sh", "name": "Download CelebA pretrain", "desc": "Downloads CelebA pretraining zip file."},
        {"file": "4.1_download_FFHQ.sh", "name": "Download FFHQ pretrain", "desc": "Downloads FFHQ pretraining zip file."}
    ]
}

def log_reader(fd, proc):
    global log_buffer, current_process, master_fd, running_script
    while True:
        try:
            r, w, x = select.select([fd], [], [], 0.1)
            if fd in r:
                data = os.read(fd, 4096)
                if not data:
                    break
                decoded_chunk = data.decode('utf-8', errors='replace')
                with log_lock:
                    log_buffer += decoded_chunk
        except (OSError, ValueError):
            break
        except Exception:
            break
        
        # Check if process died
        if proc.poll() is not None:
            # Final drain of buffer
            try:
                r, w, x = select.select([fd], [], [], 0.1)
                if fd in r:
                    data = os.read(fd, 4096)
                    if data:
                        decoded_chunk = data.decode('utf-8', errors='replace')
                        with log_lock:
                            log_buffer += decoded_chunk
            except Exception:
                pass
            break
            
    try:
        os.close(fd)
    except Exception:
        pass
        
    with log_lock:
        if current_process == proc:
            current_process = None
            master_fd = None
            running_script = None

def get_gpu_info():
    try:
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=gpu_name,utilization.gpu,utilization.memory,memory.total,memory.used,temperature.gpu', '--format=csv,noheader,nounits'],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=2
        )
        if result.returncode == 0:
            lines = result.stdout.strip().split('\n')
            gpus = []
            for line in lines:
                parts = [p.strip() for p in line.split(',')]
                if len(parts) >= 6:
                    gpus.append({
                        'name': parts[0],
                        'gpu_util': parts[1] + '%',
                        'mem_util': parts[2] + '%',
                        'mem_total': parts[3] + ' MB',
                        'mem_used': parts[4] + ' MB',
                        'temp': parts[5] + '°C'
                    })
            return gpus
    except Exception:
        pass
    return None

def get_system_stats():
    # Basic system utilization (cpu / memory)
    # Reading from /proc/stat
    cpu_util = "0%"
    mem_util = "0%"
    try:
        # CPU calculation
        with open('/proc/stat', 'r') as f:
            fields = [float(column) for column in f.readline().strip().split()[1:]]
        idle, total = fields[3], sum(fields)
        time.sleep(0.1)
        with open('/proc/stat', 'r') as f:
            fields2 = [float(column) for column in f.readline().strip().split()[1:]]
        idle2, total2 = fields2[3], sum(fields2)
        
        diff_idle = idle2 - idle
        diff_total = total2 - total
        if diff_total > 0:
            cpu_util = f"{int((1.0 - diff_idle / diff_total) * 100)}%"
            
        # Memory calculation
        with open('/proc/meminfo', 'r') as f:
            lines = f.readlines()
        mem_total = 1.0
        mem_free = 0.0
        mem_cached = 0.0
        mem_buffers = 0.0
        for line in lines:
            if line.startswith('MemTotal:'):
                mem_total = float(line.split()[1])
            elif line.startswith('MemFree:'):
                mem_free = float(line.split()[1])
            elif line.startswith('Cached:'):
                mem_cached = float(line.split()[1])
            elif line.startswith('Buffers:'):
                mem_buffers = float(line.split()[1])
        mem_used = mem_total - mem_free - mem_cached - mem_buffers
        mem_util = f"{int((mem_used / mem_total) * 100)}%"
    except Exception:
        pass
        
    return {
        'cpu': cpu_util,
        'mem': mem_util
    }

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/status')
def status():
    global current_process, running_script
    is_running = current_process is not None and current_process.poll() is None
    
    # Get workspace stats
    workspace_dir = os.path.join(os.path.dirname(__file__), 'workspace')
    stats = {
        'src_frames': 0,
        'src_faces': 0,
        'dst_frames': 0,
        'dst_faces': 0,
        'models': []
    }
    
    if os.path.exists(workspace_dir):
        def count_files(path, ext=None):
            if not os.path.exists(path):
                return 0
            count = 0
            for entry in os.scandir(path):
                if entry.is_file():
                    if ext is None or entry.name.lower().endswith(ext):
                        count += 1
            return count

        stats['src_frames'] = count_files(os.path.join(workspace_dir, 'data_src'), '.png') + count_files(os.path.join(workspace_dir, 'data_src'), '.jpg')
        stats['src_faces'] = count_files(os.path.join(workspace_dir, 'data_src', 'aligned'), '.jpg')
        stats['dst_frames'] = count_files(os.path.join(workspace_dir, 'data_dst'), '.png') + count_files(os.path.join(workspace_dir, 'data_dst'), '.jpg')
        stats['dst_faces'] = count_files(os.path.join(workspace_dir, 'data_dst', 'aligned'), '.jpg')
        
        model_dir = os.path.join(workspace_dir, 'model')
        if os.path.exists(model_dir):
            for entry in os.scandir(model_dir):
                if entry.is_file() and entry.name.endswith('.dat'):
                    stats['models'].append(entry.name)

    return jsonify({
        'running': is_running,
        'script': running_script,
        'gpu': get_gpu_info(),
        'system': get_system_stats(),
        'workspace': stats
    })

@app.route('/api/scripts')
def list_scripts():
    return jsonify(SCRIPT_CATEGORIES)

@app.route('/api/run', methods=['POST'])
def run_script():
    global current_process, master_fd, log_buffer, running_script
    if current_process is not None and current_process.poll() is None:
        return jsonify({'error': 'A script is already running'}), 400
        
    data = request.json
    script_name = data.get('script')
    if not script_name:
        return jsonify({'error': 'No script name specified'}), 400
        
    script_path = os.path.join(os.path.dirname(__file__), 'scripts', script_name)
    if not os.path.exists(script_path):
        return jsonify({'error': f'Script not found: {script_name}'}), 404
        
    # Reset log buffer
    with log_lock:
        log_buffer = f"=== Starting {script_name} ===\n"
        running_script = script_name
        
    # Spawn the script in a pseudo-terminal
    try:
        m_fd, s_fd = pty.openpty()
        
        # We spawn the script inside the conda env's python path by launching /bin/bash in pty
        current_process = subprocess.Popen(
            ['/bin/bash', script_name],
            cwd=os.path.join(os.path.dirname(__file__), 'scripts'),
            stdin=s_fd,
            stdout=s_fd,
            stderr=s_fd,
            preexec_fn=os.setsid, # Put subprocess in its own process group to kill children
            env=os.environ.copy()
        )
        
        # Close the slave file descriptor in the parent
        os.close(s_fd)
        
        master_fd = m_fd
        
        # Start background reader thread
        t = threading.Thread(target=log_reader, args=(m_fd, current_process))
        t.daemon = True
        t.start()
        
        return jsonify({'status': 'started', 'script': script_name})
    except Exception as e:
        with log_lock:
            log_buffer += f"\nFailed to launch script: {str(e)}\n"
            current_process = None
            master_fd = None
            running_script = None
        return jsonify({'error': f'Launch error: {str(e)}'}), 500

@app.route('/api/stop', methods=['POST'])
def stop_script():
    global current_process
    if current_process is None or current_process.poll() is not None:
        return jsonify({'error': 'No script is currently running'}), 400
        
    try:
        # Kill the entire process group
        pgid = os.getpgid(current_process.pid)
        os.killpg(pgid, signal.SIGTERM)
        time.sleep(0.5)
        if current_process.poll() is None:
            os.killpg(pgid, signal.SIGKILL)
            
        with log_lock:
            log_buffer += "\n=== Process terminated by user ===\n"
            
        return jsonify({'status': 'terminated'})
    except Exception as e:
        return jsonify({'error': f'Error terminating process: {str(e)}'}), 500

@app.route('/api/input', methods=['POST'])
def send_input():
    global master_fd
    data = request.json
    text = data.get('text', '')
    if master_fd is None:
        return jsonify({'error': 'No active process terminal to write to'}), 400
        
    try:
        os.write(master_fd, (text + '\n').encode('utf-8'))
        return jsonify({'status': 'sent'})
    except Exception as e:
        return jsonify({'error': f'Write error: {str(e)}'}), 500

@app.route('/api/logs')
def get_logs():
    global log_buffer
    with log_lock:
        return jsonify({'logs': log_buffer})

@app.route('/api/preview')
def list_previews():
    # Scan for any images in the workspace/model directory
    model_dir = os.path.join(os.path.dirname(__file__), 'workspace', 'model')
    if not os.path.exists(model_dir):
        return jsonify([])
        
    previews = []
    try:
        for root, dirs, files in os.walk(model_dir):
            for f in files:
                if f.lower().endswith(('.jpg', '.png')):
                    path = os.path.join(root, f)
                    mtime = os.path.getmtime(path)
                    rel_path = os.path.relpath(path, model_dir)
                    # Exclude huge files if any
                    size = os.path.getsize(path)
                    if size < 5 * 1024 * 1024: # Less than 5MB
                        previews.append({
                            'name': f,
                            'path': rel_path,
                            'mtime': mtime,
                            'size': size
                        })
        # Sort by modification time (most recent first)
        previews.sort(key=lambda x: x['mtime'], reverse=True)
    except Exception as e:
        return jsonify({'error': str(e)}), 500
        
    return jsonify(previews)

@app.route('/api/preview/file/<path:filename>')
def serve_preview_file(filename):
    model_dir = os.path.join(os.path.dirname(__file__), 'workspace', 'model')
    return send_from_directory(model_dir, filename)

@app.route('/api/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    target = request.form.get('target') # 'src' or 'dst'
    
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
        
    if target not in ('src', 'dst'):
        return jsonify({'error': 'Invalid upload target'}), 400
        
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in ('.mp4', '.avi', '.mkv', '.mov'):
        return jsonify({'error': 'Unsupported file format'}), 400
        
    workspace_dir = os.path.join(os.path.dirname(__file__), 'workspace')
    os.makedirs(workspace_dir, exist_ok=True)
    
    # Save as data_src.<ext> or data_dst.<ext>
    filename = f"data_{target}{ext}"
    dest_path = os.path.join(workspace_dir, filename)
    
    # Remove existing ones with other extensions if necessary, to keep DFL working cleanly
    for existing_ext in ('.mp4', '.avi', '.mkv', '.mov'):
        try:
            os.remove(os.path.join(workspace_dir, f"data_{target}{existing_ext}"))
        except FileNotFoundError:
            pass
            
    try:
        file.save(dest_path)
        return jsonify({'status': 'uploaded', 'filename': filename})
    except Exception as e:
        return jsonify({'error': f'Failed to save file: {str(e)}'}), 500

if __name__ == '__main__':
    # Make sure workspace directories exist
    os.makedirs('templates', exist_ok=True)
    os.makedirs('static', exist_ok=True)
    
    port = 8082
    print(f"Starting DeepFaceLab Web UI Server on port {port}...")
    print(f"Access it at http://localhost:{port}/")
    app.run(host='0.0.0.0', port=port, debug=False)