CLOUDY COLLABORATIONS

"""
THE SOVEREIGN OMNISTUDIO — CENTRAL ORCHESTRATOR
FastAPI + WebSockets — Single brain for all engines.
Run: python orchestrator.py
Requires: pip install fastapi uvicorn websockets
Ports: REST :9000 | WebSocket ws://127.0.0.1:9000/ws
"""
import os
import json
import time
import uuid
import asyncio
import threading
import urllib.request
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

app = FastAPI(title="Sovereign OmniStudio Orchestrator", version="3.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

DRIVE_ROOT = r"E:\CLOUDY_DRIVE"
COMFY_URL  = "http://127.0.0.1:8188"
BLENDER_WS = "ws://127.0.0.1:8765"
RIG_URL    = "http://127.0.0.1:8189"
LIPSYNC_URL= "http://127.0.0.1:8190"
DAEMON_URL = "http://127.0.0.1:8191"
LLM_URL    = "http://127.0.0.1:11434"   # Ollama

active_ws_clients: list[WebSocket] = []
job_registry: dict = {}

# ─── WebSocket hub ───────────────────────────────────────────────────────────
async def broadcast(msg: dict):
    dead = []
    for ws in active_ws_clients:
        try:
            await ws.send_json(msg)
        except Exception:
            dead.append(ws)
    for d in dead:
        active_ws_clients.remove(d)

@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    active_ws_clients.append(ws)
    try:
        while True:
            data = await ws.receive_json()
            await route_command(data, ws)
    except WebSocketDisconnect:
        active_ws_clients.remove(ws)

async def route_command(cmd: dict, ws: WebSocket):
    action = cmd.get("action", "")
    job_id = str(uuid.uuid4())[:8]
    job_registry[job_id] = {"status": "running", "action": action}
    await ws.send_json({"job_id": job_id, "status": "started", "action": action})
    try:
        if action == "RENDER_IMAGE":
            result = await dispatch_comfy(cmd.get("workflow", {}), job_id)
        elif action == "RENDER_SPRITES":
            result = await dispatch_blender(cmd, job_id)
        elif action == "AUTO_RIG":
            result = dispatch_rig(cmd, job_id)
        elif action == "LIP_SYNC":
            result = dispatch_lipsync(cmd, job_id)
        elif action == "LLM_CHAT":
            result = await dispatch_llm(cmd.get("prompt", ""), job_id)
        elif action == "STATUS":
            result = await get_all_status()
        else:
            result = {"error": f"Unknown action: {action}"}
        job_registry[job_id] = {"status": "done", "result": result}
        await ws.send_json({"job_id": job_id, "status": "done", "result": result})
    except Exception as e:
        job_registry[job_id] = {"status": "error", "error": str(e)}
        await ws.send_json({"job_id": job_id, "status": "error", "error": str(e)})

# ─── Engine dispatchers ───────────────────────────────────────────────────────
async def dispatch_comfy(workflow: dict, job_id: str) -> dict:
    if not workflow:
        workflow = default_txt2img_workflow()
    payload = json.dumps({"prompt": workflow}).encode()
    req = urllib.request.Request(f"{COMFY_URL}/prompt", data=payload,
                                  headers={"Content-Type": "application/json"})
    resp = urllib.request.urlopen(req)
    return {"comfy_job": json.loads(resp.read()).get("prompt_id", "queued")}

async def dispatch_blender(cmd: dict, job_id: str) -> dict:
    import websockets
    async with websockets.connect(BLENDER_WS) as ws:
        await ws.send(json.dumps({
            "action": "AUTO_RIG_AND_RENDER_PNG",
            "mesh_path": cmd.get("mesh_path", ""),
            "angles": cmd.get("angles", 8),
            "output_dir": os.path.join(DRIVE_ROOT, "05_OUTPUT_STREAM", "sprites"),
        }))
        return json.loads(await ws.recv())

def dispatch_rig(cmd: dict, job_id: str) -> dict:
    data = json.dumps({"mesh_path": cmd.get("mesh_path", "")}).encode()
    req = urllib.request.Request(f"{RIG_URL}/auto-rig", data=data,
                                  headers={"Content-Type": "application/json"})
    resp = urllib.request.urlopen(req)
    return json.loads(resp.read())

def dispatch_lipsync(cmd: dict, job_id: str) -> dict:
    data = json.dumps({"audio_path": cmd.get("audio_path", "")}).encode()
    req = urllib.request.Request(f"{LIPSYNC_URL}/transcribe", data=data,
                                  headers={"Content-Type": "application/json"})
    resp = urllib.request.urlopen(req)
    return json.loads(resp.read())

async def dispatch_llm(prompt: str, job_id: str) -> dict:
    """Call local Ollama — completely offline, no API key needed."""
    data = json.dumps({"model": "deepseek-coder:6.7b", "prompt": prompt, "stream": False}).encode()
    req = urllib.request.Request(f"{LLM_URL}/api/generate", data=data,
                                  headers={"Content-Type": "application/json"})
    try:
        resp = urllib.request.urlopen(req, timeout=60)
        result = json.loads(resp.read())
        return {"response": result.get("response", "")}
    except Exception as e:
        return {"error": f"LLM offline: {e}. Start Ollama: ollama serve"}

async def get_all_status() -> dict:
    def ping(url):
        try:
            urllib.request.urlopen(url, timeout=1)
            return "online"
        except:
            return "offline"
    return {
        "comfyui":  ping(f"{COMFY_URL}/system_stats"),
        "blender":  "online" if True else "offline",  # ws check
        "accurig":  ping(f"{RIG_URL}/status"),
        "lipsync":  ping(f"{LIPSYNC_URL}/status"),
        "daemon":   ping(f"{DAEMON_URL}/status"),
        "ollama":   ping(f"{LLM_URL}/api/tags"),
    }

# ─── REST endpoints (for frontend polling) ──────────────────────────────────
@app.get("/status")
async def status():
    return await get_all_status()

@app.get("/jobs")
async def list_jobs():
    return {"jobs": list(job_registry.values())[-20:]}

@app.post("/render")
async def rest_render(body: dict):
    job_id = str(uuid.uuid4())[:8]
    result = await dispatch_comfy(body.get("workflow", {}), job_id)
    return {"job_id": job_id, **result}

# ─── Default ComfyUI workflow (txt2img with SD1.5) ──────────────────────────
def default_txt2img_workflow():
    return {
        "4": {"inputs": {"ckpt_name": "v1-5-pruned-emaonly.ckpt"}, "class_type": "CheckpointLoaderSimple"},
        "6": {"inputs": {"text": "a professional 3D character, full body, transparent background", "clip": ["4", 1]}, "class_type": "CLIPTextEncode"},
        "7": {"inputs": {"text": "ugly, blurry, watermark", "clip": ["4", 1]}, "class_type": "CLIPTextEncode"},
        "5": {"inputs": {"width": 512, "height": 512, "batch_size": 1}, "class_type": "EmptyLatentImage"},
        "3": {"inputs": {"seed": 0, "steps": 20, "cfg": 7, "sampler_name": "euler", "scheduler": "normal",
                          "denoise": 1, "model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0],
                          "latent_image": ["5", 0]}, "class_type": "KSampler"},
        "8": {"inputs": {"samples": ["3", 0], "vae": ["4", 2]}, "class_type": "VAEDecode"},
        "9": {"inputs": {"filename_prefix": "sovereign_render", "images": ["8", 0]}, "class_type": "SaveImage"},
    }

if __name__ == "__main__":
    import uvicorn
    print("=" * 60)
    print(" SOVEREIGN OMNISTUDIO — CENTRAL ORCHESTRATOR")
    print(" REST:      http://127.0.0.1:9000")
    print(" WebSocket: ws://127.0.0.1:9000/ws")
    print(" Drive:     " + DRIVE_ROOT)
    print("=" * 60)
    uvicorn.run(app, host="127.0.0.1", port=9000, log_level="warning")