from __future__ import annotations

from typing import TYPE_CHECKING, ClassVar, Dict, List, Literal, Optional

from pydantic import Field


from .models import BaseTrait, CountingRule

if TYPE_CHECKING:
    pass


def _uncounted_rules_factory():
    return [CountingRule(counter_name="initial", on_event="received")]


class UncountedTrait(BaseTrait):
    """Marker trait to identify a task that is not yet counted."""

    trait_name: Literal["uncounted"] = "uncounted"
    counting_rules: Optional[List[CountingRule]] = Field(
        default_factory=_uncounted_rules_factory
    )


def _generated_rules_factory():
    return [CountingRule(counter_name="generated", on_event="received")]


class GeneratedTrait(BaseTrait):
    """Marker trait to identify a task that was generated by the worker."""

    trait_name: Literal["generated"] = "generated"
    counting_rules: Optional[List[CountingRule]] = Field(
        default_factory=_generated_rules_factory
    )
    forbids_traits: list[str] = ["uncounted"]


class AccountingTrait(BaseTrait):
    """Data component for the unified, hook-driven accounting system for tasks."""

    _decorator_fields: ClassVar[list[str]] = []
    trait_name: Literal["accounting"] = "accounting"
    counted: Dict[str, bool] = Field(default_factory=dict)
