"""
NATS Publisher — публікація jobs у NATS JetStream
"""

from typing import Optional, Dict, Any
from nats.aio.client import Client as NATS
from nats.js import JetStreamContext


class NATSPublisher:
    def __init__(self, nats_url: str):
        self.nats_url = nats_url
        self.nc: Optional[NATS] = None
        self.js: Optional[JetStreamContext] = None
        
    async def connect(self):
        """Підключення до NATS"""
        try:
            self.nc = NATS()
            await self.nc.connect(self.nats_url)
            self.js = self.nc.jetstream()
            print(f"✅ Підключено до NATS: {self.nats_url}")
        except Exception as e:
            print(f"❌ Помилка підключення до NATS: {e}")
            raise
            
    async def disconnect(self):
        """Відключення від NATS"""
        if self.nc:
            await self.nc.close()
            print("✅ Відключено від NATS")
            
    async def publish_job(self, job: Dict[str, Any]):
        """Публікація job у NATS stream"""
        if not self.js:
            raise RuntimeError("NATS не підключено")
            
        job_type = job.get("type", "unknown")
        priority = job.get("priority", "offline")
        
        # Визначення subject залежно від типу та пріоритету
        if priority == "online":
            if job_type == "embed":
                subject = "mm.embed.online"
            elif job_type == "retrieve":
                subject = "mm.retrieve.online"
            elif job_type == "summarize":
                subject = "mm.summarize.online"
            else:
                subject = f"mm.{job_type}.online"
        else:
            if job_type == "embed":
                subject = "mm.embed.offline"
            elif job_type == "index":
                subject = "mm.index.offline"
            elif job_type == "backfill":
                subject = "mm.backfill.offline"
            else:
                subject = f"mm.{job_type}.offline"
                
        # Публікація
        import json
        job_json = json.dumps(job).encode()
        
        ack = await self.js.publish(subject, job_json)
        print(f"📤 Job опубліковано: {subject} (seq: {ack.seq})")
        
        return ack