#////////// PROJECT FILE SEPARATION LINE ////////// CODE AFTER THIS LINE ARE FROM: <ProjectRoot>/Data/Agent/borealis-agent.py

import sys
import uuid
import socket
import os
import json
import asyncio
import concurrent.futures
from functools import partial
from io import BytesIO
import base64

import socketio
from qasync import QEventLoop
from PyQt5 import QtCore, QtGui, QtWidgets
from PIL import ImageGrab

# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: CONFIG MANAGER (do not modify unless you know what you’re doing)
# //////////////////////////////////////////////////////////////////////////
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "agent_settings.json")
DEFAULT_CONFIG = {
    "borealis_server_url": "http://localhost:5000",
    "max_task_workers": 8,
    "config_file_watcher_interval": 2,
    "agent_id": "",
    "regions": {}
}

class ConfigManager:
    def __init__(self, path):
        self.path = path
        self._last_mtime = None
        self.data = {}
        self.load()

    def load(self):
        # load or initialize
        if not os.path.exists(self.path):
            self.data = DEFAULT_CONFIG.copy()
            self._write()
        else:
            try:
                with open(self.path, 'r') as f:
                    loaded = json.load(f)
                # merge defaults
                self.data = {**DEFAULT_CONFIG, **loaded}
            except Exception as e:
                print(f"[WARN] Failed to parse config: {e}")
                self.data = DEFAULT_CONFIG.copy()
        # track mtime
        try:
            self._last_mtime = os.path.getmtime(self.path)
        except Exception:
            self._last_mtime = None

    def _write(self):
        try:
            with open(self.path, 'w') as f:
                json.dump(self.data, f, indent=2)
        except Exception as e:
            print(f"[ERROR] Could not write config: {e}")

    def watch(self):
        try:
            mtime = os.path.getmtime(self.path)
            if self._last_mtime is None or mtime != self._last_mtime:
                print("[CONFIG] Detected config change, reloading.")
                self.load()
                return True
        except Exception:
            pass
        return False

CONFIG = ConfigManager(CONFIG_PATH)
# Purge saved regions on startup (fresh run)
CONFIG.data['regions'] = {}
CONFIG._write()
# //////////////////////////////////////////////////////////////////////////
# END CORE SECTION: CONFIG MANAGER
# //////////////////////////////////////////////////////////////////////////

# Assign or persist agent_id
host = socket.gethostname().lower()
stored_id = CONFIG.data.get('agent_id')
if stored_id:
    AGENT_ID = stored_id
else:
    AGENT_ID = f"{host}-agent-{uuid.uuid4().hex[:8]}"
    CONFIG.data['agent_id'] = AGENT_ID
    CONFIG._write()

# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: WEBSOCKET SETUP & HANDLERS (do not modify unless absolutely necessary)
# //////////////////////////////////////////////////////////////////////////

sio = socketio.AsyncClient(reconnection=True, reconnection_attempts=0, reconnection_delay=5)
role_tasks = {}
overlay_widgets = {}

@sio.event
async def connect():
    print(f"[WebSocket] Connected to Agent ID: {AGENT_ID}.")
    await sio.emit('connect_agent', {"agent_id": AGENT_ID})
    await sio.emit('request_config', {"agent_id": AGENT_ID})

@sio.event
async def disconnect():
    print("[WebSocket] Disconnected from Borealis server.")
    # reset tasks and overlays
    for task in list(role_tasks.values()):
        task.cancel()
    role_tasks.clear()
    for widget in list(overlay_widgets.values()):
        try: widget.close()
        except: pass
    overlay_widgets.clear()
    # purge regions on intentional disconnect
    CONFIG.data['regions'].clear()
    CONFIG._write()
    # reload settings
    CONFIG.load()

@sio.on('agent_config')
async def on_agent_config(cfg):
    print(f"[CONNECTED] Received config with {len(cfg.get('roles',[]))} roles.")
    # determine removed roles
    new_ids = {r.get('node_id') for r in cfg.get('roles', []) if r.get('node_id')}
    old_ids = set(role_tasks.keys())
    removed = old_ids - new_ids
    for rid in removed:
        # remove region config
        if rid in CONFIG.data['regions']:
            CONFIG.data['regions'].pop(rid, None)
        # close overlay
        w = overlay_widgets.pop(rid, None)
        if w:
            try: w.close()
            except: pass
    if removed:
        CONFIG._write()
    # cancel existing and start new
    for task in list(role_tasks.values()):
        task.cancel()
    role_tasks.clear()
    for role_cfg in cfg.get('roles', []):
        if role_cfg.get('role') == 'screenshot':
            nid = role_cfg.get('node_id')
            t = asyncio.create_task(screenshot_task(role_cfg))
            role_tasks[nid] = t
# //////////////////////////////////////////////////////////////////////////
# END CORE SECTION: WEBSOCKET SETUP & HANDLERS
# //////////////////////////////////////////////////////////////////////////

