from __future__ import annotations

import ast
from typing import TYPE_CHECKING

if TYPE_CHECKING:
  from hipr.stubs.models import FunctionInfo


def extract_imports(source_tree: ast.Module) -> tuple[list[str], list[str]]:
  """Extract import statements from source tree.

  Args:
    source_tree: Parsed AST module

  Returns:
    Tuple of (future_imports, other_imports)
  """
  future_imports: list[str] = []
  other_imports: list[str] = []

  for node in source_tree.body:
    if isinstance(node, (ast.Import, ast.ImportFrom)):
      code = ast.unparse(node)
      if isinstance(node, ast.ImportFrom) and node.module == "__future__":
        future_imports.append(code)
      else:
        other_imports.append(code)
    elif isinstance(node, ast.If):
      # Check for TYPE_CHECKING blocks and include them
      # This handles: if TYPE_CHECKING:, if typing.TYPE_CHECKING:, etc.
      test_source = ast.unparse(node.test)
      if "TYPE_CHECKING" in test_source:
        other_imports.append(ast.unparse(node))

  return future_imports, other_imports


def has_default_usage(source_tree: ast.Module) -> bool:
  """Check if DEFAULT sentinel is used in source code.

  Args:
    source_tree: Parsed AST module

  Returns:
    True if DEFAULT is used, False otherwise
  """
  for node in ast.walk(source_tree):
    if isinstance(node, ast.Name) and node.id == "DEFAULT":
      return True
  return False


def collect_used_names_from_stub(stub_content: str) -> set[str]:
  """Collect all names referenced in stub content.

  Args:
    stub_content: Generated stub content (without header/imports)

  Returns:
    Set of names used in the stub (type annotations, base classes, etc.)
  """
  used_names: set[str] = set()

  try:
    tree = ast.parse(stub_content)
  except SyntaxError:
    # If we can't parse, return empty set (will keep all imports)
    return used_names

  for node in ast.walk(tree):
    # Collect simple names (e.g., DataFrame, Series)
    if isinstance(node, ast.Name):
      used_names.add(node.id)
    # Collect attribute access (e.g., pd.DataFrame -> "pd")
    elif isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
      used_names.add(node.value.id)

  return used_names


def _process_import_node(
  node: ast.Import, import_stmt: str, used_names: set[str]
) -> str | None:
  """Process 'import x' statements.

  Returns:
    Import statement if used, None otherwise
  """
  for alias in node.names:
    name_to_check = alias.asname if alias.asname else alias.name
    if name_to_check in used_names or alias.name.split(".")[0] in used_names:
      return import_stmt
  return None


def _process_import_from_node(
  node: ast.ImportFrom, import_stmt: str, used_names: set[str]
) -> str | None:
  """Process 'from x import y' statements.

  Returns:
    Filtered import statement if any names are used, None otherwise
  """
  # Keep wildcard imports
  if node.names[0].name == "*":
    return import_stmt

  # Collect used aliases
  used_aliases: list[ast.alias] = []
  for alias in node.names:
    name_to_check = alias.asname if alias.asname else alias.name
    # Skip DEFAULT and MakeableModel when imported from 'hipr'
    if node.module == "hipr" and alias.name in {"DEFAULT", "MakeableModel"}:
      continue
    if name_to_check in used_names:
      used_aliases.append(alias)

  if not used_aliases:
    return None

  # Return original if all names used, otherwise create filtered import
  if len(used_aliases) == len(node.names):
    return import_stmt

  module = node.module or ""
  name_parts = [
    f"{alias.name} as {alias.asname}" if alias.asname else alias.name
    for alias in used_aliases
  ]
  return f"from {module} import {', '.join(name_parts)}"


def filter_imports_by_usage(
  imports: list[str],
  used_names: set[str],
) -> list[str]:
  """Filter imports to only those used in the stub.

  For multi-name imports like 'from x import a, b, c', only keeps the names
  that are actually used in the stub.

  Args:
    imports: List of import statements
    used_names: Set of names actually used in stub

  Returns:
    Filtered list of import statements
  """
  filtered: list[str] = []

  for import_stmt in imports:
    # Always keep __future__ imports and TYPE_CHECKING blocks
    if "__future__" in import_stmt or "TYPE_CHECKING" in import_stmt:
      filtered.append(import_stmt)
      continue

    try:
      tree = ast.parse(import_stmt)
      node = tree.body[0]

      result = None
      if isinstance(node, ast.Import):
        result = _process_import_node(node, import_stmt, used_names)
      elif isinstance(node, ast.ImportFrom):
        result = _process_import_from_node(node, import_stmt, used_names)

      if result:
        filtered.append(result)

    except SyntaxError:
      # If we can't parse, keep the import to be safe
      filtered.append(import_stmt)

  return filtered


def _extract_names_from_default(default_value: str) -> set[str]:
  """Extract Name nodes from a default value string.

  Args:
    default_value: String representation of default value

  Returns:
    Set of names referenced in the default value
  """
  names: set[str] = set()
  try:
    default_tree = ast.parse(default_value, mode="eval")
    for node in ast.walk(default_tree):
      if isinstance(node, ast.Name):
        names.add(node.id)
  except SyntaxError:
    pass
  return names


