import time
from typing import Any

from celery import signals

from src.lib.tracing.tracker import tracer

_task_starts: dict[str, float] = {}
_task_context_tokens: dict[str, Any] = {}


def _task_id(task_id: str | None, task=None) -> str:
    if task_id:
        return task_id
    request = getattr(task, "request", None)
    return getattr(request, "id", None) or "unknown"


def _task_name(task=None, sender=None) -> str:
    return getattr(task, "name", None) or getattr(sender, "name", None) or "unknown"


def _tenant_id(args=None, kwargs=None) -> str | None:
    args = args or ()
    kwargs = kwargs or {}
    candidates = list(args)
    candidates.extend(kwargs.values())
    for candidate in candidates:
        if isinstance(candidate, dict):
            brain_id = candidate.get("brain_id")
            if brain_id:
                return str(brain_id)
    brain_id = kwargs.get("brain_id")
    return str(brain_id) if brain_id else None


def _trace_id(task_id: str, kwargs=None) -> str:
    kwargs = kwargs or {}
    trace_id = kwargs.get("trace_id")
    return str(trace_id) if trace_id else task_id


def _metadata(task_id: str, task_name: str, args=None, kwargs=None) -> dict[str, Any]:
    return {
        "task_id": task_id,
        "task_name": task_name,
        "args_count": len(args or ()),
        "kwargs_keys": sorted((kwargs or {}).keys()),
    }


def install_celery_tracing(_app=None, service_name: str = "brainapi-worker") -> None:
    signals.task_prerun.connect(
        _trace_task_prerun,
        weak=False,
        dispatch_uid=f"{service_name}.trace.task_prerun",
    )
    signals.task_postrun.connect(
        _trace_task_postrun,
        weak=False,
        dispatch_uid=f"{service_name}.trace.task_postrun",
    )
    signals.task_failure.connect(
        _trace_task_failure,
        weak=False,
        dispatch_uid=f"{service_name}.trace.task_failure",
    )


def _trace_task_prerun(sender=None, task_id=None, task=None, args=None, kwargs=None, **_):
    resolved_task_id = _task_id(task_id, task)
    resolved_task_name = _task_name(task, sender)
    _task_starts[resolved_task_id] = time.perf_counter()
    _task_context_tokens[resolved_task_id] = tracer.set_context(
        trace_id=_trace_id(resolved_task_id, kwargs),
        tenant_id=_tenant_id(args, kwargs),
    )
    tracer.publish(
        "latency",
        f"{resolved_task_name}.started",
        service="brainapi-worker",
        operation=resolved_task_name,
        metadata=_metadata(resolved_task_id, resolved_task_name, args, kwargs),
    )


def _trace_task_postrun(sender=None, task_id=None, task=None, args=None, kwargs=None, retval=None, state=None, **_):
    resolved_task_id = _task_id(task_id, task)
    resolved_task_name = _task_name(task, sender)
    started_at = _task_starts.pop(resolved_task_id, None)
    duration_ms = (time.perf_counter() - started_at) * 1000 if started_at else None
    metadata = {
        **_metadata(resolved_task_id, resolved_task_name, args, kwargs),
        "state": state,
        "returned": retval is not None,
    }
    if duration_ms is not None and duration_ms >= tracer.default_sla_ms:
        tracer.publish(
            "sla_breach",
            f"{resolved_task_name}.sla_breach",
            service="brainapi-worker",
            operation=resolved_task_name,
            duration_ms=duration_ms,
            threshold_ms=tracer.default_sla_ms,
            metadata=metadata,
        )
    token = _task_context_tokens.pop(resolved_task_id, None)
    if token is not None:
        tracer.reset_context(token)


def _trace_task_failure(sender=None, task_id=None, exception=None, args=None, kwargs=None, einfo=None, **_):
    resolved_task_id = _task_id(task_id, sender)
    resolved_task_name = _task_name(sender=sender)
    started_at = _task_starts.get(resolved_task_id)
    duration_ms = (time.perf_counter() - started_at) * 1000 if started_at else None
    metadata = _metadata(resolved_task_id, resolved_task_name, args, kwargs)
    if einfo is not None:
        metadata["traceback"] = str(einfo)
    tracer.exception(
        f"{resolved_task_name}.failed",
        exception or RuntimeError("Celery task failed"),
        service="brainapi-worker",
        operation=resolved_task_name,
        duration_ms=duration_ms,
        metadata=metadata,
    )