import os
import time

from celery import Celery
from celery.signals import task_postrun, task_prerun

# Set the default Django settings module for the 'celery' program.
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "core.settings")

app = Celery("core")

# Using a string here means the worker doesn't have to serialize
# the configuration object to child processes.
# - namespace='CELERY' means all celery-related configuration keys
#   should have a `CELERY_` prefix.
app.config_from_object("django.conf:settings", namespace="CELERY")

# Load task modules from all registered Django apps.
app.autodiscover_tasks()


# log2 ambient context per task; tokens restore the caller's context, which
# matters in sync mode (CELERY_TASK_ALWAYS_EAGER) where tasks run inside a request
_log_tokens = {}
_task_t0 = {}


@task_prerun.connect
def _log_task_prerun(task_id=None, task=None, **kwargs):
    """Bind log2 ambient context (channel=celery) for one task run."""
    from app_home.log import log2

    _task_t0[task_id] = time.monotonic()
    _log_tokens[task_id] = log2.reset(channel="celery", origin=getattr(task, "name", ""),
                                      actor="system", actor_kind="system",
                                      request_id=(task_id or "")[:16])  # fmt: skip
    aaa = _task_aaa(kwargs.get("args"), kwargs.get("kwargs"))
    if aaa:
        log2.bind_aaa(aaa)


def _task_aaa(args, kwargs):
    """Return the requester's aaa dict a task was given (aaa= kwarg, or a params dict), or None."""
    candidates = [(kwargs or {}).get("aaa")]
    candidates += [a.get("aaa") for a in (args or ()) if isinstance(a, dict)]
    for aaa in candidates:
        if isinstance(aaa, dict) and aaa.get("username"):
            return aaa
    return None


@task_postrun.connect
def _log_task_postrun(task_id=None, task=None, **kwargs):
    """Record the task duration (telemetry task.duration), flush the telemetry counters the
    task raised, and restore the log2 ambient context that was active before the task."""
    from app_data.telemetry import telemetry, telemetry_flush
    from app_home.log import log2

    t0 = _task_t0.pop(task_id, None)
    if t0 is not None:
        name = getattr(task, "name", "").rsplit(".", 1)[-1]
        telemetry("task.duration", id=name, value=(time.monotonic() - t0) * 1000)

    telemetry_flush()

    token = _log_tokens.pop(task_id, None)
    if token is not None:
        log2.restore(token)


# @app.task
# def cron1(arg):
#     print("cron " + arg)


# @app.task(bind=True)
# def debug_task(self):
#     print(f'Request: {self.request!r}')