import json
import time
import uuid
from collections.abc import Iterator
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Optional, Literal, List, Dict, Union

import yaml
import secrets
from jose import jwt
from jinja2 import Template
from passlib.context import CryptContext
from pydantic import BaseModel, Field, model_validator
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError


class TransientStorage(BaseModel):
    storage: dict[str, Any] = {}

    def get(self, key: str, default: Any = None, ensure_exists: bool = False) -> Any:
        if ensure_exists and key not in self.storage:
            raise KeyError(f"Key '{key}' not found in storage.")
        return self.storage.get(key, default)

    def set(self, key: str, value: Any) -> "TransientStorage":
        self.storage[key] = value
        return self

    def delete(self, key: str) -> "TransientStorage":
        if key in self.storage:
            del self.storage[key]
        return self

    def clear(self) -> "TransientStorage":
        self.storage.clear()
        return self

    def keys(self) -> list[str]:
        return list(self.storage.keys())

    def values(self) -> list[Any]:
        return list(self.storage.values())

    def items(self) -> list[tuple[str, Any]]:
        return list(self.storage.items())

    def __contains__(self, key: str) -> bool:
        return key in self.storage

    def __iter__(self) -> Iterator[str]:
        return iter(self.storage)


class TransientStorageMixin:
    def get_or_create_transient_storage(self) -> TransientStorage:
        if "_transient_storage" not in self.__dict__:
            self.__dict__["_transient_storage"] = TransientStorage()
        return self.__dict__["_transient_storage"]


### SECURITY

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
ALGORITHM = "HS256"


def validate_token(token: str, secret_key: str) -> dict[str, Any]:
    payload = jwt.decode(token, secret_key, algorithms=[ALGORITHM])
    return payload


def create_access_token(
    subject: str | Any, expires_delta: timedelta, secret_key: str
) -> str:
    expire = datetime.utcnow() + expires_delta
    to_encode = {"exp": expire, "sub": str(subject)}
    encoded_jwt = jwt.encode(to_encode, secret_key, algorithm=ALGORITHM)
    return encoded_jwt


def verify_password(plain_password: str, hashed_password: str) -> bool:
    return pwd_context.verify(plain_password, hashed_password)


def get_password_hash(password: str) -> str:
    return pwd_context.hash(password)


def get_time(
    in_future: Optional[float] = None,
    units: Literal["hours", "minutes", "seconds", None] = None,
) -> float:
    if in_future is not None:
        if units == "hours":
            time_delta = in_future * 3600
        elif units == "minutes":
            time_delta = in_future * 60
        elif units == "seconds":
            time_delta = in_future
        else:
            time_delta = in_future
    else:
        time_delta = 0
    return time.time() + time_delta


def to_seconds(duration: float, units: Literal["hours", "minutes", "seconds"]) -> float:
    if units == "hours":
        return duration * 3600
    elif units == "minutes":
        return duration * 60
    elif units == "seconds":
        return duration
    else:
        raise ValueError(f"Invalid units: {units}")


def from_seconds(
    duration: float, units: Literal["hours", "minutes", "seconds"]
) -> float:
    if units == "hours":
        return duration / 3600
    elif units == "minutes":
        return duration / 60
    elif units == "seconds":
        return duration
    else:
        raise ValueError(f"Invalid units: {units}")


def _parse_timezone(tz_str: str) -> ZoneInfo | timedelta | None:
    """Parse a timezone string into a ZoneInfo or fixed-offset timedelta.

    Handles:
      - IANA names: "America/New_York", "Europe/London"
      - Simple names: "UTC", "GMT"
      - Offset strings: "GMT-04:00", "UTC+5:30", "+0530", "-04:00"

    Returns None if the string is unrecognizable.
    """
    if not tz_str:
        return None

    # IANA names contain "/" (or are exactly UTC/GMT)
    if "/" in tz_str or tz_str in {"UTC", "GMT"}:
        try:
            return ZoneInfo(tz_str)
        except (ZoneInfoNotFoundError, KeyError):
            return None

    # Strip GMT/UTC prefix to get the offset part
    offset_str = tz_str
    if offset_str.startswith(("GMT", "UTC")):
        offset_str = offset_str[3:]

    if not offset_str:
        return ZoneInfo("UTC")

    # Parse offset: "+HH:MM", "-HH:MM", "+HHMM", "-HHMM", "+H"
    try:
        sign = 1 if offset_str[0] == "+" else -1
        rest = offset_str[1:]
        if ":" in rest:
            hh, mm = rest.split(":", 1)
        elif len(rest) <= 2:
            hh, mm = rest, "0"
        else:
            hh, mm = rest[:2], rest[2:]
        return timedelta(hours=sign * int(hh), minutes=sign * int(mm))
    except (ValueError, IndexError):
        return None


