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 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
  """
  slice_node = annotation.slice
  if isinstance(slice_node, ast.Tuple):
    # Hyper[T, Constraint1, ...] - extract just T
    return ast.unparse(slice_node.elts[0])
  return ast.unparse(slice_node)


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


def transform_type_annotation(
  type_ann: str,
  configurable_names: set[str],
  primitive_types: set[str],
  use_config: bool,
) -> str:
  """Transform type annotation to use .Config or component type.

  Args:
    type_ann: The type annotation string to transform
    configurable_names: Set of names of configurable classes/functions
    primitive_types: Set of primitive types to ignore (not transform)
    use_config: Whether to append .Config to the type name

  Returns:
    Transformed type string
  """
  clean_ann = type_ann.strip("'\"")

  # Already a Config type
  if ".Config" in clean_ann or clean_ann.endswith("Config"):
    return clean_ann

  # Handle .Type suffix (e.g., my_func.Type -> my_func[.Config])
  # Only transform if the base is a known configurable
  if clean_ann.endswith(".Type") and len(clean_ann) > len(".Type"):
    base = clean_ann.removesuffix(".Type")
    if base in configurable_names or base.isidentifier():
      return f"{base}.Config" if use_config else base

  # Handle local configurables or external identifiers
  if clean_ann in configurable_names or (
    clean_ann.isidentifier() and clean_ann not in primitive_types
  ):
    # Even for instance params, we often want the component type (which is the same name)
    # But if use_config is True, we append .Config
    if use_config:
      return f"{clean_ann}.Config"
    # For instance params, we keep original name (Component type)
    return clean_ann

  return clean_ann
