"""Task dependency resolution and trigger rule evaluation."""
from __future__ import annotations

import logging

from crondag.sdk.state import TaskState
from crondag.sdk.trigger_rule import TriggerRule

logger = logging.getLogger(__name__)

# Terminal states
TERMINAL_STATES = {TaskState.SUCCESS, TaskState.FAILED, TaskState.UPSTREAM_FAILED}


def evaluate_task_readiness(
    task_id: str,
    current_state: TaskState,
    trigger_rule: TriggerRule,
    upstream_states: dict[str, TaskState],
) -> TaskState:
    """Evaluate whether a task is ready to be scheduled.

    Args:
        task_id: The task being evaluated.
        current_state: Current state of this task.
        trigger_rule: The task's trigger rule.
        upstream_states: {upstream_task_id: state} for all upstream tasks.

    Returns:
        New state: SCHEDULED if ready, UPSTREAM_FAILED if blocked,
        or the current state if not yet determined.
    """
    if current_state != TaskState.PENDING:
        return current_state

    # No upstream dependencies — immediately schedulable
    if not upstream_states:
        return TaskState.SCHEDULED

    # Check if all upstreams are in terminal states
    all_done = all(s in TERMINAL_STATES for s in upstream_states.values())
    if not all_done:
        return TaskState.PENDING  # Still waiting

    all_success = all(s == TaskState.SUCCESS for s in upstream_states.values())
    any_failed = any(s in {TaskState.FAILED, TaskState.UPSTREAM_FAILED} for s in upstream_states.values())

    if trigger_rule == TriggerRule.ALL_SUCCESS:
        if all_success:
            return TaskState.SCHEDULED
        return TaskState.UPSTREAM_FAILED

    if trigger_rule == TriggerRule.ALL_DONE:
        return TaskState.SCHEDULED  # Don't care about upstream outcomes

    if trigger_rule == TriggerRule.ONE_FAILED:
        if any_failed:
            return TaskState.SCHEDULED
        if all_success:
            return TaskState.UPSTREAM_FAILED  # No failures, won't trigger
        return TaskState.PENDING

    return current_state
