from pathlib import Path
from typing import Callable, Optional

import pytest
from celery import Celery
from celery.contrib.testing.worker import TestWorkController
from pydantic import BaseModel

from acoupi.programs.core.base import AcoupiProgram
from acoupi.programs.core.workers import AcoupiWorker, WorkerConfig

WaitForCondition = Callable[..., None]


class Config(BaseModel):
    path: Path
    message: str


class ProgramWithCallbacks(AcoupiProgram):
    config: Config

    def setup(self, config: Config):
        def task_1() -> Path:
            return config.path

        def task_2(path: Optional[Path]) -> None:
            if path is None:
                return

            path.write_text(config.message)

        self.add_task(task_1, callbacks=[task_2])  # ty: ignore


@pytest.mark.timeout(10)
@pytest.mark.usefixtures("celery_app")
@pytest.mark.usefixtures("celery_worker")
def test_does_run_callbacks(
    tmp_path: Path,
    celery_app: Celery,
    celery_worker: TestWorkController,
    wait_for_condition: WaitForCondition,
):
    path = tmp_path / "test.txt"
    message = "Hello, world!"
    config = Config(path=path, message=message)

    assert not path.exists()

    program = ProgramWithCallbacks(
        program_config=config,
        app=celery_app,
    )
    assert celery_app.conf["accept_content"] == ["pickle"]

    assert "task_1" in program.tasks
    assert "task_2" in program.tasks
    assert "task_1" in program.app.tasks
    assert "task_2" in program.app.tasks

    celery_worker.reload()
    celery_worker.ensure_started()

    output = program.tasks["task_1"].delay()
    output.get(timeout=3)

    wait_for_condition(path.exists, timeout=3)
    wait_for_condition(lambda: path.read_text() == message, timeout=3)

    assert path.exists()
    assert path.read_text() == message


class ProgramWithQueuedCallbacks(AcoupiProgram):
    config: Config

    worker_config = WorkerConfig(
        workers=[
            AcoupiWorker(
                queues=["special"],
                name="special",
                concurrency=1,
            ),
            AcoupiWorker(
                queues=["default"],
                name="default",
            ),
        ]
    )

    def setup(self, config: Config):
        def task_1() -> Path:
            return config.path

        def task_2(path: Optional[Path]):
            if path is None:
                return

            path.write_text(config.message)

        self.add_task(task_1, callbacks=[task_2], queue="special")  # ty: ignore


@pytest.mark.timeout(10)
def test_does_run_callbacks_in_other_queues(
    tmp_path: Path,
    celery_app: Celery,
    celery_worker: TestWorkController,
    wait_for_condition: WaitForCondition,
):
    path = tmp_path / "test.txt"
    message = "Hello, world!"
    config = Config(path=path, message=message)

    assert not path.exists()

    program = ProgramWithQueuedCallbacks(
        program_config=config,
        app=celery_app,
    )

    celery_worker.reload()
    celery_worker.ensure_started()

    assert "task_1" in program.tasks
    assert "task_2" in program.tasks

    output = program.tasks["task_1"].delay()
    output.get(timeout=3)

    wait_for_condition(path.exists)
    wait_for_condition(lambda: path.read_text() == message, timeout=3)

    assert path.exists()
    assert path.read_text() == message