def format_time(
    seconds_since_epoch: float | None = None,
    date_time_seperator: str = " ",
    include_day_of_week: bool = False,
    day_date_separator: str = " ",
    timezone: str | None = None,
) -> str:
    """Format timestamp with optional timezone and day of week.

    Args:
        seconds_since_epoch: Unix timestamp (defaults to current time)
        date_time_seperator: Separator between date and time
        include_day_of_week: Whether to include day name
        day_date_separator: Separator between day name and date
        timezone: Timezone string — IANA name (e.g. 'America/New_York'),
                  or offset (e.g. 'GMT-04:00', 'UTC+5:30'). If None, returns UTC with 'Z' suffix.
    """
    if seconds_since_epoch is None:
        seconds_since_epoch = get_time()

    # Convert to datetime object in UTC
    dt = datetime.fromtimestamp(seconds_since_epoch, tz=ZoneInfo("UTC"))

    # Convert to target timezone if specified
    resolved_tz = _parse_timezone(timezone) if timezone else None
    if isinstance(resolved_tz, ZoneInfo):
        dt = dt.astimezone(resolved_tz)
    elif isinstance(resolved_tz, timedelta):
        dt = dt + resolved_tz
        # Strip the original tzinfo since we manually shifted
        dt = dt.replace(tzinfo=None)

    # Build format string
    format_str = "%Y-%m-%d"
    if include_day_of_week:
        format_str = f"%a{day_date_separator}" + format_str
    format_str += f"{date_time_seperator}%H:%M:%S"

    # Format the time
    formatted = dt.strftime(format_str)

    # Append timezone indicator
    if isinstance(resolved_tz, ZoneInfo):
        formatted += dt.strftime("%z")  # +HHMM or -HHMM format
    elif isinstance(resolved_tz, timedelta):
        total_minutes = int(resolved_tz.total_seconds()) // 60
        sign_char = "+" if total_minutes >= 0 else "-"
        abs_minutes = abs(total_minutes)
        formatted += f"{sign_char}{abs_minutes // 60:02d}{abs_minutes % 60:02d}"
    else:
        formatted += "Z"  # UTC/Zulu time

    return formatted


def sanitize_for_jsonb(input_dict):
    def sanitize_string(s):
        result = []
        for char in s:
            code_point = ord(char)
            if (
                (0 <= code_point <= 8)
                or (11 <= code_point <= 31)
                or (code_point == 127)
            ):
                result.append(f"\\u{code_point:04x}")
            else:
                result.append(char)
        return "".join(result)

    def sanitize_value(value):
        if isinstance(value, str):
            return sanitize_string(value)
        elif isinstance(value, dict):
            return {k: sanitize_value(v) for k, v in value.items()}
        elif isinstance(value, list):
            return [sanitize_value(v) for v in value]
        else:
            return value

    sanitized_dict = sanitize_value(input_dict)

    # Convert to JSON string and back to ensure proper escaping of quotes, backslashes, etc.
    return json.loads(json.dumps(sanitized_dict))


def generate_otp(length: int = 6) -> str:
    """Generate a random OTP code"""
    return "".join([str(secrets.randbelow(10)) for _ in range(length)])


