Budget Lifecycle

from uuid import uuid4
from datetime import date, datetime, UTC

from psp.components.budget.domain import (
    TokenType, BudgetPlan, Period, PeriodType,
    ScheduledTask, TaskCost, ScheduleStatus,
)
from psp.components.budget.adapters.persistence import (
    InMemoryTokenTypeRepository,
    InMemoryBudgetPlanRepository,
    InMemoryScheduledTaskRepository,
)

# Setup repositories
token_repo = InMemoryTokenTypeRepository()
plan_repo = InMemoryBudgetPlanRepository()
scheduled_repo = InMemoryScheduledTaskRepository()

owner_id = uuid4()
now = datetime.now(UTC)

# Create token types
deep = TokenType(id=uuid4(), owner_id=owner_id, name="deep", color="#3b82f6")
admin = TokenType(id=uuid4(), owner_id=owner_id, name="admin", color="#f59e0b")
token_repo.save(deep)
token_repo.save(admin)
print(f"Created token types: {deep.name}, {admin.name}")

# Create budget plan
plan = BudgetPlan(
    id=uuid4(),
    owner_id=owner_id,
    period_type=PeriodType.DAY,
    effective_from=date.today(),
    allocations={deep.id: 4, admin.id: 2},  # 4 deep, 2 admin per day
)
plan_repo.save(plan)
print(f"Budget plan: {plan.allocations}")

# Schedule a task
period = Period.today(PeriodType.DAY)
scheduled = ScheduledTask(
    id=uuid4(),
    task_id=uuid4(),  # Reference to task in Todos
    owner_id=owner_id,
    period=period,
    cost=TaskCost(costs={deep.id: 2}),  # 2 deep tokens
)
scheduled_repo.save(scheduled)
print(f"Scheduled task with cost: {scheduled.cost.costs}")

# Complete task (mark budget as spent)
spent = scheduled.mark_spent(now)
scheduled_repo.save(spent)
print(f"Task completed, status: {spent.status.value}")