"""``@task`` 装饰器与 ``TaskContext`` 的能力测试。

三块：

1. **``@task``**——元数据从装饰器来；不产生全局状态（导入无副作用）；
2. **进度**——进快照、进存储、**不碰版本号**；
3. **上游数据**——管道节点之间真正能传值，且失败的上游不传 ``None``。
"""

from __future__ import annotations

import logging
import time
from typing import Any

import pytest

from chongming_task.controller import Controller, TaskRegistry
from chongming_task.decorators import get_task_spec, task
from chongming_task.errors import ConfigurationError
from chongming_task.pipeline import Pipeline
from chongming_task.policies import RetryPolicy, Timeout
from chongming_task.states import TaskStatus
from chongming_task.stores.memory import InMemoryTaskStore

_WINDOW = 15.0


# ---------------------------------------------------------------------------
# 1. @task
# ---------------------------------------------------------------------------


@task
def bare(ctx: Any) -> str:
    """裸装饰器。"""
    del ctx
    return "bare"


@task(tags={"io", "nightly"}, retry=RetryPolicy(max_attempts=4), timeout=Timeout(per_attempt=9.0))
def decorated(ctx: Any) -> str:
    """带元数据的任务。

    第二行不该出现在描述里。
    """
    del ctx
    return "decorated"


@task(name="renamed")
def original_name(ctx: Any) -> str:
    del ctx
    return "renamed"


def undecorated(ctx: Any) -> str:
    del ctx
    return "nope"


def test_bare_decorator_works() -> None:
    """``@task`` 不带括号也太自然，不支持会让人反复踩空。"""
    spec = get_task_spec(bare)
    assert spec is not None
    assert spec.name == "bare"
    assert spec.description == "裸装饰器。"


def test_decorator_captures_metadata() -> None:
    spec = get_task_spec(decorated)
    assert spec is not None
    assert spec.tags == frozenset({"io", "nightly"})
    assert spec.retry is not None
    assert spec.retry.max_attempts == 4
    assert spec.timeout is not None
    assert spec.timeout.per_attempt == 9.0


def test_description_defaults_to_the_first_docstring_line() -> None:
    spec = get_task_spec(decorated)
    assert spec is not None
    assert spec.description == "带元数据的任务。", "只取第一行"
    assert "第二行" not in spec.description


def test_name_override() -> None:
    spec = get_task_spec(original_name)
    assert spec is not None
    assert spec.name == "renamed"


def test_undecorated_function_has_no_spec() -> None:
    assert get_task_spec(undecorated) is None


def test_decorating_does_not_create_global_state() -> None:
    """规格只挂在函数对象上。

    全局注册表意味着 **import 即注册**，而 import 可能发生在任何地方（包括测试收集
    阶段），那种隐式行为极难排查。
    """
    registry = TaskRegistry()
    assert registry.names() == ()
    assert get_task_spec(bare) is not None, "规格依然取得到，但没有登记进任何注册表"


def test_registry_registers_decorated_functions() -> None:
    registry = TaskRegistry()
    registration = registry.register_decorated(decorated)
    assert registration.name == "decorated"
    assert registration.tags == frozenset({"io", "nightly"})
    assert registration.retry.max_attempts == 4


def test_registry_rejects_an_undecorated_function() -> None:
    registry = TaskRegistry()
    with pytest.raises(ConfigurationError, match="没有用 @task 装饰过"):
        registry.register_decorated(undecorated)


def test_register_module_picks_up_every_decorated_function() -> None:
    import sys

    module = sys.modules[__name__]
    registry = TaskRegistry()
    registered = registry.register_module(module)

    names = {item.name for item in registered}
    assert {"bare", "decorated", "renamed"} <= names
    assert "undecorated" not in names


def test_register_module_rejects_a_module_with_nothing_decorated() -> None:
    import json

    registry = TaskRegistry()
    with pytest.raises(ConfigurationError, match="没有用 @task 装饰过的函数"):
        registry.register_module(json)


def test_controller_registers_a_module_end_to_end() -> None:
    import sys

    controller = Controller(max_parallel=2)
    controller.register_module(sys.modules[__name__])
    run = controller.submit("bare")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.ok
    assert outcome.value == "bare"
    controller.close(timeout=_WINDOW)