def collect_constant_references(functions: list[FunctionInfo]) -> set[str]:
  """Collect names of constants referenced in default values.

  Args:
    functions: List of FunctionInfo objects

  Returns:
    Set of constant names referenced in defaults (e.g., {"A", "B"})
  """
  constants: set[str] = set()

  for func in functions:
    # Check hyperparameters
    for param in func.params:
      constants.update(_extract_names_from_default(param.default_value))

    # Check all parameters (for non-Hyper parameters)
    for _, _, default_val in func.all_params:
      if default_val and default_val != "...":
        constants.update(_extract_names_from_default(default_val))

  return constants


def extract_module_constants(
  source_tree: ast.Module, constant_names: set[str]
) -> list[str]:
  """Extract module-level constant assignments from source.

  Args:
    source_tree: Parsed AST module
    constant_names: Set of constant names to extract

  Returns:
    List of constant assignment strings (e.g., ["A = 3", "B = 'hello'"])
  """
  constants: list[str] = []

  for node in source_tree.body:
    # Look for simple assignments: A = value
    if isinstance(node, ast.Assign):
      for target in node.targets:
        if isinstance(target, ast.Name) and target.id in constant_names:
          # Found a constant assignment
          constants.append(ast.unparse(node))
          break

    # Look for annotated assignments: A: int = value
    if (
      isinstance(node, ast.AnnAssign)
      and isinstance(node.target, ast.Name)
      and node.target.id in constant_names
    ):
      constants.append(ast.unparse(node))

  return constants


def collect_type_references(functions: list[FunctionInfo]) -> set[str]:
  """Collect all type names referenced in function signatures.

  Args:
    functions: List of FunctionInfo objects

  Returns:
    Set of type names used in signatures
  """
  type_names: set[str] = set()

  for func in functions:
    # Collect from return type
    try:
      return_tree = ast.parse(func.return_type, mode="eval")
      for node in ast.walk(return_tree):
        if isinstance(node, ast.Name):
          type_names.add(node.id)
    except SyntaxError:
      pass

    # Collect from all parameters
    for _, param_type, _ in func.all_params:
      try:
        type_tree = ast.parse(param_type, mode="eval")
        for node in ast.walk(type_tree):
          if isinstance(node, ast.Name):
            type_names.add(node.id)
      except SyntaxError:
        pass

  # Filter out built-in types and common imports
  builtin_types = {
    "int",
    "float",
    "str",
    "bool",
    "None",
    "list",
    "dict",
    "set",
    "tuple",
    "Any",
    "Optional",
    "Union",
  }
  return type_names - builtin_types


def collect_type_references_from_ast(func_nodes: list[ast.FunctionDef]) -> set[str]:
  """Collect type references from AST function nodes.

  Args:
    func_nodes: List of AST FunctionDef nodes

  Returns:
    Set of type names referenced in function signatures
  """
  type_names: set[str] = set()

  for func_node in func_nodes:
    # Collect from return type annotation
    if func_node.returns:
      for node in ast.walk(func_node.returns):
        if isinstance(node, ast.Name):
          type_names.add(node.id)

    # Collect from parameter annotations
    all_args = (
      func_node.args.args + func_node.args.posonlyargs + func_node.args.kwonlyargs
    )
    for arg in all_args:
      if arg.annotation:
        for node in ast.walk(arg.annotation):
          if isinstance(node, ast.Name):
            type_names.add(node.id)

  # Filter out built-in types
  builtin_types = {
    "int",
    "float",
    "str",
    "bool",
    "None",
    "list",
    "dict",
    "set",
    "tuple",
    "Any",
    "Optional",
    "Union",
  }
  return type_names - builtin_types


def _extract_class_definition(
  node: ast.ClassDef, configurable_names: set[str]
) -> tuple[str, set[str]] | None:
  """Extract class definition for stub file.

  Returns:
    Tuple of (definition string, base class refs) or None if should skip
  """
  # Skip @configurable classes - they're handled by Config generation
  if node.name in configurable_names:
    return None

  base_class_refs: set[str] = set()
  for base in node.bases:
    if isinstance(base, ast.Name):
      base_class_refs.add(base.id)
    elif isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name):
      base_class_refs.add(base.value.id)

  # Remove decorators for stub files
  node_copy = ast.ClassDef(
    name=node.name,
    bases=node.bases,
    keywords=node.keywords,
    body=node.body,
    decorator_list=[],  # Remove decorators
    type_params=getattr(node, "type_params", []),  # Python 3.12+
    lineno=node.lineno,
    col_offset=node.col_offset,
  )
  return ast.unparse(node_copy), base_class_refs


