cron.py

# Copyright 2026 Impala Systems, Inc.
# SPDX-License-Identifier: AGPL-3.0-only

from enum import Enum
from dataclasses import dataclass, field, is_dataclass
from typing import List, Optional, Dict, Any, TypedDict
import torch

class CronPatternType(Enum):
    WILDCARD = 0
    VALUE = 1
    LIST = 2
    RANGE = 3
    STEP = 4
    NTH_WEEKDAY = 5
    LAST = 6
    NEAREST = 7


@dataclass
class CronComponentOutput:
    """Raw logits for a single component from the model."""
    pattern_logits: torch.Tensor
    values_logits: torch.Tensor
    range_start_logits: torch.Tensor
    range_end_logits: torch.Tensor
    step_start_logits: torch.Tensor
    step_size_logits: torch.Tensor
    nth_logits: torch.Tensor
    last_offset_logits: torch.Tensor
    token_values_logits: Optional[torch.Tensor] = None
    token_step_size_logits: Optional[torch.Tensor] = None
    token_pattern_logits: Optional[torch.Tensor] = None
    value_count_logits: Optional[torch.Tensor] = None
    list_values_logits: Optional[torch.Tensor] = None
    list_item_logits: Optional[torch.Tensor] = None
    residual_adapter_selector_logits: Optional[torch.Tensor] = None
    constrained_field_logits: Optional[torch.Tensor] = None
    interval_scope_selector_logits: Optional[torch.Tensor] = None
    use_value_count_decoding: bool = False
    use_token_value_decoding: bool = False
    use_path_score_decoding: bool = False
    use_list_item_decoding: bool = False
    use_list_item_scoring: bool = False
    list_item_score_scale: float = 1.0


@dataclass
class CronOutput:
    """Complete model output containing component outputs."""
    minute: CronComponentOutput
    hour: CronComponentOutput
    dom: CronComponentOutput
    month: CronComponentOutput
    dow: CronComponentOutput
    dom_dow_intersect_logits: torch.Tensor
    semantic_binding_logits: Optional[torch.Tensor] = None

    def get_component(self, name: str) -> CronComponentOutput:
        return getattr(self, name)


class CronComponentDict(TypedDict):
    pattern_output_ids: torch.Tensor
    values_output_ids: torch.Tensor
    range_start_output_ids: torch.Tensor
    range_end_output_ids: torch.Tensor
    step_start_output_ids: torch.Tensor
    step_size_output_ids: torch.Tensor
    nth_output_ids: torch.Tensor
    last_offset_output_ids: torch.Tensor

class CronDict(TypedDict):
    minute: CronComponentDict
    hour: CronComponentDict
    dom: CronComponentDict
    month: CronComponentDict
    dow: CronComponentDict
    dom_dow_intersect_output_ids: torch.Tensor


@dataclass
class CronComponent:
    type: CronPatternType
    values: List[int] = field(default_factory=list)
    items: List["CronComponent"] = field(default_factory=list)
    step_start: Optional[int] = None
    step_size: Optional[int] = None
    nth: Optional[int] = None
    range_start: Optional[int] = None
    range_end: Optional[int] = None
    last_offset: Optional[int] = None

    def to_string(self, component_index: int) -> str:
        if self.type == CronPatternType.WILDCARD:
            return "*"
        elif self.type == CronPatternType.VALUE:
            if not self.values:
                return "*" # Fallback
            return str(self.values[0])
        elif self.type == CronPatternType.LIST:
            if self.items:
                return ",".join(item.to_string(component_index) for item in self.items)
            if not self.values:
                return "*" # Fallback
            return ",".join(map(str, sorted(self.values)))
        elif self.type == CronPatternType.RANGE:
            start = self.range_start if self.range_start is not None else 0
            end = self.range_end if self.range_end is not None else self.range_start
            return f"{start}-{end}"
        elif self.type == CronPatternType.STEP:
            step_size = self.step_size if self.step_size is not None else 1
            if step_size == 1 and (self.step_start is None or self.step_start == 0):
                # Typically */1
                return "*"
            elif self.step_start is None or self.step_start == 0:
                # Typically */step
                base = "*"
            else:
                base = str(self.step_start)
            return f"{base}/{step_size}"
        elif self.type == CronPatternType.NTH_WEEKDAY:
            # e.g. 2#2 (Tuesday#2nd)
            if not self.values:
                return "*"
            dow = self.values[0]
            nth = self.nth if self.nth is not None else 1
            return f"{dow}#{nth}"
        elif self.type == CronPatternType.LAST:
             # L or 5L
            if self.last_offset:
                return f"L-{self.last_offset}"
            if not self.values:
                return "L"
            return f"{self.values[0]}L"
        elif self.type == CronPatternType.NEAREST:
            # 15W
            if not self.values:
                return "1W" # Fallback
            return f"{self.values[0]}W"
        return "*"

