#!/usr/bin/env python3
"""
Emergent Language Collector v2 - High Performance Sidecar

Key optimizations:
1. Fire-and-forget: Agents don't wait for translation response
2. Parallel batchers: Multiple workers sending batches concurrently
3. Shared queue: All agents feed one queue, batches fill faster

Performance: 9,721 msg/s (10.5x vs individual, 2.8x vs client batching)
"""

import asyncio
import time
import httpx
from typing import Dict, List
from dataclasses import dataclass, field
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn

# ═══════════════════════════════════════════════════════════════════════════════
# Configuration
# ═══════════════════════════════════════════════════════════════════════════════

TRANSLATOR_API = "https://emergent-language.fly.dev"
AUTH_TOKEN = "eudaimonia-translator-demo"
BATCH_SIZE = 75
NUM_BATCHERS = 4
BATCH_TIMEOUT_MS = 10


# ═══════════════════════════════════════════════════════════════════════════════
# Collector
# ═══════════════════════════════════════════════════════════════════════════════

@dataclass
class Stats:
    messages_received: int = 0
    messages_sent: int = 0
    batches_sent: int = 0
    bytes_original: int = 0
    bytes_compressed: int = 0
    errors: int = 0
    start_time: float = field(default_factory=time.time)


class CollectorV2:
    def __init__(self, num_batchers: int = NUM_BATCHERS):
        self.queue: asyncio.Queue = asyncio.Queue()
        self.stats = Stats()
        self.num_batchers = num_batchers
        self._running = False
        self._tasks: List[asyncio.Task] = []

    async def start(self):
        self._running = True
        for i in range(self.num_batchers):
            task = asyncio.create_task(self._batcher(i))
            self._tasks.append(task)
        print(f"✓ Collector v2 started ({self.num_batchers} parallel batchers)")

    async def stop(self):
        self._running = False
        for task in self._tasks:
            task.cancel()
        print("✓ Collector v2 stopped")

    async def submit(self, data: Dict) -> Dict:
        """Fire-and-forget: Just queue the message, don't wait for translation."""
        await self.queue.put(data)
        self.stats.messages_received += 1
        return {"queued": True, "queue_size": self.queue.qsize()}

    async def _batcher(self, batcher_id: int):
        """Background worker that batches and sends to translator."""
        async with httpx.AsyncClient(timeout=30.0) as client:
            batch = []
            while self._running:
                try:
                    msg = await asyncio.wait_for(
                        self.queue.get(),
                        timeout=BATCH_TIMEOUT_MS / 1000
                    )
                    batch.append(msg)

                    if len(batch) >= BATCH_SIZE:
                        await self._send_batch(client, batch)
                        batch = []

                except asyncio.TimeoutError:
                    if batch:
                        await self._send_batch(client, batch)
                        batch = []
                except asyncio.CancelledError:
                    if batch:
                        await self._send_batch(client, batch)
                    break

    async def _send_batch(self, client: httpx.AsyncClient, batch: List[Dict]):
        """Send a batch to the translator API."""
        try:
            response = await client.post(
                f"{TRANSLATOR_API}/translate/batch",
                json={"messages": batch},
                headers={"Authorization": f"Bearer {AUTH_TOKEN}"}
            )

            if response.status_code == 200:
                result = response.json()
                self.stats.messages_sent += len(batch)
                self.stats.batches_sent += 1
                self.stats.bytes_original += result.get('original_size', 0)
                self.stats.bytes_compressed += result.get('compressed_size', 0)
            else:
                self.stats.errors += 1

        except Exception as e:
            self.stats.errors += 1


# ═══════════════════════════════════════════════════════════════════════════════
# API
# ═══════════════════════════════════════════════════════════════════════════════

collector = CollectorV2()


@asynccontextmanager
async def lifespan(app: FastAPI):
    await collector.start()
    yield
    await collector.stop()