# ---------------- Overlay Widget ----------------
class ScreenshotRegion(QtWidgets.QWidget):
    def __init__(self, node_id, x=100, y=100, w=300, h=200):
        super().__init__()
        self.node_id = node_id
        self.setGeometry(x, y, w, h)
        self.setWindowFlags(QtCore.Qt.FramelessWindowHint|QtCore.Qt.WindowStaysOnTopHint)
        self.setAttribute(QtCore.Qt.WA_TranslucentBackground)
        self.drag_offset = None
        self.resizing = False
        self.resize_handle_size = 12
        self.setVisible(True)
        self.label = QtWidgets.QLabel(self)
        self.label.setText(f"{node_id[:8]}")
        self.label.setStyleSheet("color: lime; background: transparent; font-size: 10px;")
        self.label.move(8,4)
        self.setMouseTracking(True)

    def paintEvent(self, event):
        p = QtGui.QPainter(self)
        p.setRenderHint(QtGui.QPainter.Antialiasing)
        p.setBrush(QtCore.Qt.transparent)
        p.setPen(QtGui.QPen(QtGui.QColor(0,255,0),2))
        p.drawRect(self.rect())
        hr = self.resize_handle_size
        hrect = QtCore.QRect(self.width()-hr, self.height()-hr, hr, hr)
        p.fillRect(hrect, QtGui.QColor(0,255,0))

    def mousePressEvent(self, e):
        if e.button()==QtCore.Qt.LeftButton:
            x,y = e.pos().x(), e.pos().y()
            if x>self.width()-self.resize_handle_size and y>self.height()-self.resize_handle_size:
                self.resizing=True
            else:
                self.drag_offset = e.globalPos() - self.frameGeometry().topLeft()

    def mouseMoveEvent(self, e):
        if self.resizing:
            nw = max(e.pos().x(),100)
            nh = max(e.pos().y(),80)
            self.resize(nw,nh)
        elif e.buttons()&QtCore.Qt.LeftButton and self.drag_offset:
            self.move(e.globalPos()-self.drag_offset)

    def mouseReleaseEvent(self,e):
        self.resizing=False; self.drag_offset=None

    def get_geometry(self):
        g=self.geometry(); return (g.x(),g.y(),g.width(),g.height())

# ---------------- Screenshot Task ----------------
async def screenshot_task(cfg):
    nid = cfg.get('node_id');
    # initial region from config or payload
    r = CONFIG.data['regions'].get(nid)
    region = (r['x'],r['y'],r['w'],r['h']) if r else (cfg.get('x',100),cfg.get('y',100),cfg.get('w',300),cfg.get('h',200))
    if nid not in overlay_widgets:
        widget = ScreenshotRegion(nid,*region)
        overlay_widgets[nid] = widget; widget.show()
    interval = cfg.get('interval',1000)/1000.0
    loop = asyncio.get_event_loop()
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=CONFIG.data.get('max_task_workers',DEFAULT_CONFIG['max_task_workers']))
    try:
        while True:
            x,y,w,h = overlay_widgets[nid].get_geometry()
            # persist if changed
            prev = CONFIG.data['regions'].get(nid)
            if prev!={'x':x,'y':y,'w':w,'h':h}:
                CONFIG.data['regions'][nid]={'x':x,'y':y,'w':w,'h':h}
                CONFIG._write()
            grab = partial(ImageGrab.grab,bbox=(x,y,x+w,y+h))
            img = await loop.run_in_executor(executor,grab)
            buf = BytesIO(); img.save(buf,format='PNG')
            encoded = base64.b64encode(buf.getvalue()).decode('utf-8')
            await sio.emit('agent_screenshot_task',{'agent_id':AGENT_ID,'node_id':nid,'image_base64':encoded})
            await asyncio.sleep(interval)
    except asyncio.CancelledError:
        return
    except Exception as e:
        print(f"[ERROR] Screenshot task {nid} failed: {e}")

# ---------------- Config Watcher ----------------
async def config_watcher():
    while True:
        if CONFIG.watch(): pass
        await asyncio.sleep(CONFIG.data.get('config_file_watcher_interval',DEFAULT_CONFIG['config_file_watcher_interval']))

# //////////////////////////////////////////////////////////////////////////
# CORE SECTION: MAIN & EVENT LOOP (do not modify unless you know what you’re doing)
# //////////////////////////////////////////////////////////////////////////
async def connect_loop():
    retry=5
    while True:
        try:
            url=CONFIG.data.get('borealis_server_url',DEFAULT_CONFIG['borealis_server_url'])
            print(f"[WebSocket] Connecting to {url}...")
            await sio.connect(url,transports=['websocket'])
            break
        except:
            print(f"[WebSocket] Server unavailable, retrying in {retry}s...")
            await asyncio.sleep(retry)

if __name__=='__main__':
    app = QtWidgets.QApplication(sys.argv)
    loop = QEventLoop(app)
    asyncio.set_event_loop(loop)
    with loop:
        loop.create_task(config_watcher())
        loop.create_task(connect_loop())
        loop.run_forever()
# //////////////////////////////////////////////////////////////////////////
# END CORE SECTION: MAIN & EVENT LOOP
# //////////////////////////////////////////////////////////////////////////