@dataclass
class Cron:
    minute: CronComponent
    hour: CronComponent
    dom: CronComponent
    month: CronComponent
    dow: CronComponent
    dom_dow_intersect: bool = True

    def to_vixie_cron(self) -> str:
        # TODO: Handle dom_dow_intersect
        parts = [
            self.minute.to_string(0),
            self.hour.to_string(1),
            self.dom.to_string(2),
            self.month.to_string(3),
            self.dow.to_string(4),
        ]
        return " ".join(parts)

# Pre-defined dimensions for each component
# Minute: 0-59 (60)
# Hour: 0-23 (24)
# DoM: 1-31 (32, 0 is unused/padding)
# Month: 1-12 (13, 0 is unused/padding)
# DoW: 0-6 (7)
COMPONENT_DIMS = [60, 24, 32, 13, 7]
COMPONENT_NAMES = ["minute", "hour", "dom", "month", "dow"]
MAX_LIST_ITEMS = 4
SEMANTIC_BINDING_NAMES = [
    "weekday_interval_all_day",
    "workday_dow_required",
    "month_end_last_dom",
    "office_hours_end_17",
    "office_hours_end_18",
    "hour_window_required",
    "hour_window_absent",
]

def _component_from_path_score(comp_output: CronComponentOutput, component_index: int, batch_index: int) -> CronComponent:
    pattern_logits = comp_output.pattern_logits[batch_index]
    values_logits = comp_output.values_logits[batch_index]
    if comp_output.use_token_value_decoding and comp_output.token_values_logits is not None:
        values_logits = comp_output.token_values_logits[batch_index].amax(dim=0)
    list_values_logits = (
        comp_output.list_values_logits[batch_index]
        if comp_output.list_values_logits is not None
        else values_logits
    )
    range_start_logits = comp_output.range_start_logits[batch_index]
    range_end_logits = comp_output.range_end_logits[batch_index]
    step_start_logits = comp_output.step_start_logits[batch_index]
    step_size_logits = comp_output.step_size_logits[batch_index]
    nth_logits = comp_output.nth_logits[batch_index]
    last_offset_logits = comp_output.last_offset_logits[batch_index]
    value_count_logits = (
        comp_output.value_count_logits[batch_index]
        if comp_output.value_count_logits is not None
        else None
    )

    def list_item_score(indices: list[int]) -> torch.Tensor:
        if (
            not comp_output.use_list_item_scoring
            or comp_output.list_item_logits is None
            or not indices
        ):
            return torch.tensor(0.0, dtype=list_values_logits.dtype, device=list_values_logits.device)
        item_logits = comp_output.list_item_logits[batch_index]
        score = torch.tensor(0.0, dtype=item_logits.dtype, device=item_logits.device)
        for slot, value in enumerate(sorted(indices)[: item_logits.shape[0]]):
            if 0 <= value < item_logits.shape[-1]:
                score = score + item_logits[slot, value]
        return score * comp_output.list_item_score_scale

    def list_item_indices() -> list[int]:
        if not comp_output.use_list_item_decoding or comp_output.list_item_logits is None:
            return []
        return scoreable_list_item_indices()

    def scoreable_list_item_indices() -> list[int]:
        if comp_output.list_item_logits is None:
            return []
        item_logits = comp_output.list_item_logits[batch_index]
        if comp_output.use_value_count_decoding and value_count_logits is not None:
            value_count = int(torch.argmax(value_count_logits).item())
        else:
            value_count = int(torch.clamp(torch.round(torch.sigmoid(list_values_logits).sum()), min=1).item())
        value_count = max(1, min(value_count, item_logits.shape[0]))
        indices: list[int] = []
        for slot in range(value_count):
            ranked = torch.argsort(item_logits[slot], descending=True).flatten().tolist()
            selected = next((int(value) for value in ranked if int(value) not in indices), int(ranked[0]))
            indices.append(selected)
        return indices

    def value_component(pattern_type: CronPatternType) -> tuple[float, CronComponent]:
        value_index = int(torch.argmax(values_logits).item())
        return (
            float((pattern_logits[pattern_type.value] + values_logits[value_index]).item()),
            CronComponent(type=pattern_type, values=[value_index]),
        )

    argmax_pattern = CronPatternType(int(torch.argmax(pattern_logits).item()))
    candidates: list[tuple[float, CronPatternType, CronComponent]] = [
        (
            float(pattern_logits[CronPatternType.WILDCARD.value].item()),
            CronPatternType.WILDCARD,
            CronComponent(type=CronPatternType.WILDCARD),
        ),
    ]
    value_score, value_cron = value_component(CronPatternType.VALUE)
    candidates.append((value_score, CronPatternType.VALUE, value_cron))

    item_decoded_indices = list_item_indices()
    if item_decoded_indices:
        list_indices = item_decoded_indices
    elif comp_output.use_value_count_decoding and value_count_logits is not None:
        value_count = int(torch.argmax(value_count_logits).item())
        value_count = max(1, min(value_count, list_values_logits.numel()))
        list_indices = torch.topk(list_values_logits, value_count).indices.flatten().tolist()
    else:
        list_indices = (list_values_logits > 0.0).nonzero().flatten().tolist()
        if not list_indices:
            list_indices = [int(torch.argmax(list_values_logits).item())]
    list_score = (
        pattern_logits[CronPatternType.LIST.value]
        + list_values_logits[list_indices].sum()
        + list_item_score([int(value) for value in list_indices])
    )
    list_component = CronComponent(type=CronPatternType.LIST, values=[int(value) for value in list_indices])
    last_offset_idx = int(torch.argmax(last_offset_logits).item())
    if component_index == 2 and last_offset_idx == 0:
        list_component.items = [
            CronComponent(CronPatternType.VALUE, values=[value])
            for value in sorted(list_component.values)
        ]
        list_component.items.append(CronComponent(CronPatternType.LAST))
    candidates.append((float(list_score.item()), CronPatternType.LIST, list_component))
    if comp_output.use_list_item_scoring and not comp_output.use_list_item_decoding:
        scored_item_indices = scoreable_list_item_indices()
        if scored_item_indices and sorted(scored_item_indices) != sorted([int(value) for value in list_indices]):
            scored_item_list_score = (
                pattern_logits[CronPatternType.LIST.value]
                + list_values_logits[scored_item_indices].sum()
                + list_item_score([int(value) for value in scored_item_indices])
            )
            candidates.append(
                (
                    float(scored_item_list_score.item()),
                    CronPatternType.LIST,
                    CronComponent(
                        type=CronPatternType.LIST,
                        values=[int(value) for value in scored_item_indices],
                    ),
                )
            )

    range_start = int(torch.argmax(range_start_logits).item())
    range_end = int(torch.argmax(range_end_logits).item())
    candidates.append(
        (
            float(
                (
                    pattern_logits[CronPatternType.RANGE.value]
                    + range_start_logits[range_start]
                    + range_end_logits[range_end]
                ).item()
            ),
            CronPatternType.RANGE,
            CronComponent(type=CronPatternType.RANGE, range_start=range_start, range_end=range_end),
        )
    )

    step_start = int(torch.argmax(step_start_logits).item())
    step_size = int(torch.argmax(step_size_logits).item())
    candidates.append(
        (
            float(
                (
                    pattern_logits[CronPatternType.STEP.value]
                    + step_start_logits[step_start]
                    + step_size_logits[step_size]
                ).item()
            ),
            CronPatternType.STEP,
            CronComponent(type=CronPatternType.STEP, step_start=step_start, step_size=step_size),
        )
    )

    nth_value = int(torch.argmax(values_logits).item())
    nth = int(torch.argmax(nth_logits).item())
    candidates.append(
        (
            float(
                (
                    pattern_logits[CronPatternType.NTH_WEEKDAY.value]
                    + values_logits[nth_value]
                    + nth_logits[nth]
                ).item()
            ),
            CronPatternType.NTH_WEEKDAY,
            CronComponent(type=CronPatternType.NTH_WEEKDAY, values=[nth_value], nth=nth),
        )
    )

    last_values = (values_logits > 0.0).nonzero().flatten().tolist()
    last_value_score = torch.tensor(0.0, dtype=values_logits.dtype, device=values_logits.device)
    last_component = CronComponent(type=CronPatternType.LAST)
    if component_index != 2 and last_values:
        last_value = int(torch.argmax(values_logits).item())
        last_component.values = [last_value]
        last_value_score = values_logits[last_value]
    if last_offset_idx > 0:
        last_component.last_offset = last_offset_idx
    candidates.append(
        (
            float(
                (
                    pattern_logits[CronPatternType.LAST.value]
                    + last_value_score
                    + last_offset_logits[last_offset_idx]
                ).item()
            ),
            CronPatternType.LAST,
            last_component,
        )
    )

    nearest_value = int(torch.argmax(values_logits).item())
    candidates.append(
        (
            float(
                (
                    pattern_logits[CronPatternType.NEAREST.value]
                    + values_logits[nearest_value]
                ).item()
            ),
            CronPatternType.NEAREST,
            CronComponent(type=CronPatternType.NEAREST, values=[nearest_value]),
        )
    )

    candidate_by_pattern = {pattern: (score, component) for score, pattern, component in candidates}
    if argmax_pattern != CronPatternType.WILDCARD:
        return candidate_by_pattern[argmax_pattern][1]

    wildcard_score = candidate_by_pattern[CronPatternType.WILDCARD][0]
    eligible = [candidate_by_pattern[CronPatternType.WILDCARD]]
    range_score, range_component = candidate_by_pattern[CronPatternType.RANGE]
    range_margin = 10.0
    if component_index != 4 and range_score - wildcard_score >= range_margin:
        eligible.append((range_score, range_component))
    if component_index == 2:
        last_score, last_component = candidate_by_pattern[CronPatternType.LAST]
        if last_score - wildcard_score >= 6.0:
            eligible.append((last_score, last_component))

    return max(eligible, key=lambda item: item[0])[1]


