import eventlet
eventlet.monkey_patch()

from flask import Flask
from flask_socketio import SocketIO, emit
import time
import torch
import cv2
import os
import base64
from core.dreamer import Dreamer
from train import train_with_socket
import core.generator as gen

app = Flask(__name__)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='eventlet')

dreamer = Dreamer()
is_dreaming = True
is_paused = False

@socketio.on('connect')
def handle_connect():
    print('!!! [SOCKET] NEW_CLIENT_CONNECTED !!!', flush=True)
    emit('status', {'data': 'Connected to DREAMER_ENGINE_V2'})

@socketio.on('set_speed')
def handle_set_speed(data):
    speed = float(data.get('speed', 0.03))
    dreamer.speed = speed
    print(f"[ENGINE] SPEED_ADJUSTED: {speed}", flush=True)

@socketio.on('set_pause')
def handle_set_pause(data):
    global is_paused
    is_paused = data.get('paused', False)
    print(f"[ENGINE] PAUSE_STATE: {is_paused}", flush=True)

@socketio.on('ping_engine')
def handle_ping(data=None):
    print(f'>>> [SOCKET] PING_RECEIVED <<<', flush=True)
    emit('gen_status', {'msg': 'PONG: ENGINE_ALIVE_ON_5001'})
    return "PONG"

@socketio.on('list_models')
def handle_list_models(data=None):
    weights_dir = "data/weights"
    if not os.path.exists(weights_dir):
        os.makedirs(weights_dir)
    files = [f for f in os.listdir(weights_dir) if f.endswith('.pth')]
    return {'models': files}

@socketio.on('select_model')
def handle_select_model(data=None):
    if data is None: data = {}
    model_name = data.get('name')
    print(f"[ENGINE] HOT_SWAPPING_MODEL: {model_name}", flush=True)
    try:
        weights_path = os.path.join("data/weights", model_name)
        if os.path.exists(weights_path):
            dreamer.model.load_state_dict(torch.load(weights_path, map_location=dreamer.device))
            emit('model_status', {'msg': f'SWAP_SUCCESS: {model_name}'})
            return "SUCCESS"
        else:
            return "FILE_NOT_FOUND"
    except Exception as e:
        emit('model_status', {'msg': f'ERR: {str(e)}'})
        return str(e)

@socketio.on('save_model')
def handle_save_model(data=None):
    if data is None: data = {}
    custom_name = data.get('name', 'vae_circuit_backup')
    if not custom_name.endswith('.pth'): custom_name += '.pth'
    
    print(f"[ENGINE] ARCHIVING_WEIGHTS_AS: {custom_name}", flush=True)
    try:
        import shutil
        src = "data/weights/vae_circuit.pth"
        dst = os.path.join("data/weights", custom_name)
        if os.path.exists(src):
            shutil.copy(src, dst)
            emit('model_status', {'msg': f'SAVED: {custom_name}'})
            return "SUCCESS"
        else:
            emit('model_status', {'msg': 'ERR: NO_WEIGHTS_FOUND'})
            return "ERROR"
    except Exception as e:
        emit('model_status', {'msg': f'ERR: {str(e)}'})
        return "ERROR"

@socketio.on('load_model')
def handle_load_model(data=None):
    if data is None: data = {}
    model_name = data.get('name', 'vae_circuit.pth')
    
    print(f"[ENGINE] RELOADING_WEIGHTS_FROM: {model_name}", flush=True)
    try:
        weights_path = os.path.join("data/weights", model_name)
        if os.path.exists(weights_path):
            dreamer.model.load_state_dict(torch.load(weights_path, map_location=dreamer.device))
            emit('model_status', {'msg': f'LOADED: {model_name}'})
            return "SUCCESS"
        else:
            emit('model_status', {'msg': 'ERR: FILE_NOT_FOUND'})
            return "ERROR"
    except Exception as e:
        emit('model_status', {'msg': f'ERR: {str(e)}'})
        return "ERROR"

@socketio.on('start_generation')
def handle_start_generation(data=None):
    print(f"!!! [SOCKET] START_GEN_RECEIVED !!!", flush=True)
    if data is None: data = {}
    count = int(data.get('count', 1000))
    print(f"[ENGINE] UI_REQUEST: SYNTHESIZE {count} SAMPLES", flush=True)
    socketio.emit('gen_status', {'msg': 'INITIALIZING...'})
    socketio.start_background_task(run_generation, count)
    return "ACK"

