"""
rabbit_integration.py
---------------------
Drop-in RabbitMQ integration for the transcription → summarization → SEO pipeline.

Architecture: three queues, three consumers, each stage is an independent worker.
Stages are the same pure functions from pipeline.py — the queue layer is just transport.

Queue topology:
    transcription_jobs      → TranscriptionJobMessage
    summarization_jobs      → TranscriptionDoneMessage
    seo_jobs                → SummarizationDoneMessage
    pipeline_results        → PipelineDoneMessage

Each consumer:
  1. Deserializes the message into a typed schema
  2. Calls the stage function
  3. Publishes the result to the next queue
  4. Acks the original message (only on success)

Environment variables required:
    RABBITMQ_URL    — amqp connection string, e.g. amqp://user:pass@localhost:5672/vhost

Usage — publish a job:
    from .rabbit_integration import publish_transcription_job
    from .schemas import TranscriptionInput, MediaMeta, TranscriptionJobMessage
    import uuid

    publish_transcription_job(TranscriptionJobMessage(
        job_id=str(uuid.uuid4()),
        inp=TranscriptionInput(audio_path="/data/audio.wav"),
        meta=MediaMeta(domain="https://typicallyoutliers.com"),
    ))

Usage — run workers (one per stage, each in its own process/container):
    python -m yourpackage.rabbit_integration transcription
    python -m yourpackage.rabbit_integration summarization
    python -m yourpackage.rabbit_integration seo
"""
from __future__ import annotations

import json
import logging
import os
import sys
from typing import Callable

logger = logging.getLogger(__name__)

# ── Queue names ───────────────────────────────────────────────────────────────
# Centralised here so they never drift between publisher and consumer.

QUEUE_TRANSCRIPTION_JOBS  = "transcription_jobs"
QUEUE_SUMMARIZATION_JOBS  = "summarization_jobs"
QUEUE_SEO_JOBS            = "seo_jobs"
QUEUE_PIPELINE_RESULTS    = "pipeline_results"
QUEUE_DEAD_LETTER         = "pipeline_dead_letter"

ALL_QUEUES = [
    QUEUE_TRANSCRIPTION_JOBS,
    QUEUE_SUMMARIZATION_JOBS,
    QUEUE_SEO_JOBS,
    QUEUE_PIPELINE_RESULTS,
    QUEUE_DEAD_LETTER,
]


# ── Connection ────────────────────────────────────────────────────────────────

def _get_rabbit_url() -> str:
    url = os.environ.get("RABBITMQ_URL", "").strip()
    if not url:
        raise EnvironmentError(
            "[RabbitMQ] RABBITMQ_URL is not set. "
            "Set it to your amqp connection string before starting a worker."
        )
    return url


def _get_channel():
    """Return a pika channel with all pipeline queues declared."""
    import pika
    params  = pika.URLParameters(_get_rabbit_url())
    conn    = pika.BlockingConnection(params)
    channel = conn.channel()
    for queue in ALL_QUEUES:
        channel.queue_declare(queue=queue, durable=True)
    return channel


# ── Publisher ─────────────────────────────────────────────────────────────────

def _publish(channel, queue: str, payload: dict) -> None:
    import pika
    channel.basic_publish(
        exchange="",
        routing_key=queue,
        body=json.dumps(payload),
        properties=pika.BasicProperties(
            delivery_mode=2,        # persistent
            content_type="application/json",
        ),
    )
    logger.debug("[RabbitMQ] Published to %s", queue)


def publish_transcription_job(msg: "TranscriptionJobMessage") -> None:  # noqa: F821
    """Enqueue a new pipeline job. Called by API routes or external triggers."""
    channel = _get_channel()
    _publish(channel, QUEUE_TRANSCRIPTION_JOBS, msg.model_dump())


# ── Consumer base ─────────────────────────────────────────────────────────────