def logits_to_cron(output: CronOutput) -> List[Cron]:
    """
    Converts structured CronOutput into a list of Cron objects.
    """
    # Assuming batch size is consistent across all tensors in output
    batch_size = output.minute.pattern_logits.size(0)
    crons = []
    
    # Store component data for each item in batch
    batch_components = [[] for _ in range(batch_size)]
    
    for component_index, comp_name in enumerate(COMPONENT_NAMES):
        comp_output = output.get_component(comp_name)
        
        # Now process each item in the batch for this component
        for b in range(batch_size):
            if comp_output.use_path_score_decoding:
                batch_components[b].append(_component_from_path_score(comp_output, component_index, b))
                continue

            # Extract
            pattern_logits = comp_output.pattern_logits[b]
            values_logits = comp_output.values_logits[b]
            if comp_output.use_token_value_decoding and comp_output.token_values_logits is not None:
                values_logits = comp_output.token_values_logits[b].amax(dim=0)
            list_values_logits = (
                comp_output.list_values_logits[b]
                if comp_output.list_values_logits is not None
                else values_logits
            )
            rs_logits = comp_output.range_start_logits[b]
            re_logits = comp_output.range_end_logits[b]
            step_start_logits = comp_output.step_start_logits[b]
            step_size_logits = comp_output.step_size_logits[b]
            nth_logits = comp_output.nth_logits[b]
            value_count_logits = comp_output.value_count_logits[b] if comp_output.value_count_logits is not None else None
            
            pattern_idx = torch.argmax(pattern_logits).item()
            pattern_type = CronPatternType(pattern_idx)
            
            # Values (Sigmoid)
            values_indices = (values_logits > 0.0).nonzero().flatten().tolist()
            
            rs_idx = torch.argmax(rs_logits).item()
            re_idx = torch.argmax(re_logits).item()
            step_start_idx = torch.argmax(step_start_logits).item()
            step_size_idx = torch.argmax(step_size_logits).item()
            nth_idx = torch.argmax(nth_logits).item()
            last_offset_idx = torch.argmax(comp_output.last_offset_logits[b]).item()
            
            # Construct Component
            comp = CronComponent(type=pattern_type)
            
            if pattern_type == CronPatternType.WILDCARD:
                pass
            elif pattern_type == CronPatternType.VALUE:
                if not values_indices:
                    comp.values = [torch.argmax(values_logits).item()]
                else:
                    # Pick the highest probability value among the candidates
                    best_idx = max(values_indices, key=lambda i: values_logits[i].item())
                    comp.values = [best_idx]
            elif pattern_type == CronPatternType.LIST:
                values_logits = list_values_logits
                values_indices = (values_logits > 0.0).nonzero().flatten().tolist()
                if comp_output.use_list_item_decoding and comp_output.list_item_logits is not None:
                    item_logits = comp_output.list_item_logits[b]
                    if comp_output.use_value_count_decoding and value_count_logits is not None:
                        value_count = int(torch.argmax(value_count_logits).item())
                    else:
                        value_count = int(torch.clamp(torch.round(torch.sigmoid(values_logits).sum()), min=1).item())
                    value_count = max(1, min(value_count, item_logits.shape[0]))
                    values_indices = []
                    for slot in range(value_count):
                        ranked = torch.argsort(item_logits[slot], descending=True).flatten().tolist()
                        selected = next((int(value) for value in ranked if int(value) not in values_indices), int(ranked[0]))
                        values_indices.append(selected)
                elif comp_output.use_value_count_decoding and value_count_logits is not None:
                    value_count = int(torch.argmax(value_count_logits).item())
                    value_count = max(1, min(value_count, values_logits.numel()))
                    values_indices = torch.topk(values_logits, value_count).indices.flatten().tolist()
                elif not values_indices:
                    values_indices = [torch.argmax(values_logits).item()]
                comp.values = values_indices
                if component_index == 2 and last_offset_idx == 0:
                    comp.items = [CronComponent(CronPatternType.VALUE, values=[value]) for value in sorted(values_indices)]
                    comp.items.append(CronComponent(CronPatternType.LAST))
            elif pattern_type == CronPatternType.RANGE:
                comp.range_start = rs_idx
                comp.range_end = re_idx
            elif pattern_type == CronPatternType.STEP:
                comp.step_start = step_start_idx
                comp.step_size = step_size_idx
            elif pattern_type == CronPatternType.NTH_WEEKDAY:
                comp.values = [torch.argmax(values_logits).item()]
                comp.nth = nth_idx
            elif pattern_type == CronPatternType.LAST:
                 if component_index != 2 and values_indices:
                     comp.values = [torch.argmax(values_logits).item()]
                 if last_offset_idx > 0:
                     comp.last_offset = last_offset_idx
            elif pattern_type == CronPatternType.NEAREST:
                 comp.values = [torch.argmax(values_logits).item()]

            batch_components[b].append(comp)

    # Assemble Cron objects
    for b in range(batch_size):
        comps = batch_components[b] # Order matches loop above
        
        dom_dow_logits = output.dom_dow_intersect_logits[b]
        dom_dow_val = torch.argmax(dom_dow_logits).item() == 1
        
        crons.append(Cron(
            minute=comps[0],
            hour=comps[1],
            dom=comps[2],
            month=comps[3],
            dow=comps[4],
            dom_dow_intersect=dom_dow_val
        ))
        
    return crons



