"""Example: Job and RPC services + clients using nats_common.

Requires a running NATS server with JetStream enabled::

    nats-server --jetstream

This script starts the server-side services *and* immediately exercises
them with the matching client classes to demonstrate the full round-trip.
"""

from __future__ import annotations

import asyncio
import logging
import signal

import nats
from pydantic import BaseModel

from nats_common.services import (
    JobStatus,
    JobStatusUpdate,
    NatsJob,
    NatsJobClient,
    NatsJobService,
    NatsRPC,
    NatsRPCClient,
    NatsRPCService,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ─── Models ───────────────────────────────────────────────────────────────────


class TranscribeRequest(BaseModel):
    audio_url: str
    language: str = "en"


class TranscribeChunk(BaseModel):
    text: str
    offset_ms: int


class TranscribeResult(BaseModel):
    full_text: str


class EchoRequest(BaseModel):
    message: str


class EchoResponse(BaseModel):
    echoed: str


# ─── Job handler (server-side) ────────────────────────────────────────────────


class TranscribeJob(NatsJob[TranscribeRequest, TranscribeResult, TranscribeChunk]):
    """Simulates a long-running transcription job with progress updates."""

    async def handle(self) -> None:
        req: TranscribeRequest | None = self.request.data  # type: ignore[assignment]
        if req is None:
            return

        logger.info("Transcribing %s (lang=%s)", req.audio_url, req.language)

        chunks = ["Hello", "world", "from", "NATS"]
        for i, word in enumerate(chunks):
            if self.cancel_requested:
                logger.info("Job %s cancelled", self.job_id)
                return

            self.set_local_status(
                JobStatus.RUNNING,
                progress=(i + 1) / len(chunks),
                message=f"Processing chunk {i + 1}/{len(chunks)}",
            )

            await self.publish_intermediate_result(
                TranscribeChunk(text=word, offset_ms=i * 500),
            )
            await asyncio.sleep(0.5)

        await self.publish_final_result(
            TranscribeResult(full_text=" ".join(chunks)),
        )


# ─── RPC handler (server-side) ────────────────────────────────────────────────


class EchoRpc(NatsRPC[EchoRequest, EchoResponse]):
    """Simple echo RPC that returns the received message."""

    async def handle(self) -> EchoResponse:
        req: EchoRequest | None = self.request.data  # type: ignore[assignment]
        return EchoResponse(echoed=req.message if req else "")


# ─── Server-only demo ────────────────────────────────────────────────────────


async def run_server_only() -> None:
    """Start services and wait for Ctrl+C.  Use the nats CLI to test."""
    quit_event = asyncio.Event()
    loop = asyncio.get_event_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, lambda *_: quit_event.set())

    nats_client = await nats.connect("nats://localhost:4222")

    job_svc = NatsJobService(
        nats_client=nats_client,
        subject="demo.transcribe",
        job_cls=TranscribeJob,
    )
    rpc_svc = NatsRPCService(
        nats_client=nats_client,
        subject="demo.echo",
        rpc_cls=EchoRpc,
        request_model=EchoRequest,
        response_model=EchoResponse,
    )

    await job_svc.start()
    await rpc_svc.start()
    logger.info("Services started — press Ctrl+C to stop")

    await quit_event.wait()

    await job_svc.stop()
    await rpc_svc.stop()
    await nats_client.close()


# ─── Full round-trip demo (server + client) ───────────────────────────────────


async def run_round_trip() -> None:
    """Start services, exercise them with the client classes, then shut down."""
    nats_client = await nats.connect("nats://localhost:4222")

    # --- Start server-side ---
    job_svc = NatsJobService(
        nats_client=nats_client,
        subject="demo.transcribe",
        job_cls=TranscribeJob,
    )
    rpc_svc = NatsRPCService(
        nats_client=nats_client,
        subject="demo.echo",
        rpc_cls=EchoRpc,
        request_model=EchoRequest,
        response_model=EchoResponse,
    )
    await job_svc.start()
    await rpc_svc.start()
    logger.info("Server-side services started")

    # --- RPC client demo ---
    rpc_client = NatsRPCClient(
        nats_client=nats_client,
        subject="demo.echo",
        response_model=EchoResponse,
    )
    resp = await rpc_client.call(data=EchoRequest(message="hello world"))
    logger.info("RPC response: %s (error=%s)", resp.data, resp.error)

    # --- Job client demo ---
    job_client = NatsJobClient[TranscribeRequest, TranscribeChunk, TranscribeResult](
        nats_client=nats_client,
        subject="demo.transcribe",
    )

    async def _on_status(update: JobStatusUpdate) -> None:
        logger.info(
            "  status: %s  progress=%s  msg=%s",
            update.status.value,
            update.progress,
            update.message,
        )

    async def _on_chunk(chunk: TranscribeChunk) -> None:
        logger.info("  intermediate chunk: %r @ %d ms", chunk.text, chunk.offset_ms)

    # Subscribe first, then submit — so we never miss a message.
    completion = asyncio.create_task(
        job_client.wait_for_completion(
            "test-job-1",
            timeout_seconds=30,
            on_status=_on_status,
            on_intermediate=_on_chunk,
        ),
    )

    # Give subscriptions a moment to be set up.
    await asyncio.sleep(0.1)

    await job_client.submit(
        job_id="test-job-1",
        data=TranscribeRequest(audio_url="https://example.com/audio.wav"),
    )
    logger.info("Job submitted")

    result = await completion
    logger.info(
        "Job completed: succeeded=%s  final_result=%s",
        result.succeeded,
        result.final_result,
    )

    # --- Tear down ---
    await job_svc.stop()
    await rpc_svc.stop()
    await nats_client.close()


if __name__ == "__main__":
    import sys

    mode = sys.argv[1] if len(sys.argv) > 1 else "round-trip"
    if mode == "server":
        asyncio.run(run_server_only())
    else:
        asyncio.run(run_round_trip())