class InMemoryKVStore:
    """In-memory key-value store with support for list operations."""

    def __init__(self):
        self._store: dict[str, Any] = {}

    def write(
        self, key: str, value: Any, append: bool = False, extend: bool = False
    ) -> None:
        """Write a value to the key-value store.

        Args:
            key: The key to store the value under
            value: The value to store
            append: If True, append the value to an existing list (creates list if key doesn't exist)
            extend: If True, extend an existing list with the value (creates list if key doesn't exist)
        """
        if append and extend:
            raise ValueError("Cannot use both append and extend at the same time")

        if append:
            # Append single value to a list
            existing = self._store.get(key, [])
            if not isinstance(existing, list):
                # Convert non-list to list if needed
                existing = [existing] if existing is not None else []
            self._store[key] = existing + [value]
        elif extend:
            # Extend list with multiple values
            existing = self._store.get(key, [])
            if not isinstance(existing, list):
                # Convert non-list to list if needed
                existing = [existing] if existing is not None else []
            if not isinstance(value, (list, tuple)):
                # Ensure value is iterable for extend
                value = [value]
            else:
                value = list(value)
            self._store[key] = existing + value
        else:
            # Default behavior: directly set the value
            self._store[key] = value

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from the key-value store.

        Args:
            key: The key to retrieve
            default: The default value if the key doesn't exist

        Returns:
            The value associated with the key, or the default value
        """
        return self._store.get(key, default)

    def delete(self, key: str) -> None:
        """Delete a key from the store if it exists."""
        self._store.pop(key, None)

    def clear(self) -> None:
        """Clear all entries from the store."""
        self._store.clear()

    def keys(self) -> list[str]:
        """Get all keys in the store."""
        return list(self._store.keys())

    def values(self) -> list[Any]:
        """Get all values in the store."""
        return list(self._store.values())

    def items(self) -> list[tuple[str, Any]]:
        """Get all key-value pairs in the store."""
        return list(self._store.items())

    def __contains__(self, key: str) -> bool:
        """Check if a key exists in the store."""
        return key in self._store

    def __getitem__(self, key: str) -> Any:
        """Get item using bracket notation."""
        return self._store[key]

    def __setitem__(self, key: str, value: Any) -> None:
        """Set item using bracket notation."""
        self._store[key] = value

    def __delitem__(self, key: str) -> None:
        """Delete item using bracket notation."""
        del self._store[key]

    def __iter__(self) -> Iterator[str]:
        """Iterate over keys."""
        return iter(self._store)

    def __len__(self) -> int:
        """Get the number of items in the store."""
        return len(self._store)

    def __repr__(self) -> str:
        """String representation of the store."""
        return f"InMemoryKVStore({self._store!r})"


class TodoStatus(str, Enum):
    """Status states for TodoItem."""

    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    WAITING = "waiting"  # Started but temporarily on hold
    BLOCKED = "blocked"  # Cannot start until blocking condition is resolved
    COMPLETED = "completed"
    CANCELLED = "cancelled"


class TodoItem(TransientStorageMixin, BaseModel):
    id: str = Field(default_factory=lambda: f"tdo_{uuid.uuid4().hex}")

    created_at: float = Field(default_factory=get_time)
    updated_at: float = Field(default_factory=get_time)

    task: str
    notes: str | None = None

    status: TodoStatus = TodoStatus.PENDING

    # Discriminator field for proper polymorphic serialization
    item_type: Literal["TodoItem"] = "TodoItem"

    # --- Engine bindings ---

    def bind_engine(
        self, engine: "TodoManagementEngine", add: bool = False
    ) -> "TodoItem":
        storage = self.get_or_create_transient_storage()
        storage.set("engine", engine)
        if add:
            engine.add_todo(todo=self)
        return self

    def get_engine(self) -> Optional["TodoManagementEngine"]:
        storage = self.get_or_create_transient_storage()
        return storage.get("engine", None)

    @property
    def engine(self):
        return self.get_engine()

    # --- Status marking ---

    def mark_status(self, new_status: TodoStatus | str) -> "TodoItem":
        # Convert string to enum if necessary
        if isinstance(new_status, str):
            new_status = TodoStatus(new_status)
        self.status = new_status
        self.updated_at = get_time()
        return self

    def mark_as_completed(self):
        return self.mark_status(TodoStatus.COMPLETED)

    def mark_as_pending(self):
        return self.mark_status(TodoStatus.PENDING)

    def mark_as_in_progress(self):
        return self.mark_status(TodoStatus.IN_PROGRESS)

    def mark_as_waiting(self):
        return self.mark_status(TodoStatus.WAITING)

    def mark_as_blocked(self):
        return self.mark_status(TodoStatus.BLOCKED)

    def mark_as_cancelled(self):
        return self.mark_status(TodoStatus.CANCELLED)

    # --- Engine ops ---

    def follow_up_with(self, item: "TodoItem"):
        if self.engine is None:
            raise ValueError("No engine is bound.")
        self.engine.add_todo(item, depends_on=self)
        return item

    def cancel(self, cascade: bool = True):
        if self.engine is None:
            raise ValueError("No engine is bound.")
        self.engine.cancel_todo(self, cascade=cascade)
        return self

    @property
    def follow_up_todo_items(self) -> List["TodoItem"]:
        """Returns a list of TodoItems that depend on this TodoItem (i.e., follow-up items)."""
        if self.engine is None:
            return []

        # Use the engine's get_follow_up_items method
        return self.engine.get_follow_up_items(self)

    @property
    def next_items(self) -> List["TodoItem"]:
        """Alias for follow_up_todo_items - returns items that depend on this TodoItem."""
        if self.engine is None:
            return []

        # Use the engine's get_follow_up_items method
        return self.engine.get_follow_up_items(self, exclude_not_runnable=True)

    @property
    def previous_items(self) -> List["TodoItem"]:
        """Returns a list of TodoItems that this TodoItem depends on (i.e., prerequisite items)."""
        if self.engine is None:
            return []

        # Use the engine's get_previous_items method
        return self.engine.get_previous_items(self)

    @classmethod
    def load(
        cls,
        name: str,
        *,
        version_string: str | None = None,
        prompt_dir: str,
        template_kwargs: dict[str, Any] | None = None,
    ):
        """Load a DecisionNode from a YAML file.

        Args:
            name: Name of the decision node file (without extension)
            version_string: Optional version string (e.g., "v0.0.1")
            prompt_dir: Directory containing the decision node files
            template_kwargs: Optional kwargs for rendering Jinja2 templates (required for .yml.j2 files)

        Returns:
            A DecisionNode instance loaded from the YAML file

        Raises:
            ValueError: If file extension is unsupported or if .j2 file is loaded without template_kwargs
        """
        from common.intel.utils import find_prompt_path

        path = find_prompt_path(
            name=name,
            prompt_dir=prompt_dir,
            version_string=version_string,
            extensions=[".yml", ".yml.j2"],
        )

        # Read the file content
        with open(path, "r") as f:
            content = f.read()

        # Handle Jinja templates
        if path.endswith(".yml.j2"):
            if template_kwargs is None:
                raise ValueError(
                    f"template_kwargs must be provided for Jinja template file: {path}"
                )
            # Render the Jinja template
            template = Template(content)
            content = template.render(**template_kwargs)
        elif not path.endswith(".yml"):
            raise ValueError(f"Unsupported file extension for decision node: {path}")

        # Parse YAML and validate with Pydantic
        d = yaml.safe_load(content)
        return cls.model_validate(d)


class DecisionBranch(BaseModel):
    name: str
    ids: List[str] = Field(default_factory=list)
    description: str | None = None


class DecisionNode(TodoItem):
    """A todo_ that represents a decision point with multiple branches.

    A DecisionNode completes by choosing one of several branches, which
    activates todos in that branch and cancels todos in other branches.
    """

    # Override discriminator for proper serialization
    item_type: Literal["DecisionNode"] = "DecisionNode"

    # Decision-specific fields - using DecisionBranch for clean organization
    branches: Dict[str, DecisionBranch] = Field(
        default_factory=dict
    )  # branch_name -> DecisionBranch
    chosen_branch: str | None = Field(default=None)

    @property
    def branch_names(self) -> List[str]:
        return list(self.branches.keys())

    def decide(self, branch: str, mark_as_completed: bool = True) -> "DecisionNode":
        """Make a decision by choosing a specific branch.

        This method must be called before marking the decision as completed.

        Args:
            branch: The name of the branch to choose.
            mark_as_completed: If True, automatically mark the decision as completed after choosing the branch.

        Raises:
            ValueError: If the branch is invalid or decision was already made.
            RuntimeError: If the decision node is already completed.
        """
        if self.status == TodoStatus.COMPLETED:
            raise RuntimeError(f"Decision '{self.task}' has already been completed")

        if self.chosen_branch is not None:
            raise RuntimeError(
                f"Decision '{self.task}' has already been decided (chosen: {self.chosen_branch})"
            )

        if branch not in self.branches:
            available = ", ".join(self.branches.keys())
            raise ValueError(f"Invalid branch '{branch}'. Available: {available}")

        self.chosen_branch = branch

        # Let the engine handle branch activation/cancellation
        if self.engine:
            self.engine.activate_decision_branch(self, branch)

        if mark_as_completed:
            self.mark_as_completed()

        return self

    def mark_as_completed(self):
        """Mark the decision as completed.

        A decision can only be completed after a branch has been chosen
        using the decide() method.

        Raises:
            RuntimeError: If no decision has been made yet.
        """
        if self.chosen_branch is None:
            raise RuntimeError(
                f"Cannot complete DecisionNode '{self.task}' without making a decision. "
                f"Call decide(branch) first with one of: {', '.join(self.branches.keys())}"
            )

        # Now mark as completed using parent's method
        return super().mark_as_completed()

    # Decision-specific properties
    @property
    def is_decided(self) -> bool:
        """Check if a decision has been made."""
        return self.status == TodoStatus.COMPLETED and self.chosen_branch is not None

    @property
    def available_branches(self) -> List[str]:
        """Get branches that can still be chosen."""
        if self.is_decided:
            return []
        return list(self.branches.keys())

    @property
    def chosen_branch_todos(self) -> List[str]:
        """Get the todo_ IDs for the chosen branch."""
        if self.chosen_branch and self.chosen_branch in self.branches:
            return self.branches[self.chosen_branch].ids
        return []

    # Helper methods for building branches
    def add_to_branch(
        self, branch: str, todos: TodoItem | List[TodoItem]
    ) -> "DecisionNode":
        """Add todos to a specific branch of this decision.

        Args:
            branch: The branch name to add todos to
            todos: A TodoItem or list of TodoItems to add to the branch

        Returns:
            Self for method chaining

        Raises:
            ValueError: If the decision is already completed or branch doesn't exist
        """
        if self.is_decided:
            raise ValueError("Cannot add todos to a decided decision")

        # Ensure todos is a list
        if isinstance(todos, TodoItem):
            todos = [todos]

        # Create branch if it doesn't exist
        if branch not in self.branches:
            self.branches[branch] = DecisionBranch(name=branch)

        # If engine is bound, add todos to it
        if self.engine:
            for todo in todos:
                # Set todo to BLOCKED status
                todo.status = TodoStatus.BLOCKED
                # Add to engine
                self.engine.add_todo(todo)
                # Add ID to branch
                self.branches[branch].ids.append(todo.id)
        else:
            # If no engine, just store the todo IDs
            # (assumes todos will be added to engine later)
            for todo in todos:
                todo.status = TodoStatus.BLOCKED
                self.branches[branch].ids.append(todo.id)

        return self

    def create_branch(
        self, branch: str, description: str | None = None
    ) -> "DecisionNode":
        """Create a new empty branch for this decision.

        Args:
            branch: The name of the branch to create
            description: Optional description of what this branch represents

        Returns:
            Self for method chaining

        Raises:
            ValueError: If the branch already exists or decision is completed
        """
        if self.is_decided:
            raise ValueError("Cannot create branch on a decided decision")

        if branch in self.branches:
            raise ValueError(f"Branch '{branch}' already exists")

        self.branches[branch] = DecisionBranch(name=branch, description=description)

        return self

    # Make follow_up_with work naturally
    def follow_up_with(self, item: "TodoItem"):
        """Add a follow-up that depends on this decision being made.

        The follow-up will become runnable after any branch is chosen.
        """
        if self.engine is None:
            raise ValueError("No engine is bound.")

        # The follow-up depends on this decision being completed
        self.engine.add_todo(item, depends_on=self)
        return item


class TodoManagementEngine(BaseModel):
    # Use Field with discriminator for polymorphic deserialization
    # Note: DecisionNode must come before TodoItem in the Union for proper discrimination
    todo_list: List[Union[DecisionNode, TodoItem]] = Field(
        default=[], discriminator="item_type"
    )
    dependencies: Dict[str, List[str]] = {}

    @model_validator(mode="after")
    def bind_initial_todos(self) -> "TodoManagementEngine":
        """Bind engine to all todos in the initial todo_list."""
        for todo in self.todo_list:
            if todo.engine is None:
                todo.bind_engine(self)
        return self

    def get_todo_by_id(self, todo_id: str) -> Optional[TodoItem]:
        for todo in self.todo_list:
            if todo.id == todo_id:
                return todo
        return None

    def get_follow_up_items(
        self, todo: TodoItem | DecisionNode | str, exclude_not_runnable: bool = False
    ) -> List[TodoItem | DecisionNode]:
        """Returns a list of TodoItems that depend on the given TodoItem (i.e., follow-up items).

        Args:
            todo: Either a TodoItem instance or a todo ID string
            exclude_not_runnable: If True, exclude todos that are not runnable (i.e., have incomplete dependencies)

        Returns:
            List of TodoItem instances that depend on the given todo_
        """
        todo_id = todo.id if isinstance(todo, TodoItem) else todo

        follow_ups = []
        # Find all todos that have this todo_'s ID in their dependency list
        for dependent_todo_id, dependency_ids in self.dependencies.items():
            if todo_id in dependency_ids:
                # This todo_ depends on the given todo_, so it's a follow-up
                dependent_todo = self.get_todo_by_id(dependent_todo_id)
                if dependent_todo is not None:
                    # Check if we should exclude not runnable todos
                    if exclude_not_runnable:
                        # A todo_ is not runnable if any of its
                        # dependencies are not completed
                        all_deps_complete = True
                        for dependency_id in dependency_ids:
                            dep_todo = self.get_todo_by_id(dependency_id)
                            if dep_todo and dep_todo.status != TodoStatus.COMPLETED:
                                all_deps_complete = False
                                break

                        # Skip this todo_ if not all dependencies are complete
                        if not all_deps_complete:
                            continue

                    follow_ups.append(dependent_todo)

        return follow_ups

    def get_previous_items(
        self, todo: TodoItem | DecisionNode | str
    ) -> List[TodoItem | DecisionNode]:
        """Returns a list of TodoItems that the given TodoItem depends on (i.e., prerequisite items).

        Args:
            todo: Either a TodoItem instance or a todo ID string

        Returns:
            List of TodoItem instances that the given todo_ depends on
        """
        todo_id = todo.id if isinstance(todo, TodoItem) else todo

        previous_items = []
        # Get the dependency IDs for this todo_
        if todo_id in self.dependencies:
            dependency_ids = self.dependencies[todo_id]
            # Convert dependency IDs to TodoItem instances
            for dependency_id in dependency_ids:
                dep_todo = self.get_todo_by_id(dependency_id)
                if dep_todo is not None:
                    previous_items.append(dep_todo)

        return previous_items

    def in_todo_list(self, todo: TodoItem | DecisionNode | str) -> bool:
        todo_id = todo.id if isinstance(todo, TodoItem) else todo
        return any(t.id == todo_id for t in self.todo_list)

    def add_todo(
        self,
        todo: TodoItem | DecisionNode,
        depends_on: Optional[
            Union[List[TodoItem | DecisionNode], TodoItem | DecisionNode]
        ] = None,
    ) -> "TodoManagementEngine":
        if not self.in_todo_list(todo):
            self.todo_list.append(todo)
            todo.bind_engine(self)

        if depends_on is not None:
            # Record dependencies - dependencies maps todo_id to list of todo_ids it depends on
            if isinstance(depends_on, TodoItem):
                depends_on = [depends_on]

            # Make sure all dependencies are in the todo_ list
            for dep in depends_on:
                if not self.in_todo_list(dep):
                    self.todo_list.append(dep)
                dep.bind_engine(self)

            # Get the IDs of the dependencies
            dependency_ids = [dep.id for dep in depends_on]

            # Record in dependencies graph
            if todo.id in self.dependencies:
                # Merge with existing dependencies, avoiding duplicates
                existing = set(self.dependencies[todo.id])
                existing.update(dependency_ids)
                self.dependencies[todo.id] = list(existing)
            else:
                self.dependencies[todo.id] = dependency_ids
        return self

    def cancel_todo(self, todo: TodoItem | DecisionNode, cascade: bool = True):
        # Make sure todo_ is bound to this engine
        todo.bind_engine(self)
        # Mark the todo_ as cancelled (don't remove from list)
        todo.mark_as_cancelled()

        # Find all todos that depend on this one
        dependent_todo_ids = []
        for todo_id, deps in self.dependencies.items():
            if todo.id in deps:
                dependent_todo_ids.append(todo_id)

        if cascade:
            # Recursively cancel all dependent todos
            for dependent_id in dependent_todo_ids:
                dependent_todo = self.get_todo_by_id(dependent_id)
                if dependent_todo and dependent_todo.status != TodoStatus.CANCELLED:
                    self.cancel_todo(dependent_todo, cascade=True)
        else:
            # Remove the dependency link for todos that depended on this one
            for dependent_id in dependent_todo_ids:
                if dependent_id in self.dependencies:
                    self.dependencies[dependent_id] = [
                        dependency_id
                        for dependency_id in self.dependencies[dependent_id]
                        if dependency_id != todo.id
                    ]
                    # Clean up empty dependency lists
                    if not self.dependencies[dependent_id]:
                        del self.dependencies[dependent_id]

        # Clean up this todo_'s dependencies if it had any
        if todo.id in self.dependencies:
            del self.dependencies[todo.id]

    def get_todos(
        self,
        status_in: List[TodoStatus] | None = None,
        include_dependents: bool = False,
        include_completed: bool = False,
        include_cancelled: bool = False,
        include_waiting: bool = False,
        include_blocked: bool = False,
    ) -> List[TodoItem | DecisionNode]:
        result = []

        for todo in self.todo_list:
            # Filter out cancelled todos unless explicitly included
            if todo.status == TodoStatus.CANCELLED and not include_cancelled:
                continue

            # Filter out completed todos unless explicitly included
            if todo.status == TodoStatus.COMPLETED and not include_completed:
                continue

            # Filter out waiting todos unless explicitly included
            if todo.status == TodoStatus.WAITING and not include_waiting:
                continue

            # Filter out blocked todos unless explicitly included
            if todo.status == TodoStatus.BLOCKED and not include_blocked:
                continue

            # Check dependencies
            if todo.id in self.dependencies and not include_dependents:
                # This todo_ has dependencies - check if all are complete
                dependencies = self.dependencies[todo.id]
                all_deps_complete = True

                for dependency_id in dependencies:
                    dep_todo = self.get_todo_by_id(dependency_id)
                    if dep_todo and dep_todo.status != TodoStatus.COMPLETED:
                        all_deps_complete = False
                        break

                # Skip this todo_ if not all dependencies are complete
                if not all_deps_complete:
                    continue

            result.append(todo)

        # Filter by status if specified
        if status_in is not None:
            result = [todo for todo in result if todo.status in status_in]

        for todo in result:
            todo.bind_engine(self)

        return result

    def add_decision(
        self,
        decision: DecisionNode,
        branches: Dict[str, List[TodoItem]] | None = None,
        depends_on: Optional[Union[List[TodoItem], TodoItem]] = None,
    ) -> DecisionNode:
        """Add a decision point with multiple branches.

        Args:
            decision: The DecisionNode to add
            branches: Optional dictionary mapping branch names to lists of TodoItems for that branch.
                     If not provided, assumes decision.branches is already configured with DecisionBranch objects.
            depends_on: Optional dependencies that must complete before the decision

        Returns:
            The DecisionNode that was added
        """
        # If branches are provided, set up the branch todos
        if branches is not None:
            # Clear existing branches and create new DecisionBranch objects
            decision.branches = {}

            for branch_name, branch_todos in branches.items():
                # Create a new DecisionBranch
                branch = DecisionBranch(name=branch_name)

                # Add all todos for this branch
                for todo in branch_todos:
                    # Set todos to BLOCKED status - they can't run until decision is made
                    todo.status = TodoStatus.BLOCKED
                    self.add_todo(todo)
                    branch.ids.append(todo.id)

                # Add the branch to the decision
                decision.branches[branch_name] = branch
        else:
            # Ensure all branch todos are already in the engine and set to BLOCKED
            for branch_name, branch in decision.branches.items():
                for todo_id in branch.ids:
                    todo = self.get_todo_by_id(todo_id)
                    if todo and todo.status != TodoStatus.BLOCKED:
                        todo.status = TodoStatus.BLOCKED

        # Add the decision itself with any dependencies
        self.add_todo(decision, depends_on=depends_on)

        return decision

    def activate_decision_branch(
        self, decision: "DecisionNode", branch: str
    ) -> "DecisionNode":
        """Activate todos in the chosen branch and cancel todos in other branches.

        Internal method called by DecisionNode.mark_as_completed().

        Args:
            decision: The DecisionNode being completed
            branch: The chosen branch name
        """
        # Activate todos in the chosen branch (BLOCKED → PENDING)
        if branch in decision.branches:
            for todo_id in decision.branches[branch].ids:
                todo = self.get_todo_by_id(todo_id)
                if todo and todo.status == TodoStatus.BLOCKED:
                    todo.mark_as_pending()

        # Cancel todos in other branches (BLOCKED → CANCELLED)
        for other_branch_name, other_branch in decision.branches.items():
            if other_branch_name != branch:
                for todo_id in other_branch.ids:
                    todo = self.get_todo_by_id(todo_id)
                    if todo and todo.status == TodoStatus.BLOCKED:
                        # Use cancel_todo to handle dependency cascading
                        self.cancel_todo(todo, cascade=True)

        return decision

    def complete_todo(
        self, todo: TodoItem | str, branch: str | None = None
    ) -> "TodoManagementEngine":
        """Complete a todo, with special handling for DecisionNodes.

        Args:
            todo: The todo to complete (or its ID)
            branch: For DecisionNodes, the branch to choose

        Returns:
            The engine for chaining
        """
        # Get the actual todo object if an ID was passed
        if isinstance(todo, str):
            todo = self.get_todo_by_id(todo)
            if todo is None:
                raise ValueError(f"Todo with ID {todo} not found")

        # Handle DecisionNodes specially
        if isinstance(todo, DecisionNode):
            if branch is None:
                raise ValueError(
                    "DecisionNode requires a branch choice. Use decide(branch) first or provide branch parameter."
                )
            todo.decide(branch, mark_as_completed=True)
        else:
            if branch is not None:
                raise ValueError("Branch parameter only valid for DecisionNodes")
            todo.mark_as_completed()

        return self

    def get_decision_by_id(self, decision_id: str) -> Optional[DecisionNode]:
        """Get a DecisionNode by its ID.

        Args:
            decision_id: The ID of the decision

        Returns:
            The DecisionNode if found and is a decision, None otherwise
        """
        todo = self.get_todo_by_id(decision_id)
        if isinstance(todo, DecisionNode):
            return todo
        return None

    def get_decisions(
        self, include_completed: bool = False, include_cancelled: bool = False
    ) -> List[DecisionNode]:
        """Get all DecisionNodes in the todo_ list.

        Args:
            include_completed: Whether to include decided decisions
            include_cancelled: Whether to include cancelled decisions

        Returns:
            List of DecisionNodes
        """
        decisions = []
        for todo in self.todo_list:
            if isinstance(todo, DecisionNode):
                if todo.status == TodoStatus.COMPLETED and not include_completed:
                    continue
                if todo.status == TodoStatus.CANCELLED and not include_cancelled:
                    continue
                decisions.append(todo)
        return decisions