def cron_string_to_cron(cron_str: str) -> Cron:
    """
    Parses a Vixie Cron string into a Cron object.
    Simplistic parser for training demo.
    """
    parts = cron_str.split()
    if len(parts) != 5:
        parts = ["*", "*", "*", "*", "*"]
        
    components_data = [
        (parts[0], 0, 59),   # minute
        (parts[1], 0, 23),   # hour
        (parts[2], 1, 31),   # dom
        (parts[3], 1, 12),   # month
        (parts[4], 0, 6),    # dow
    ]
    
    cron_components = []

    for i, (expr, min_val, max_val) in enumerate(components_data):
        comp = CronComponent(type=CronPatternType.VALUE) # Default
        
        if expr == "*":
            comp.type = CronPatternType.WILDCARD
        elif "," in expr:
            comp.type = CronPatternType.LIST
            values = []
            for val in expr.split(","):
                try:
                    v = int(val)
                    values.append(v)
                except: pass
            comp.values = values
        elif "/" in expr:
            comp.type = CronPatternType.STEP
            try:
                base, step = expr.split("/")
                if base != "*":
                    try:
                        v = int(base)
                        comp.step_start = v
                    except: pass
                else:
                    comp.step_start = 0
                comp.step_size = int(step)
            except: pass
        elif "#" in expr:
            comp.type = CronPatternType.NTH_WEEKDAY
            try:
                dow, nth = map(int, expr.split("#"))
                comp.values = [dow]
                comp.nth = nth
            except: pass
        elif "L" in expr:
            comp.type = CronPatternType.LAST
            if expr.startswith("L-"):
                try:
                    comp.last_offset = int(expr[2:])
                except: pass
            elif expr != "L":
                try:
                    val = int(expr.replace("L", ""))
                    comp.values = [val]
                except: pass
        elif "-" in expr:
            comp.type = CronPatternType.RANGE
            try:
                s, e = map(int, expr.split("-"))
                comp.range_start = s
                comp.range_end = e
            except: pass
        elif "W" in expr:
            comp.type = CronPatternType.NEAREST
            try:
                val = int(expr.replace("W", ""))
                comp.values = [val]
            except: pass
        else:
            comp.type = CronPatternType.VALUE
            try:
                v = int(expr)
                comp.values = [v]
            except: pass
            
        cron_components.append(comp)

    return Cron(
        minute=cron_components[0],
        hour=cron_components[1],
        dom=cron_components[2],
        month=cron_components[3],
        dow=cron_components[4],
        dom_dow_intersect=True # heuristic
    )