def run_generation(count):
    print(f"[ENGINE] PURGING_OLD_DATA...", flush=True)
    train_dir = "data/training"
    try:
        import shutil
        if os.path.exists(train_dir):
            shutil.rmtree(train_dir)
        os.makedirs(train_dir)
    except Exception as e:
        print(f"[ENGINE] PURGE_ERR: {e}", flush=True)

    print(f"[ENGINE] STARTING_SYNTHESIS: {count} iterations", flush=True)
    
    try:
        for i in range(count):
            img = gen.generate_circuit_image()
            cv2.imwrite(os.path.join("data/training", f'circuit_{i:04d}.png'), img)
            if (i+1) % 50 == 0 or i == count - 1:
                print(f"[ENGINE] PROGRESS: {i+1}/{count}", flush=True)
                socketio.emit('gen_status', {'msg': f'PROG: {i+1}/{count}'})
                socketio.sleep(0.01)
        
        print("[ENGINE] SYNTHESIS_COMPLETE", flush=True)
        socketio.emit('gen_complete', {'msg': f'SUCCESS: {count} IMG'})
    except Exception as e:
        print(f"[ENGINE] ERROR_IN_SYNTH: {str(e)}", flush=True)
        socketio.emit('gen_error', {'msg': f'ERR: {str(e)}'})

@socketio.on('start_training')
def handle_start_training(data=None):
    global is_dreaming
    print(f"!!! [SOCKET] START_TRAIN_RECEIVED !!!", flush=True)
    if not is_dreaming:
        print("[ENGINE] REJECT_TRAIN: Engine busy.", flush=True)
        return "BUSY"

    if data is None: data = {}
    train_dir = "data/training"
    if not os.path.exists(train_dir) or len(os.listdir(train_dir)) == 0:
        print("[ENGINE] ABORT_TRAIN: No data found.", flush=True)
        emit('train_error', {'msg': 'NO_SAMPLES_FOUND. RUN_SYNTHESIZER_FIRST.'})
        return "ERROR"

    epochs = int(data.get('epochs', 10))
    batch_size = int(data.get('batch_size', 16))
    print(f"\n[ENGINE] UI_REQUEST: START_TRAINING ({epochs} epochs)", flush=True)
    socketio.start_background_task(run_training, epochs, batch_size)
    return "ACK"

def run_training(epochs, batch_size):
    global is_dreaming
    is_dreaming = False
    print(f"[ENGINE] OFFLOADING_DREAMER_TO_CPU", flush=True)
    dreamer.model.to("cpu")
    
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    import gc
    gc.collect()
    
    print(f"[ENGINE] RESET_EPOCH. STARTING AT 0.", flush=True)
    
    try:
        train_with_socket(socketio, epochs=epochs, batch_size=batch_size)
        weights_path = os.path.join("data/weights/vae_circuit.pth")
        if os.path.exists(weights_path):
            print(f"[ENGINE] SYNCING_WEIGHTS_TO_GPU", flush=True)
            dreamer.model.load_state_dict(torch.load(weights_path, map_location="cpu"))
            dreamer.model.to(dreamer.device)
    except Exception as e:
        print(f"[ENGINE] ERROR_IN_TRAIN: {str(e)}", flush=True)
        socketio.emit('train_error', {'msg': str(e)})
        dreamer.model.to(dreamer.device)
    finally:
        is_dreaming = True

def heartbeat():
    while True:
        socketio.emit('heartbeat', {'time': time.time()})
        socketio.sleep(5)

def dream_loop():
    print("[ENGINE] STARTING_DREAM_LOOP", flush=True)
    while True:
        if is_dreaming and not is_paused:
            try:
                frame_b64, latent_coords = dreamer.step()
                socketio.emit('frame', {
                    'image': frame_b64,
                    'z': latent_coords
                })
            except Exception as e:
                print(f"[DREAM_ERR] {e}", flush=True)
        socketio.sleep(0.033)

if __name__ == '__main__':
    socketio.start_background_task(dream_loop)
    socketio.start_background_task(heartbeat)
    print("!!! ENGINE_SVR_RUNNING_ON_PORT_5001 !!!", flush=True)
    socketio.run(app, host='0.0.0.0', port=5001)