def extract_local_type_definitions(
  source_tree: ast.Module, type_names: set[str], configurable_names: set[str]
) -> tuple[list[str], set[str]]:
  """Extract local class and type alias definitions from source.

  Args:
    source_tree: Parsed AST module
    type_names: Set of type names to extract
    configurable_names: Set of @configurable names to exclude (handled separately)

  Returns:
    Tuple of (type definition strings, base class names that need importing)
  """
  definitions: list[str] = []
  base_class_refs: set[str] = set()

  for node in source_tree.body:
    # Extract class definitions (without decorators for stub files)
    if isinstance(node, ast.ClassDef) and node.name in type_names:
      result = _extract_class_definition(node, configurable_names)
      if result:
        definitions.append(result[0])
        base_class_refs.update(result[1])

    # Extract type alias (TypeAlias annotation or simple assignment)
    elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
      if node.target.id in type_names:
        definitions.append(ast.unparse(node))

    # Extract simple type aliases (e.g., MyType = dict[str, int])
    elif isinstance(node, ast.Assign):
      for target in node.targets:
        if isinstance(target, ast.Name) and target.id in type_names:
          definitions.append(ast.unparse(node))
          break

  return definitions, base_class_refs


def extract_public_functions(
  source_tree: ast.Module, configurable_names: set[str]
) -> tuple[list[str], list[ast.FunctionDef]]:
  """Extract public function signatures (non-configurable functions).

  In stub files, ALL public functions must be included, not just @configurable ones,
  because .pyi files override .py files completely for type checkers.

  Args:
    source_tree: Parsed AST module
    configurable_names: Set of @configurable function names to skip

  Returns:
    Tuple of (function signature strings, original AST nodes for type collection)
  """
  function_sigs: list[str] = []
  func_nodes: list[ast.FunctionDef] = []

  for node in source_tree.body:
    if isinstance(node, ast.FunctionDef):
      # Skip private functions and @configurable ones (handled separately)
      if node.name.startswith("_") or node.name in configurable_names:
        continue

      # Store original node for type reference collection
      func_nodes.append(node)

      # Generate stub signature (function signature with ... body)
      # Remove decorators and replace body with ...
      stub_func = ast.FunctionDef(
        name=node.name,
        args=node.args,
        body=[ast.Expr(value=ast.Constant(value=Ellipsis))],
        decorator_list=[],
        returns=node.returns,
        type_comment=node.type_comment,
        type_params=getattr(node, "type_params", []),
        lineno=node.lineno,
        col_offset=node.col_offset,
      )
      function_sigs.append(ast.unparse(stub_func))

  return function_sigs, func_nodes


def extract_public_classes(  # noqa: C901, PLR0912
  source_tree: ast.Module, configurable_names: set[str]
) -> tuple[list[str], set[str]]:
  """Extract all public class definitions for stub file.

  In stub files, ALL public classes must be included, not just those
  referenced in @configurable signatures, because .pyi files override
  .py files completely for type checkers.

  Args:
    source_tree: Parsed AST module
    configurable_names: Set of @configurable class names to skip

  Returns:
    Tuple of (class definition strings, base class names that need importing)
  """
  class_defs: list[str] = []
  base_class_refs: set[str] = set()

  for node in source_tree.body:
    if isinstance(node, ast.ClassDef):
      # Skip private classes
      if node.name.startswith("_"):
        continue

      # Skip @configurable classes - they're handled by Config generation
      if node.name in configurable_names:
        continue

      # Extract base class references for import
      for base in node.bases:
        if isinstance(base, ast.Name):
          base_class_refs.add(base.id)
        elif isinstance(base, ast.Attribute) and isinstance(base.value, ast.Name):
          base_class_refs.add(base.value.id)

      # Generate stub class (remove decorators, keep body with ... for methods)
      stub_body: list[ast.stmt] = []
      for item in node.body:
        if isinstance(item, ast.FunctionDef):
          # Replace method body with ...
          stub_method = ast.FunctionDef(
            name=item.name,
            args=item.args,
            body=[ast.Expr(value=ast.Constant(value=Ellipsis))],
            decorator_list=[],
            returns=item.returns,
            type_comment=item.type_comment,
            type_params=getattr(item, "type_params", []),
            lineno=item.lineno,
            col_offset=item.col_offset,
          )
          stub_body.append(stub_method)
        elif isinstance(item, ast.AnnAssign):
          # Keep type annotations
          stub_body.append(item)
        elif isinstance(item, ast.Assign):
          # Keep assignments (e.g., class attributes)
          stub_body.append(item)
        elif isinstance(item, ast.Pass | ast.Expr):
          # Keep pass and docstrings
          stub_body.append(item)

      # If body is empty, add pass
      if not stub_body:
        stub_body = [ast.Pass()]

      stub_class = ast.ClassDef(
        name=node.name,
        bases=node.bases,
        keywords=node.keywords,
        body=stub_body,
        decorator_list=[],  # Remove decorators
        type_params=getattr(node, "type_params", []),
        lineno=node.lineno,
        col_offset=node.col_offset,
      )
      class_defs.append(ast.unparse(stub_class))

  return class_defs, base_class_refs