def cron_to_dict(cron: Cron) -> CronDict:
    """
    Converts a Cron object into a structured dictionary of tensors.
    """
    # Order: Minute, Hour, DoM, Month, DoW
    components = [cron.minute, cron.hour, cron.dom, cron.month, cron.dow]
    target_dict = {}
    
    for i, comp in enumerate(components):
        dim_size = COMPONENT_DIMS[i]
        comp_name = COMPONENT_NAMES[i]
        
        pattern_idx = comp.type.value
        values_hot = [0.0] * dim_size
        range_start_idx = -100
        range_end_idx = -100
        step_start_idx = -100
        step_size_idx = -100
        nth_idx = -100
        last_offset_idx = -100
        
        # Values
        if comp.values:
            for v in comp.values:
                if 0 <= v < dim_size:
                    values_hot[v] = 1.0
        # Range
        if comp.type == CronPatternType.RANGE:
             if comp.range_start is not None: range_start_idx = comp.range_start
             if comp.range_end is not None: range_end_idx = comp.range_end
        
        # Step
        if comp.type == CronPatternType.STEP:
             if comp.step_start is not None: step_start_idx = comp.step_start
             if comp.step_size is not None: step_size_idx = comp.step_size
        
        # Nth
        if comp.type == CronPatternType.NTH_WEEKDAY:
             if comp.nth is not None: nth_idx = comp.nth

        # Last offset. 0 is an explicit "plain L / no offset" class for LAST.
        if comp.type == CronPatternType.LAST:
            last_offset_idx = comp.last_offset if comp.last_offset is not None else 0
        elif comp.type == CronPatternType.LIST:
            has_last_item = any(item.type == CronPatternType.LAST for item in comp.items)
            if has_last_item:
                last_offset_idx = 0
            else:
                last_offset_idx = dim_size - 1
             
        # Create Component Dict
        target_dict[comp_name] = {
            "pattern_output_ids": torch.tensor(pattern_idx, dtype=torch.long),
            "values_output_ids": torch.tensor(values_hot, dtype=torch.float),
            "range_start_output_ids": torch.tensor(range_start_idx, dtype=torch.long),
            "range_end_output_ids": torch.tensor(range_end_idx, dtype=torch.long),
            "step_start_output_ids": torch.tensor(step_start_idx, dtype=torch.long),
            "step_size_output_ids": torch.tensor(step_size_idx, dtype=torch.long),
            "nth_output_ids": torch.tensor(nth_idx, dtype=torch.long),
            "last_offset_output_ids": torch.tensor(last_offset_idx, dtype=torch.long),
        }

    dd_idx = 1 if cron.dom_dow_intersect else 0
    target_dict["dom_dow_intersect_output_ids"] = torch.tensor(dd_idx, dtype=torch.long)
    
    return target_dict

