from __future__ import annotations

from typing import Optional, Self

from celery import Celery
from celery.canvas import Signature

from minix.core.scheduler.task import PeriodicTask, Task
from minix.core.scheduler.workflow import Workflow
from minix.core.scheduler.task.workflow_tasks import (
    InitContextTask,
    ExecuteNodeTask,
    ExtractOneTask,
    ExtractSinksTask,
)


class SchedulerConfig:
    def __init__(self):
        self.broker_url = ''
        self.result_backend = ''
        self.task_serializer = 'json'
        self.result_serializer = 'json'
        self.accept_content = ['json']
        self.timezone = 'GMT'

    def set_broker_url(self, url) -> Self:
        self.broker_url = url
        return self

    def set_result_backend(self, url) -> Self:
        self.result_backend = url
        return self

    def set_task_serializer(self, serializer) -> Self:
        self.task_serializer = serializer
        return self

    def set_result_serializer(self, serializer) -> Self:
        self.result_serializer = serializer
        return self

    def set_accept_content(self, content) -> Self:
        self.accept_content = content
        return self

    def set_timezone(self, timezone) -> Self:
        self.timezone = timezone
        return self

    def get_broker_url(self):
        return self.broker_url

    def get_result_backend(self):
        return self.result_backend

    def get_task_serializer(self):
        return self.task_serializer

    def get_result_serializer(self):
        return self.result_serializer

    def get_accept_content(self):
        return self.accept_content

    def get_timezone(self):
        return self.timezone


class Scheduler(Celery):

    def __init__(self, config: SchedulerConfig):
        super().__init__('scheduler')
        self.celery = Celery(
            "task_scheduler",
            broker=config.get_broker_url(),
            backend=config.get_result_backend(),
            timezone=config.get_timezone(),
            broker_connection_retry_on_startup=False,
            task_track_started=True,
            result_extended=True,
        )
        self.celery.beat_scheduler = 'redbeat.RedBeatScheduler'
        self.celery.redbeat_redis_url = config.get_broker_url()
        self.celery.conf.task_serializer = config.get_task_serializer()
        self.celery.conf.result_serializer = config.get_result_serializer()
        self.celery.conf.accept_content = config.get_accept_content()
        self.celery.conf.timezone = config.get_timezone()

        # Register workflow helper tasks
        self.celery.register_task(InitContextTask())
        self.celery.register_task(ExecuteNodeTask())
        self.celery.register_task(ExtractOneTask())
        self.celery.register_task(ExtractSinksTask())

    def get_app(self):
        return self.celery

    def run_task(self, task: Task):
        # Return AsyncResult so caller can track id
        return self.get_app().send_task(
            task.get_name(),
            args=task.get_args(),
            kwargs=task.get_kwargs()
        )

    def run_workflow(
            self,
            wf: Workflow,
            *,
            target_node_id: Optional[str] = None,
            link_error: Optional[Signature] = None,
    ):
        # NOTE:
        # This implementation uses a context-carrying chain and does NOT require chords,
        # so we no longer block execution if result_backend is empty.
        canvas = wf.to_canvas(app=self.get_app(), target_node_id=target_node_id)
        return canvas.apply_async(link_error=link_error) if link_error else canvas.apply_async()

    def register_async_task(self, task: Task):
        self.get_app().register_task(task)
        return self

    # def register_periodic_task(self, task: PeriodicTask):
    #     registering_task = task
    #     self.get_app().register_task(registering_task)
    #     self.get_app().conf.beat_schedule = {
    #         registering_task.get_name(): {
    #             'task': registering_task.get_name(),
    #             'schedule': registering_task.get_schedule(),
    #             'args': registering_task.get_args(),
    #             'kwargs': registering_task.get_kwargs()
    #         }
    #     }
    #     return self

    def register_periodic_task(self, task: PeriodicTask):
        app = self.get_app()
        app.register_task(task)

        # merge instead of replace
        beat = getattr(app.conf, "beat_schedule", {}) or {}
        beat[task.get_name()] = {
            "task": task.get_name(),
            "schedule": task.get_schedule(),
            "args": task.get_args(),
            "kwargs": task.get_kwargs(),
        }
        app.conf.beat_schedule = beat
        return self