def test_decorated_retry_actually_applies() -> None:
    """装饰器上写的重试次数必须真的生效，而不只是躺在元数据里。"""
    store = InMemoryTaskStore()
    controller = Controller(store=store, max_parallel=2)
    registry = TaskRegistry()
    registry.register(
        "flaky",
        _flaky,
        retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()),
    )
    controller.register("flaky", _flaky, retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()))
    run = controller.submit("flaky")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.ok, "第三次应成功"
    snapshot = store.get(run.snapshot.task_id)
    assert snapshot is not None
    assert snapshot.attempt == 3
    controller.close(timeout=_WINDOW)


# ---------------------------------------------------------------------------
# 2. TaskContext：进度与日志
# ---------------------------------------------------------------------------


def test_progress_reaches_the_snapshot_without_bumping_the_version() -> None:
    """进度若参与乐观并发，会与状态写入持续冲突。"""
    controller = Controller(max_parallel=2)
    controller.register("reporter", _reporter)
    run = controller.submit("reporter")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.ok

    snapshot = outcome.snapshot
    assert snapshot.progress is not None
    assert snapshot.progress.done == 3
    assert snapshot.progress.total == 3
    assert snapshot.progress.unit == "rows"
    # 状态迁移只有 4 次（queued/starting/running/succeeded），进度不该挤进去。
    assert snapshot.status_version == 4


def test_progress_is_persisted_and_visible_in_queries() -> None:
    """进度要经 touch 旁路落到存储，否则看板的 PROGRESS 列永远是空的。"""
    store = InMemoryTaskStore()
    controller = Controller(store=store, max_parallel=2, command_poll_interval=0.05)
    controller.register("slow_reporter", _slow_reporter)
    run = controller.submit("slow_reporter")

    assert _wait_until(
        lambda: (store.get(run.snapshot.task_id) or _stub()).progress is not None,
        timeout=_WINDOW,
    )
    persisted = store.get(run.snapshot.task_id)
    assert persisted is not None
    assert persisted.progress is not None
    assert persisted.progress.done >= 1

    controller.stop(run.snapshot.task_id)
    controller.wait_all(timeout=_WINDOW)
    controller.close(timeout=_WINDOW)


def test_progress_on_an_unsupported_context_is_silently_dropped() -> None:
    """没有回写通道时静默丢弃，而不是抛异常。

    直接构造上下文（不经过执行器）是测试与工具里的常见用法，不该因此报错。
    """
    from chongming_task.runtime.cancellation import CancellationToken
    from chongming_task.runtime.context import RuntimeContext
    from chongming_task.runtime.heartbeat import Heartbeat

    context = RuntimeContext(
        task_id="t",
        name="t",
        attempt=1,
        token=CancellationToken(),
        heartbeat=Heartbeat(),
    )
    context.progress(1, total=2)
    context.log("hello")
    assert list(context.logs) == ["hello"]


def test_log_records_into_the_buffer_and_the_stdlib_logger(
    caplog: pytest.LogCaptureFixture,
) -> None:
    # caplog 默认只捕获 WARNING 及以上，而 INFO 才是我们最常用的级别。
    caplog.set_level(logging.INFO, logger="chongming_task.task")

    controller = Controller(max_parallel=2)
    controller.register("talker", _talker)
    run = controller.submit("talker")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.ok

    assert outcome.value == "done"
    messages = [record.message for record in caplog.records]
    assert "hello" in messages, "INFO 日志应当走到标准库 logging"
    assert "careful" in messages

    record = next(r for r in caplog.records if r.message == "hello")
    assert record.task_id == run.snapshot.task_id
    assert record.task_name == "talker"
    assert record.attempt == 1


def test_task_context_is_the_runtime_context() -> None:
    """公开名字与内部名字是**同一个类**。

    任务作者只面对一个 ``ctx``；两个类会立刻带来"我该用哪个"的困惑。
    """
    from chongming_task import RuntimeContext, TaskContext

    assert TaskContext is RuntimeContext


def test_checkpoint_still_works_alongside_progress() -> None:
    controller = Controller(max_parallel=2)
    controller.register("stopme", _stopme)
    run = controller.submit("stopme")
    assert _wait_until(lambda: run.status is TaskStatus.RUNNING)

    controller.stop(run.snapshot.task_id, reason="收工")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.status is TaskStatus.STOPPED
    controller.close(timeout=_WINDOW)


# ---------------------------------------------------------------------------
# 3. 上游数据（管道传值）
# ---------------------------------------------------------------------------