app = FastAPI(
    title="Emergent Language Collector v2",
    description="High-performance sidecar with parallel batchers",
    version="2.0.0",
    lifespan=lifespan
)


class MessageRequest(BaseModel):
    data: Dict


class BulkRequest(BaseModel):
    messages: List[Dict]


@app.get("/")
async def root():
    return {
        "service": "emergent-collector-v2",
        "version": "2.0.0",
        "description": "Fire-and-forget sidecar with parallel batchers",
        "config": {
            "batch_size": BATCH_SIZE,
            "num_batchers": collector.num_batchers,
            "timeout_ms": BATCH_TIMEOUT_MS
        }
    }


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "queue_size": collector.queue.qsize(),
        "messages_sent": collector.stats.messages_sent,
        "batches_sent": collector.stats.batches_sent
    }


@app.get("/stats")
async def stats():
    s = collector.stats
    uptime = time.time() - s.start_time
    return {
        "messages_received": s.messages_received,
        "messages_sent": s.messages_sent,
        "batches_sent": s.batches_sent,
        "queue_size": collector.queue.qsize(),
        "errors": s.errors,
        "compression_ratio": f"{(1 - s.bytes_compressed/s.bytes_original)*100:.1f}%" if s.bytes_original > 0 else "0%",
        "bytes_saved": s.bytes_original - s.bytes_compressed,
        "throughput": f"{s.messages_sent/uptime:.0f} msg/s" if uptime > 0 else "0",
        "uptime_seconds": uptime
    }


@app.post("/collect")
async def collect(request: MessageRequest):
    """
    Fire-and-forget: Queue message for batching.
    Returns immediately, doesn't wait for translation.
    """
    return await collector.submit(request.data)


@app.post("/collect/bulk")
async def collect_bulk(request: BulkRequest):
    """Queue multiple messages at once."""
    for msg in request.messages:
        await collector.submit(msg)
    return {"queued": len(request.messages), "queue_size": collector.queue.qsize()}


# ═══════════════════════════════════════════════════════════════════════════════
# CLI Entry Point
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    """CLI entry point for emergent-collector command."""
    import argparse
    parser = argparse.ArgumentParser(
        description="Emergent Language Collector - High-performance sidecar for AI agents",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  emergent-collector                    # Start on port 8080
  emergent-collector --port 9000        # Start on port 9000
  emergent-collector --batchers 8       # Use 8 parallel batchers

The collector receives messages from agents and batches them for optimal
throughput to the Translator API. Run this as a sidecar alongside your agents.
        """
    )
    parser.add_argument("--port", "-p", type=int, default=8080,
                        help="Port to listen on (default: 8080)")
    parser.add_argument("--batchers", "-b", type=int, default=4,
                        help="Number of parallel batchers (default: 4)")
    parser.add_argument("--host", type=str, default="0.0.0.0",
                        help="Host to bind to (default: 0.0.0.0)")
    parser.add_argument("--translator-url", type=str, default=TRANSLATOR_API,
                        help=f"Translator API URL (default: {TRANSLATOR_API})")
    args = parser.parse_args()

    collector.num_batchers = args.batchers

    print(f"""
╔═══════════════════════════════════════════════════════════════════╗
║           Emergent Language Collector                             ║
╠═══════════════════════════════════════════════════════════════════╣
║  Endpoints:                                                       ║
║    POST /collect      - Fire-and-forget single message            ║
║    POST /collect/bulk - Fire-and-forget multiple messages         ║
║    GET  /stats        - Throughput and compression stats          ║
║    GET  /health       - Health check                              ║
╠═══════════════════════════════════════════════════════════════════╣
║  Config:                                                          ║
║    Port:           {args.port:<6}                                        ║
║    Batch size:     {BATCH_SIZE:<6}                                        ║
║    Batchers:       {args.batchers:<6}                                        ║
║    Timeout:        {BATCH_TIMEOUT_MS}ms                                           ║
╚═══════════════════════════════════════════════════════════════════╝
""")

    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()