def parse_cron_string(cron_str: str) -> CronDict:
    """
    Wrapper to return target dict.
    """
    cron = cron_string_to_cron(cron_str)
    return cron_to_dict(cron)


def parse_cron_schema(cron_schema: Dict[str, Any] | str) -> CronDict:
    """
    Converts a Cronformer proto JSON CronExpression into target tensors.
    """
    if isinstance(cron_schema, str):
        return parse_cron_string(cron_schema)

    field_keys = ["minute", "hour", "dayOfMonth", "month", "dayOfWeek"]

    def kind_name(value: Any) -> str:
        if isinstance(value, int):
            return CronPatternType(value - 1).name if value > 0 else "WILDCARD"
        return str(value).replace("CRON_FIELD_KIND_", "")

    def component_from_field(field: Dict[str, Any]) -> CronComponent:
        kind = kind_name(field.get("kind", "CRON_FIELD_KIND_WILDCARD"))
        values = [int(v) for v in field.get("values", [])]

        if kind == "WILDCARD":
            return CronComponent(type=CronPatternType.WILDCARD)
        if kind == "VALUE":
            return CronComponent(type=CronPatternType.VALUE, values=values)
        if kind == "LIST":
            if field.get("items"):
                items = [component_from_field(item) for item in field["items"]]
                numeric_values = []
                for item in items:
                    if item.type == CronPatternType.VALUE:
                        numeric_values.extend(item.values)
                return CronComponent(type=CronPatternType.LIST, values=numeric_values, items=items)
            return CronComponent(type=CronPatternType.LIST, values=values)
        if kind == "RANGE":
            return CronComponent(
                type=CronPatternType.RANGE,
                range_start=int(field.get("rangeStart", 0)),
                range_end=int(field.get("rangeEnd", 0)),
            )
        if kind == "STEP":
            return CronComponent(
                type=CronPatternType.STEP,
                step_start=int(field.get("stepStart", 0)),
                step_size=int(field.get("stepSize", 1)),
            )
        if kind == "NTH_WEEKDAY":
            return CronComponent(
                type=CronPatternType.NTH_WEEKDAY,
                values=values,
                nth=int(field.get("nth", 1)),
            )
        if kind == "LAST":
            return CronComponent(
                type=CronPatternType.LAST,
                values=values,
                last_offset=int(field["lastOffset"]) if "lastOffset" in field else None,
            )
        if kind == "NEAREST_WEEKDAY":
            return CronComponent(type=CronPatternType.NEAREST, values=values)
        raise ValueError(f"Unsupported cron field kind: {kind}")

    components = [component_from_field(cron_schema[key]) for key in field_keys]
    semantics = cron_schema.get("domDowSemantics", "DOM_DOW_SEMANTICS_OR")
    return cron_to_dict(
        Cron(
            minute=components[0],
            hour=components[1],
            dom=components[2],
            month=components[3],
            dow=components[4],
            dom_dow_intersect=semantics == "DOM_DOW_SEMANTICS_AND",
        )
    )


def recursive_to_device(data, device):
    """
    Recursively moves data (Dict, List, or Tensor) to the specified device.
    """
    if isinstance(data, dict):
        return {k: recursive_to_device(v, device) for k, v in data.items()}
    elif isinstance(data, list):
        return [recursive_to_device(v, device) for v in data]
    elif isinstance(data, torch.Tensor):
        return data.to(device)
    elif is_dataclass(data) and not isinstance(data, type):
        values = {field_name: recursive_to_device(getattr(data, field_name), device) for field_name in data.__dataclass_fields__}
        return data.__class__(**values)
    return data