from __future__ import annotations

import ast
from typing import TypeGuard


def to_cap_words(snake_case: str) -> str:
  """Convert snake_case to CapWords (PascalCase)."""
  return "".join(
    word[0].upper() + word[1:] if word else "" for word in snake_case.split("_")
  )


def convert_runtime_attr_to_stub_type(type_str: str) -> str:
  """Convert runtime attribute references to stub type names.

  In stub files, runtime attributes like Foo.Config cannot be used as types.
  We need to convert them to their actual class names like FooConfig.

  Args:
    type_str: Type annotation string (e.g., "Optimizer.Config")

  Returns:
    Stub-compatible type string (e.g., "OptimizerConfig")
  """
  # In the new nested config system, T.Config is a valid type in stubs.
  # So we don't need to convert it to TConfig anymore.
  return type_str


def extract_inner_type(annotation: ast.Subscript) -> str:
  """Extract the inner type T from Hyper[T] or Hyper[T, constraints...].

  Args:
    annotation: AST Subscript node representing Hyper[...]

  Returns:
    String representation of the inner type T, converted for stub compatibility
  """
  slice_node = annotation.slice
  if isinstance(slice_node, ast.Tuple):
    # Hyper[T, Constraint1, ...] - extract just T
    type_str = ast.unparse(slice_node.elts[0])
  else:
    type_str = ast.unparse(slice_node)

  # Convert runtime attributes to stub type names
  return convert_runtime_attr_to_stub_type(type_str)


def extract_default_value(default: ast.expr | None) -> str:
  """Extract default value as string."""
  if default is None:
    return "..."
  return ast.unparse(default)


def is_hyper_annotated(annotation: ast.expr | None) -> TypeGuard[ast.Subscript]:
  """Check if annotation is Hyper[T]."""
  return (
    isinstance(annotation, ast.Subscript)
    and isinstance(annotation.value, ast.Name)
    and annotation.value.id == "Hyper"
  )


def is_configurable_decorated(
  node: ast.FunctionDef | ast.ClassDef,
) -> bool:
  """Check if function or class has @configurable decorator."""
  for decorator in node.decorator_list:
    if isinstance(decorator, ast.Name) and decorator.id == "configurable":
      return True
  return False


def is_dataclass_decorated(class_node: ast.ClassDef) -> bool:
  """Check if class has @dataclass decorator."""
  for decorator in class_node.decorator_list:
    if isinstance(decorator, ast.Name) and decorator.id == "dataclass":
      return True
  return False
