"""Celery patching for task queue capture.

Captures Celery task invocations and results for replay.

Supported operations:
- apply_async (async task invocation)
- delay (convenience method)
- apply (sync task invocation)
- Task results

Usage:
    from autodebug_probe.celery_patch import install_celery
    install_celery()

    # Now all Celery task calls are captured
    result = my_task.delay(arg1, arg2)
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from autodebug_probe.recorder import get_recorder


def _serialize_arg(arg: Any) -> Any:
    """Serialize task argument for recording."""
    if arg is None:
        return None
    if isinstance(arg, (str, int, float, bool)):
        return arg
    if isinstance(arg, (list, tuple)):
        return [_serialize_arg(a) for a in arg]
    if isinstance(arg, dict):
        return {k: _serialize_arg(v) for k, v in arg.items()}
    if isinstance(arg, datetime):
        return arg.isoformat()
    if hasattr(arg, "__dict__"):
        return {"_type": type(arg).__name__, "_data": str(arg)}
    return str(arg)


def install_celery() -> None:
    """Install Celery patches for capture."""
    try:
        import celery
        from celery import Task
    except ImportError:
        return

    # Patch Task.apply_async
    original_apply_async = Task.apply_async

    def patched_apply_async(
        self,
        args=None,
        kwargs=None,
        task_id=None,
        producer=None,
        link=None,
        link_error=None,
        shadow=None,
        **options
    ):
        result = original_apply_async(
            self, args, kwargs, task_id, producer, link, link_error, shadow, **options
        )

        get_recorder().record_event({
            "type": "CELERY_TASK",
            "operation": "apply_async",
            "task_name": self.name,
            "task_id": str(result.id) if result else None,
            "args": _serialize_arg(args),
            "kwargs": _serialize_arg(kwargs),
            "options": {
                k: _serialize_arg(v)
                for k, v in options.items()
                if k not in ("headers", "argsrepr", "kwargsrepr")
            },
        })

        return result

    Task.apply_async = patched_apply_async

    # Patch Task.delay (which calls apply_async)
    original_delay = Task.delay

    def patched_delay(self, *args, **kwargs):
        # delay() just calls apply_async, which is now patched
        return self.apply_async(args=args, kwargs=kwargs)

    Task.delay = patched_delay

    # Patch Task.apply (sync execution)
    original_apply = Task.apply

    def patched_apply(self, args=None, kwargs=None, **options):
        result = original_apply(self, args, kwargs, **options)

        get_recorder().record_event({
            "type": "CELERY_TASK",
            "operation": "apply",
            "task_name": self.name,
            "task_id": str(result.id) if result else None,
            "args": _serialize_arg(args),
            "kwargs": _serialize_arg(kwargs),
            "sync": True,
            "result": _serialize_arg(result.result) if result and result.ready() else None,
            "status": result.status if result else None,
        })

        return result

    Task.apply = patched_apply

    # Patch AsyncResult.get to capture result retrieval
    try:
        from celery.result import AsyncResult

        original_get = AsyncResult.get

        def patched_get(self, timeout=None, propagate=True, interval=0.5, **kwargs):
            result = original_get(self, timeout, propagate, interval, **kwargs)

            get_recorder().record_event({
                "type": "CELERY_RESULT",
                "operation": "get",
                "task_id": str(self.id),
                "result": _serialize_arg(result),
                "status": self.status,
            })

            return result

        AsyncResult.get = patched_get

    except ImportError:
        pass