def _run_consumer(
    queue:       str,
    handler:     Callable[[dict], None],
    prefetch:    int = 1,
) -> None:
    """
    Blocking consumer loop. Acks only on successful handler return.
    Sends to dead-letter queue on unhandled exception.
    """
    import pika

    channel = _get_channel()
    channel.basic_qos(prefetch_count=prefetch)

    def on_message(ch, method, properties, body):
        try:
            data = json.loads(body)
            handler(data)
            ch.basic_ack(delivery_tag=method.delivery_tag)
        except Exception as exc:
            logger.exception("[RabbitMQ] Handler failed for queue=%s: %s", queue, exc)
            _publish(ch, QUEUE_DEAD_LETTER, {
                "queue":  queue,
                "error":  str(exc),
                "body":   body.decode("utf-8", errors="replace"),
            })
            ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)

    channel.basic_consume(queue=queue, on_message_callback=on_message)
    logger.info("[RabbitMQ] Worker listening on '%s' ...", queue)
    channel.start_consuming()


# ── Stage consumers ───────────────────────────────────────────────────────────

def _handle_transcription_job(data: dict) -> None:
    from .schemas import TranscriptionJobMessage, TranscriptionDoneMessage
    from .pipeline import stage_transcribe

    msg           = TranscriptionJobMessage(**data)
    transcription = stage_transcribe(msg.inp)

    done = TranscriptionDoneMessage(
        job_id=msg.job_id,
        transcription=transcription,
        meta=msg.meta,
        task=msg.task,
        keyword_count=msg.keyword_count,
        title=msg.title,
    )
    channel = _get_channel()
    _publish(channel, QUEUE_SUMMARIZATION_JOBS, done.model_dump())
    logger.info("[Worker:Transcription] job_id=%s done.", msg.job_id)


def _handle_summarization_job(data: dict) -> None:
    from .schemas import TranscriptionDoneMessage, SummarizationDoneMessage
    from .pipeline import stage_summarize

    msg           = TranscriptionDoneMessage(**data)
    summarization = stage_summarize(
        transcription=msg.transcription,
        task=msg.task,
        keyword_count=msg.keyword_count,
    )

    done = SummarizationDoneMessage(
        job_id=msg.job_id,
        summarization=summarization,
        transcription=msg.transcription,
        meta=msg.meta,
        title=msg.title,
    )
    channel = _get_channel()
    _publish(channel, QUEUE_SEO_JOBS, done.model_dump())
    logger.info("[Worker:Summarization] job_id=%s done.", msg.job_id)


def _handle_seo_job(data: dict) -> None:
    from .schemas import SummarizationDoneMessage, PipelineDoneMessage, PipelineResult
    from .pipeline import stage_seo

    msg = SummarizationDoneMessage(**data)
    seo = stage_seo(
        summary_result=msg.summarization,
        transcription=msg.transcription,
        meta=msg.meta,
        title=msg.title,
    )

    result = PipelineResult(
        video_id=msg.transcription.video_id,
        transcription=msg.transcription,
        summarization=msg.summarization,
        seo=seo,
    )
    done = PipelineDoneMessage(job_id=msg.job_id, result=result)
    channel = _get_channel()
    _publish(channel, QUEUE_PIPELINE_RESULTS, done.model_dump())
    logger.info("[Worker:SEO] job_id=%s done. title='%s'", msg.job_id, seo.seo_title)


# ── Worker entry points ───────────────────────────────────────────────────────

def run_transcription_worker(prefetch: int = 1) -> None:
    _run_consumer(QUEUE_TRANSCRIPTION_JOBS, _handle_transcription_job, prefetch)


def run_summarization_worker(prefetch: int = 1) -> None:
    _run_consumer(QUEUE_SUMMARIZATION_JOBS, _handle_summarization_job, prefetch)


def run_seo_worker(prefetch: int = 1) -> None:
    _run_consumer(QUEUE_SEO_JOBS, _handle_seo_job, prefetch)


# ── CLI entry point ───────────────────────────────────────────────────────────
# python -m yourpackage.rabbit_integration <stage>

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    _WORKERS = {
        "transcription": run_transcription_worker,
        "summarization": run_summarization_worker,
        "seo":           run_seo_worker,
    }

    stage = sys.argv[1] if len(sys.argv) > 1 else ""
    if stage not in _WORKERS:
        print(f"Usage: python -m ... rabbit_integration <stage>")
        print(f"  stages: {', '.join(_WORKERS)}")
        sys.exit(1)

    _WORKERS[stage]()
