"""
THE Z AI — Computer Mode Server
================================
سيرفر التحكم بالكمبيوتر عن بعد للذكاء الاصطناعي
يوفر:
  - WebSocket للتحكم اللحظي (الماوس، لوحة المفاتيح، الترمينال)
  - تصوير سكرين شوت وإرساله للعميل
  - بث شاشة مباشر (base64 frames)
  - تنفيذ أوامر bash وقراءة الناتج
"""

import asyncio
import base64
import io
import json
import os
import subprocess
import threading
import time
import traceback
from pathlib import Path

import pyautogui
import pyperclip
from PIL import Image, ImageGrab
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn

# ─── ENV ────────────────────────────────────────────
DISPLAY = os.environ.get("DISPLAY", ":1")
os.environ["DISPLAY"] = DISPLAY
pyautogui.FAILSAFE = False
pyautogui.PAUSE = 0.05

app = FastAPI(title="Z-Computer-Mode API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ─── Active Connections ──────────────────────────────
active_connections: list[WebSocket] = []
stream_active = False
stream_quality = 60   # JPEG quality
stream_fps    = 3     # frames per second (low to save bandwidth)
stream_scale  = 0.5   # downscale factor


# ─── Utility ─────────────────────────────────────────

def capture_screen(scale=stream_scale, quality=stream_quality) -> str:
    """Capture screen → base64 JPEG string"""
    try:
        img = ImageGrab.grab()
        if scale < 1.0:
            w = int(img.width * scale)
            h = int(img.height * scale)
            img = img.resize((w, h), Image.LANCZOS)
        buf = io.BytesIO()
        img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
        return base64.b64encode(buf.getvalue()).decode()
    except Exception as e:
        return ""


def run_command(cmd: str, timeout: int = 30) -> dict:
    """Run a shell command and return stdout/stderr"""
    try:
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True,
            timeout=timeout, env={**os.environ, "DISPLAY": DISPLAY}
        )
        return {
            "stdout": result.stdout[-8000:] if len(result.stdout) > 8000 else result.stdout,
            "stderr": result.stderr[-4000:] if len(result.stderr) > 4000 else result.stderr,
            "returncode": result.returncode,
        }
    except subprocess.TimeoutExpired:
        return {"stdout": "", "stderr": f"Command timed out after {timeout}s", "returncode": -1}
    except Exception as e:
        return {"stdout": "", "stderr": str(e), "returncode": -1}


async def broadcast(msg: dict):
    """Send JSON message to all connected WebSocket clients"""
    txt = json.dumps(msg)
    disconnected = []
    for ws in active_connections:
        try:
            await ws.send_text(txt)
        except Exception:
            disconnected.append(ws)
    for ws in disconnected:
        if ws in active_connections:
            active_connections.remove(ws)


# ─── Screen Streaming Loop ───────────────────────────

async def screen_stream_loop():
    global stream_active
    interval = 1.0 / stream_fps
    while stream_active and active_connections:
        try:
            frame = capture_screen()
            if frame:
                await broadcast({"type": "frame", "data": frame})
        except Exception:
            pass
        await asyncio.sleep(interval)
    stream_active = False


# ─── Action Handler ──────────────────────────────────