def test_downstream_receives_upstream_values() -> None:
    controller = Controller(max_parallel=4)
    controller.register("produce", _produce)
    controller.register("consume", _consume)

    pipeline = Pipeline("data")
    source = pipeline.node("produce", node_id="source")
    pipeline.node("consume", node_id="sink", depends_on=[source])

    run = controller.submit_pipeline(pipeline, pipeline_run_id="d1")
    outcome = run.wait(timeout=_WINDOW)
    assert outcome.status.value == "succeeded"
    assert outcome.nodes["sink"] is TaskStatus.SUCCEEDED

    snapshot = controller.get("d1:sink")
    assert snapshot is not None
    assert snapshot.outcome == {"value": 42}, "下游应当拿到上游的返回值 42"
    controller.close(timeout=_WINDOW)


def test_a_failed_upstream_contributes_no_value() -> None:
    """失败的上游不塞 ``None``：那会让下游以为"上游返回了 None"。"""
    controller = Controller(max_parallel=4)
    controller.register("boom", _boom)
    controller.register("inspect", _inspect)

    pipeline = Pipeline("missing")
    bad = pipeline.node("boom", node_id="bad")
    pipeline.node("inspect", node_id="probe", depends_on=[bad], join=_all_settled())

    run = controller.submit_pipeline(pipeline, pipeline_run_id="d2")
    run.wait(timeout=_WINDOW)

    snapshot = controller.get("d2:probe")
    assert snapshot is not None
    assert snapshot.outcome == {"value": []}, "上游没有值，ctx.upstream 必须是空的"
    controller.close(timeout=_WINDOW)


def test_input_can_be_read_with_a_default() -> None:
    controller = Controller()
    controller.register("use_default", _use_default)
    run = controller.submit("use_default")
    outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
    assert outcome.value == "fallback"
    controller.close(timeout=_WINDOW)


def test_inputs_cannot_be_injected_after_the_task_started() -> None:
    controller = Controller(max_parallel=2)
    controller.register("stopme", _stopme)
    run = controller.submit("stopme")
    assert _wait_until(lambda: run.status is TaskStatus.RUNNING)

    with pytest.raises(RuntimeError, match="已经开始执行"):
        run.set_inputs({"x": 1})

    controller.stop(run.snapshot.task_id)
    controller.wait_all(timeout=_WINDOW)
    controller.close(timeout=_WINDOW)


# ---------------------------------------------------------------------------
# 任务体
# ---------------------------------------------------------------------------


def _stub() -> Any:
    from chongming_task.models import TaskSnapshot

    return TaskSnapshot(task_id="stub", name="stub")


def _no_backoff() -> Any:
    from chongming_task.policies import FixedBackoff

    return FixedBackoff(0.0)


def _flaky(ctx: Any) -> str:
    state = getattr(_flaky, "_count", 0) + 1
    _flaky._count = state  # type: ignore[attr-defined]
    if state < 3:
        raise RuntimeError(f"第 {state} 次失败")
    return "ok"


def _reporter(ctx: Any) -> str:
    for index in range(1, 4):
        ctx.checkpoint()
        ctx.progress(index, total=3, unit="rows")
    return "reported"


def _slow_reporter(ctx: Any) -> str:
    for index in range(1, 200):
        ctx.checkpoint()
        ctx.progress(index, total=200, unit="rows")
        time.sleep(0.01)
    return "done"


def _talker(ctx: Any) -> str:
    ctx.log("hello")
    ctx.warning("careful")
    return "done"


def _stopme(ctx: Any) -> str:
    for _ in range(1_000_000):
        ctx.checkpoint()
        time.sleep(0.005)
    return "never"


def _produce(ctx: Any) -> int:
    del ctx
    return 42


def _consume(ctx: Any) -> int:
    value = ctx.input("source")
    assert value == 42, f"预期 42，收到 {value!r}"
    return int(value)


def _inspect(ctx: Any) -> list[str]:
    return sorted(ctx.upstream)


def _use_default(ctx: Any) -> str:
    return str(ctx.input("nobody", "fallback"))


def _boom(ctx: Any) -> str:
    del ctx
    raise RuntimeError("故意失败")


def _all_settled() -> Any:
    from chongming_task.policies import JoinMode

    return JoinMode.ALL_SETTLED


def _wait_until(predicate: Any, timeout: float = _WINDOW) -> bool:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.01)
    return False