from __future__ import annotations

import ast


def extract_constraint_value(node: ast.expr) -> float | int | str | None:
  """Extract constraint value from AST node.

  Handles Ge[10], Le[100], Pattern[r"..."], etc.
  Returns the value inside the brackets, or None if not extractable.
  """
  # Handle subscript like Ge[10]
  if isinstance(node, ast.Subscript):
    slice_node = node.slice
    # Numeric constraints: Ge[10], Le[100.5]
    if isinstance(slice_node, ast.Constant):
      return slice_node.value  # type: ignore[return-value]
    # String constraints: Pattern[r"..."]
    if isinstance(slice_node, ast.Constant) and isinstance(slice_node.value, str):
      return slice_node.value
  return None


def extract_constraints_from_annotation(
  annotation: ast.Subscript,
) -> dict[str, float | int | str]:
  """Extract constraint values from Hyper[T, Ge[2], Le[100]] annotation.

  Returns dict like {"ge": 2, "le": 100}
  """
  constraints: dict[str, float | int | str] = {}
  slice_node = annotation.slice

  if not isinstance(slice_node, ast.Tuple):
    return constraints

  # Process constraint metadata (skip first element which is the type)
  for constraint_node in slice_node.elts[1:]:
    if not isinstance(constraint_node, ast.Subscript):
      continue
    if not isinstance(constraint_node.value, ast.Name):
      continue

    constraint_name = constraint_node.value.id
    constraint_value = extract_constraint_value(constraint_node)

    if constraint_value is None:
      continue

    constraint_map = {
      "Ge": "ge",
      "Gt": "gt",
      "Le": "le",
      "Lt": "lt",
      "MinLen": "min_length",
      "MaxLen": "max_length",
      "MultipleOf": "multiple_of",
      "Pattern": "pattern",
    }

    field_name = constraint_map.get(constraint_name)
    if field_name:
      constraints[field_name] = constraint_value

  return constraints