async def handle_action(ws: WebSocket, msg: dict):
    action = msg.get("action", "")
    data   = msg.get("data", {})

    # ── SCREENSHOT ──────────────────────────────────
    if action == "screenshot":
        frame = capture_screen(scale=0.8, quality=75)
        await ws.send_text(json.dumps({
            "type": "screenshot",
            "data": frame,
            "ts": int(time.time() * 1000)
        }))

    # ── TERMINAL COMMAND ─────────────────────────────
    elif action == "terminal":
        cmd = data.get("cmd", "")
        if cmd:
            result = run_command(cmd, timeout=data.get("timeout", 30))
            await ws.send_text(json.dumps({
                "type": "terminal_result",
                "cmd": cmd,
                "stdout": result["stdout"],
                "stderr": result["stderr"],
                "returncode": result["returncode"],
            }))
            # Auto-screenshot after terminal command
            await asyncio.sleep(0.5)
            frame = capture_screen(scale=0.7, quality=70)
            if frame:
                await ws.send_text(json.dumps({
                    "type": "screenshot",
                    "data": frame,
                    "ts": int(time.time() * 1000),
                    "auto": True
                }))

    # ── MOUSE MOVE ───────────────────────────────────
    elif action == "mouse_move":
        x, y = int(data.get("x", 0)), int(data.get("y", 0))
        pyautogui.moveTo(x, y, duration=0.1)
        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_move"}))

    # ── MOUSE CLICK ──────────────────────────────────
    elif action == "mouse_click":
        x      = int(data.get("x", 0))
        y      = int(data.get("y", 0))
        button = data.get("button", "left")
        double = data.get("double", False)
        pyautogui.moveTo(x, y, duration=0.08)
        if double:
            pyautogui.doubleClick(x, y, button=button)
        else:
            pyautogui.click(x, y, button=button)
        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_click"}))

    # ── MOUSE DRAG ───────────────────────────────────
    elif action == "mouse_drag":
        x1, y1 = int(data.get("x1", 0)), int(data.get("y1", 0))
        x2, y2 = int(data.get("x2", 0)), int(data.get("y2", 0))
        pyautogui.moveTo(x1, y1)
        pyautogui.dragTo(x2, y2, duration=0.3, button="left")
        await ws.send_text(json.dumps({"type": "ack", "action": "mouse_drag"}))

    # ── KEYBOARD TYPE ────────────────────────────────
    elif action == "keyboard_type":
        text = data.get("text", "")
        pyautogui.typewrite(text, interval=0.03)
        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_type"}))

    # ── KEYBOARD HOTKEY ──────────────────────────────
    elif action == "keyboard_hotkey":
        keys = data.get("keys", [])
        if keys:
            pyautogui.hotkey(*keys)
        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_hotkey"}))

    # ── KEYBOARD PRESS ───────────────────────────────
    elif action == "keyboard_press":
        key = data.get("key", "")
        if key:
            pyautogui.press(key)
        await ws.send_text(json.dumps({"type": "ack", "action": "keyboard_press"}))

    # ── CLIPBOARD WRITE ──────────────────────────────
    elif action == "clipboard_write":
        text = data.get("text", "")
        pyperclip.copy(text)
        await ws.send_text(json.dumps({"type": "ack", "action": "clipboard_write"}))

    # ── CLIPBOARD READ ───────────────────────────────
    elif action == "clipboard_read":
        text = pyperclip.paste()
        await ws.send_text(json.dumps({"type": "clipboard_content", "text": text}))

    # ── SCROLL ───────────────────────────────────────
    elif action == "scroll":
        x      = int(data.get("x", 0))
        y      = int(data.get("y", 0))
        clicks = int(data.get("clicks", 3))
        pyautogui.scroll(clicks, x=x, y=y)
        await ws.send_text(json.dumps({"type": "ack", "action": "scroll"}))

    # ── OPEN APP ─────────────────────────────────────
    elif action == "open_app":
        app_cmd = data.get("cmd", "")
        if app_cmd:
            subprocess.Popen(
                app_cmd, shell=True,
                env={**os.environ, "DISPLAY": DISPLAY},
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
            )
        await ws.send_text(json.dumps({"type": "ack", "action": "open_app"}))
        # Auto-screenshot after app opens (wait for it to load)
        await asyncio.sleep(3)
        frame = capture_screen(scale=0.7, quality=70)
        if frame:
            await ws.send_text(json.dumps({
                "type": "screenshot",
                "data": frame,
                "ts": int(time.time() * 1000),
                "auto": True
            }))

    # ── START STREAM ─────────────────────────────────
    elif action == "start_stream":
        global stream_active, stream_fps, stream_quality, stream_scale
        stream_fps     = int(data.get("fps", 3))
        stream_quality = int(data.get("quality", 60))
        stream_scale   = float(data.get("scale", 0.5))
        if not stream_active:
            stream_active = True
            asyncio.create_task(screen_stream_loop())
        await ws.send_text(json.dumps({"type": "ack", "action": "start_stream"}))

    # ── STOP STREAM ──────────────────────────────────
    elif action == "stop_stream":
        stream_active = False
        await ws.send_text(json.dumps({"type": "ack", "action": "stop_stream"}))

    # ── GET SCREEN INFO ──────────────────────────────
    elif action == "screen_info":
        try:
            w, h = pyautogui.size()
            mx, my = pyautogui.position()
        except:
            w, h, mx, my = 1920, 1080, 0, 0
        await ws.send_text(json.dumps({
            "type": "screen_info",
            "width": w, "height": h,
            "mouse_x": mx, "mouse_y": my,
        }))

    else:
        await ws.send_text(json.dumps({"type": "error", "msg": f"Unknown action: {action}"}))


# ─── WebSocket Endpoint ──────────────────────────────

@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    active_connections.append(ws)
    # Send welcome
    try:
        w, h = pyautogui.size()
    except:
        w, h = 1920, 1080
    await ws.send_text(json.dumps({
        "type": "connected",
        "screen_width": w,
        "screen_height": h,
        "msg": "Z Computer Mode — Connected"
    }))
    try:
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            await handle_action(ws, msg)
    except WebSocketDisconnect:
        pass
    except Exception as e:
        pass
    finally:
        if ws in active_connections:
            active_connections.remove(ws)


# ─── REST Fallback ───────────────────────────────────

@app.get("/screenshot")
async def rest_screenshot():
    frame = capture_screen(scale=0.75, quality=70)
    return JSONResponse({"image": frame, "ts": int(time.time() * 1000)})


@app.post("/terminal")
async def rest_terminal(body: dict):
    cmd    = body.get("cmd", "")
    result = run_command(cmd, timeout=body.get("timeout", 30))
    return JSONResponse(result)


@app.get("/health")
async def health():
    return {"status": "ok", "display": DISPLAY}


# ─── Entry ───────────────────────────────────────────

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, log_level="warning")