# services/comfy-agent/app/nats_client.py
import json
import asyncio
from nats.aio.client import Client as NATS
from .config import settings
from .jobs import JOB_STORE
from .worker import enqueue

async def start_nats() -> NATS:
    nc = NATS()
    await nc.connect(servers=[settings.NATS_URL])

    async def handle(msg):
        subj = msg.subject
        reply = msg.reply
        payload = json.loads(msg.data.decode("utf-8"))

        # payload contract (MVP):
        # { "type": "text-to-image|text-to-video", "workflow": {...} }
        gen_type = payload.get("type", "text-to-image")
        workflow = payload.get("workflow")
        if not workflow:
            if reply:
                await nc.publish(reply, json.dumps({"error": "missing_workflow"}).encode())
            return

        job = JOB_STORE.create(gen_type)
        enqueue(job.job_id, gen_type, workflow)

        if reply:
            await nc.publish(reply, json.dumps({"job_id": job.job_id}).encode())

    await nc.subscribe(settings.NATS_SUBJECT_INVOKE, cb=handle)
    await nc.subscribe(settings.NATS_SUBJECT_IMAGE, cb=handle)
    await nc.subscribe(settings.NATS_SUBJECT_VIDEO, cb=handle)
    return nc