#!/usr/bin/env python3
"""Generate .pyi stub files for @configurable decorated functions and classes.

This script scans Python files for @configurable decorators on functions,
methods, classes (with __init__), and dataclasses. It generates corresponding
type stub files with TypedDict configs and ConfigurableIndicator wrappers.
"""

from __future__ import annotations

import ast
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import NamedTuple, TypeGuard

from loguru import logger

# Maximum file size to process (10MB)
MAX_FILE_SIZE = 10 * 1024 * 1024


class HyperParam(NamedTuple):
  """Represents a hyperparameter extracted from a function signature."""

  name: str
  type_annotation: str
  default_value: str


@dataclass
class FunctionInfo:
  """Information about a @configurable decorated function."""

  name: str
  params: list[HyperParam]
  all_params: list[tuple[str, str, str]]  # (name, type, default)
  return_type: str
  has_data_param: bool


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 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_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


def validate_constraint_conflicts(
  constraints: dict[str, float | int | str], param_name: str, func_name: str
) -> None:
  """Validate that constraints don't conflict (AST version).

  Raises:
    ValueError: If constraints are contradictory
  """
  # Check numeric bound conflicts
  ge_val = constraints.get("ge")
  gt_val = constraints.get("gt")
  le_val = constraints.get("le")
  lt_val = constraints.get("lt")

  # Determine effective lower and upper bounds
  lower_bound: float | int | None = None
  lower_inclusive = True
  if ge_val is not None and isinstance(ge_val, (int, float)):
    lower_bound = ge_val
    lower_inclusive = True
  if (
    gt_val is not None
    and isinstance(gt_val, (int, float))
    and (lower_bound is None or gt_val >= lower_bound)
  ):
    lower_bound = gt_val
    lower_inclusive = False

  upper_bound: float | int | None = None
  upper_inclusive = True
  if le_val is not None and isinstance(le_val, (int, float)):
    upper_bound = le_val
    upper_inclusive = True
  if (
    lt_val is not None
    and isinstance(lt_val, (int, float))
    and (upper_bound is None or lt_val <= upper_bound)
  ):
    upper_bound = lt_val
    upper_inclusive = False

  # Check if bounds are impossible
  if lower_bound is not None and upper_bound is not None:
    if lower_bound > upper_bound:
      msg = (
        f"In function '{func_name}', parameter '{param_name}' has "
        f"contradictory constraints: lower bound ({lower_bound}) is "
        f"greater than upper bound ({upper_bound})"
      )
      raise ValueError(msg)
    if lower_bound == upper_bound and not (lower_inclusive and upper_inclusive):
      msg = (
        f"In function '{func_name}', parameter '{param_name}' has "
        f"contradictory constraints: no value can satisfy both strict "
        f"and non-strict bounds at {lower_bound}"
      )
      raise ValueError(msg)

  # Check length constraint conflicts
  min_len = constraints.get("min_length")
  max_len = constraints.get("max_length")

  if (
    min_len is not None
    and max_len is not None
    and isinstance(min_len, int)
    and isinstance(max_len, int)
    and min_len > max_len
  ):
    msg = (
      f"In function '{func_name}', parameter '{param_name}' has "
      f"contradictory constraints: min_length ({min_len}) is greater "
      f"than max_length ({max_len})"
    )
    raise ValueError(msg)


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


def extract_hyperparams(func_node: ast.FunctionDef) -> tuple[list[HyperParam], bool]:
  """Extract Hyper[T] parameters from function signature.

  Args:
    func_node: AST FunctionDef node to extract from

  Returns:
    Tuple of (hyperparameters, has_data_param)

  Raises:
    ValueError: If constraints are contradictory
  """
  params: list[HyperParam] = []
  has_data_param = False

  for arg in func_node.args.args:
    if arg.arg in {"data", "self", "cls"}:
      has_data_param = has_data_param or arg.arg == "data"
      continue

    if is_hyper_annotated(arg.annotation):
      constraints = extract_constraints_from_annotation(arg.annotation)
      validate_constraint_conflicts(constraints, arg.arg, func_node.name)

      type_str = extract_inner_type(arg.annotation)
      default_idx = len(func_node.args.args) - len(func_node.args.defaults)
      arg_idx = func_node.args.args.index(arg)

      if arg_idx >= default_idx:
        default = func_node.args.defaults[arg_idx - default_idx]
        default_str = extract_default_value(default)
      else:
        default_str = "..."

      params.append(HyperParam(arg.arg, type_str, default_str))

  return params, has_data_param


def extract_all_params(func_node: ast.FunctionDef) -> list[tuple[str, str, str]]:
  """Extract all parameters with their types and defaults.

  Args:
    func_node: AST FunctionDef node to extract from

  Returns:
    List of (name, type_annotation, default_value) tuples
  """
  all_params: list[tuple[str, str, str]] = []
  default_idx = len(func_node.args.args) - len(func_node.args.defaults)

  for arg_idx, arg in enumerate(func_node.args.args):
    if arg.arg in {"self", "cls"}:
      continue

    if arg.annotation:
      if is_hyper_annotated(arg.annotation):
        type_str = extract_inner_type(arg.annotation)
      else:
        type_str = ast.unparse(arg.annotation)
    else:
      type_str = "Any"

    if arg_idx >= default_idx:
      default = func_node.args.defaults[arg_idx - default_idx]
      default_str = extract_default_value(default)
    else:
      default_str = ""

    all_params.append((arg.arg, type_str, default_str))

  return all_params


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 extract_hyperparams_from_init(
  class_node: ast.ClassDef,
) -> tuple[list[HyperParam], bool] | None:
  """Extract Hyper[T] parameters from class __init__ method.

  Returns:
    Tuple of (hyperparameters, has_data_param) or None if no __init__ found

  Raises:
    ValueError: If constraints are contradictory
  """
  # Find __init__ method in class body
  init_method = None
  for item in class_node.body:
    if isinstance(item, ast.FunctionDef) and item.name == "__init__":
      init_method = item
      break

  if init_method is None:
    return None

  # Extract hyperparams from __init__ signature
  return extract_hyperparams(init_method)


def extract_hyperparams_from_dataclass(
  class_node: ast.ClassDef,
) -> tuple[list[HyperParam], bool]:
  """Extract Hyper[T] fields from dataclass.

  Args:
    class_node: AST ClassDef node representing a dataclass

  Returns:
    Tuple of (hyperparameters, has_data_param)

  Raises:
    ValueError: If constraints are contradictory
  """
  params: list[HyperParam] = []
  has_data_param = False

  for item in class_node.body:
    if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
      field_name = item.target.id

      if field_name == "data":
        has_data_param = True
        continue

      if is_hyper_annotated(item.annotation):
        constraints = extract_constraints_from_annotation(item.annotation)
        validate_constraint_conflicts(constraints, field_name, class_node.name)

        type_str = extract_inner_type(item.annotation)
        default_str = extract_default_value(item.value)

        params.append(HyperParam(field_name, type_str, default_str))

  return params, has_data_param


def extract_all_params_from_class(
  class_node: ast.ClassDef,
) -> list[tuple[str, str, str]]:
  """Extract all parameters from class.

  Args:
    class_node: AST ClassDef node to extract from

  Returns:
    List of (name, type_annotation, default_value) tuples from either
    __init__ method or dataclass fields
  """
  for item in class_node.body:
    if isinstance(item, ast.FunctionDef) and item.name == "__init__":
      return extract_all_params(item)

  all_params: list[tuple[str, str, str]] = []

  for item in class_node.body:
    if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
      field_name = item.target.id

      if field_name in {"self", "cls"}:
        continue

      if item.annotation:
        if is_hyper_annotated(item.annotation):
          type_str = extract_inner_type(item.annotation)
        else:
          type_str = ast.unparse(item.annotation)
      else:
        type_str = "Any"

      default_str = extract_default_value(item.value)

      all_params.append((field_name, type_str, default_str))

  return all_params


def scan_module(source_path: Path) -> list[FunctionInfo]:
  """Scan a Python module for @configurable functions and classes.

  Args:
    source_path: Path to Python source file

  Returns:
    List of FunctionInfo objects for each @configurable item found
  """
  file_size = source_path.stat().st_size
  if file_size > MAX_FILE_SIZE:
    logger.warning(
      "Skipping {}: file too large ({} bytes, max {})",
      source_path,
      file_size,
      MAX_FILE_SIZE,
    )
    return []

  source_code = source_path.read_text(encoding="utf-8")
  tree = ast.parse(source_code)

  functions: list[FunctionInfo] = []

  for node in ast.walk(tree):
    # Handle @configurable functions
    if isinstance(node, ast.FunctionDef) and is_configurable_decorated(node):
      params, has_data = extract_hyperparams(node)
      all_params = extract_all_params(node)
      return_type = ast.unparse(node.returns) if node.returns else "Any"

      functions.append(
        FunctionInfo(
          name=node.name,
          params=params,
          all_params=all_params,
          return_type=return_type,
          has_data_param=has_data,
        )
      )

    # Handle @configurable classes
    elif isinstance(node, ast.ClassDef) and is_configurable_decorated(node):
      # Extract hyperparams based on whether it's a dataclass or regular class
      if is_dataclass_decorated(node):
        params, has_data = extract_hyperparams_from_dataclass(node)
      else:
        result = extract_hyperparams_from_init(node)
        if result is None:
          # No __init__ found, skip this class
          logger.debug(
            "Skipping class {} in {}: no __init__ method found",
            node.name,
            source_path,
          )
          continue
        params, has_data = result

      all_params = extract_all_params_from_class(node)

      # For classes, the return type is the class itself
      return_type = node.name

      functions.append(
        FunctionInfo(
          name=node.name,
          params=params,
          all_params=all_params,
          return_type=return_type,
          has_data_param=has_data,
        )
      )

  return functions


def to_cap_words(snake_case: str) -> str:
  """Convert snake_case to CapWords (PascalCase)."""
  return "".join(word.capitalize() for word in snake_case.split("_"))


def generate_config_class(
  func_name: str, params: list[HyperParam], return_type: str
) -> str:
  """Generate Pydantic config class stub.

  Args:
    func_name: Name of the function/class
    params: List of hyperparameters
    return_type: Return type annotation

  Returns:
    String containing the config class stub definition
  """
  class_name = f"{to_cap_words(func_name)}Config"

  if not params:
    return f"class {class_name}(MakeableModel[{return_type}]):\n  pass"

  fields = "\n".join(f"  {p.name}: {p.type_annotation}" for p in params)

  # Generate __init__ signature
  init_params = ", ".join(
    f"{p.name}: {p.type_annotation} = {p.default_value}" for p in params
  )
  init_sig = f"  def __init__(self, *, {init_params}) -> None: ..."

  return f"class {class_name}(MakeableModel[{return_type}]):\n{fields}\n{init_sig}"


def generate_configurable_class(func_info: FunctionInfo) -> str:
  """Generate ConfigurableIndicator subclass stub.

  Args:
    func_info: Information about the configurable function/class

  Returns:
    String containing the ConfigurableIndicator stub definition
  """
  func_name = func_info.name
  class_name = f"_{to_cap_words(func_name)}Configurable"
  config_class = f"{to_cap_words(func_name)}Config"

  # Build exact __call__ signature from all_params
  param_strs: list[str] = []
  for param_name, param_type, default_val in func_info.all_params:
    if default_val:
      param_strs.append(f"{param_name}: {param_type} = {default_val}")
    else:
      param_strs.append(f"{param_name}: {param_type}")

  if func_info.has_data_param and func_info.params:
    data_param: str = param_strs[0] if param_strs else ""
    hyper_params: list[str] = param_strs[1:] if len(param_strs) > 1 else []
    params_sig: str
    if hyper_params:
      params_sig = f"{data_param}, *, {', '.join(hyper_params)}"
    else:
      params_sig = data_param
  else:
    params_sig = ", ".join(param_strs)

  max_line_length = 88
  one_line_call = f"  def __call__(self, {params_sig}) -> {func_info.return_type}: ..."
  if len(one_line_call) > max_line_length:
    call_sig = f"""def __call__(
    self, {params_sig}
  ) -> {func_info.return_type}: ..."""
  else:
    call_sig = f"def __call__(self, {params_sig}) -> {func_info.return_type}: ..."

  return f"""class {class_name}:
  Config: type[{config_class}]
  {call_sig}

{func_name}: {class_name}"""


def find_pandas_import(source_tree: ast.Module) -> tuple[bool, str]:
  """Find pandas import in AST.

  Args:
    source_tree: Parsed AST module

  Returns:
    Tuple of (has_pandas_import, pandas_alias)
  """
  for node in ast.walk(source_tree):
    if isinstance(node, ast.ImportFrom) and node.module == "pandas":
      alias = node.names[0].asname or "pd"
      return True, alias
    if not isinstance(node, ast.Import):
      continue
    for name in node.names:
      if name.name == "pandas":
        return True, name.asname or "pandas"
  return False, "pd"


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 generate_stub_header(
  *, has_pandas: bool, pandas_alias: str, has_default: bool
) -> str:
  """Generate stub file header with imports.

  Args:
    has_pandas: Whether to include pandas import
    pandas_alias: Alias to use for pandas import
    has_default: Whether to import DEFAULT sentinel

  Returns:
    String containing the stub file header
  """
  header = '''"""Auto-generated type stubs for @configurable decorators.

This file is automatically generated by scripts/generate_stubs.py.
Do not edit manually - changes will be overwritten.
"""

'''

  if has_pandas:
    header += f"import pandas as {pandas_alias}\n\n"

  # Import MakeableModel and DEFAULT if needed
  if has_default:
    header += "from hipr.config import DEFAULT, MakeableModel\n"
  else:
    header += "from hipr.config import MakeableModel\n"

  return header


def generate_stub_content(functions: list[FunctionInfo], source_path: Path) -> str:
  """Generate complete .pyi stub file content.

  Args:
    functions: List of FunctionInfo objects
    source_path: Path to source file (for import detection)

  Returns:
    String containing the complete stub file content
  """
  if not functions:
    return ""

  source_code = source_path.read_text(encoding="utf-8")
  source_tree = ast.parse(source_code)
  has_pandas, pandas_alias = find_pandas_import(source_tree)
  has_default = has_default_usage(source_tree)

  stub = generate_stub_header(
    has_pandas=has_pandas, pandas_alias=pandas_alias, has_default=has_default
  )

  config_classes = [
    generate_config_class(func.name, func.params, func.return_type)
    for func in functions
  ]
  stub += "\n" + "\n\n".join(config_classes)

  configurable_classes = [generate_configurable_class(func) for func in functions]
  stub += "\n\n" + "\n\n".join(configurable_classes) + "\n"

  return stub


def process_file(source_path: Path) -> bool:
  """Process a single Python file and generate its stub if needed.

  Args:
    source_path: Path to Python source file

  Returns:
    True if stub was generated, False otherwise
  """
  functions = scan_module(source_path)

  if not functions:
    return False

  stub_content = generate_stub_content(functions, source_path)
  stub_path = source_path.with_suffix(".pyi")
  stub_path.write_text(stub_content, encoding="utf-8")
  logger.info("Generated: {}", stub_path)

  return True


def main(argv: list[str] | None = None) -> int:
  """Scan directory and generate stubs for all @configurable functions.

  Args:
    argv: Command line arguments (uses sys.argv if None)

  Returns:
    Exit code (0 for success, 1 for error)
  """
  import argparse

  parser = argparse.ArgumentParser(
    description="Generate .pyi stub files for @configurable decorators",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog="""
Examples:
  # Generate stubs for src/ directory (default)
  hipr-generate-stubs

  # Generate stubs for a specific directory
  hipr-generate-stubs my_package/

  # Generate stubs for current directory
  hipr-generate-stubs .
    """,
  )
  parser.add_argument(
    "path",
    nargs="?",
    default="src",
    help="Directory to scan for @configurable decorators (default: src)",
  )
  parser.add_argument(
    "--verbose",
    "-v",
    action="store_true",
    help="Enable verbose logging",
  )

  args = parser.parse_args(argv)

  if not args.verbose:
    logger.remove()  # Remove default handler
    logger.add(sys.stderr, level="INFO", format="<level>{message}</level>")

  src_dir = Path(args.path)

  if not src_dir.exists():
    logger.error("Error: Directory '{}' not found", src_dir)
    return 1

  if not src_dir.is_dir():
    logger.error("Error: '{}' is not a directory", src_dir)
    return 1

  python_files = list(src_dir.rglob("*.py"))
  generated_count = 0

  for py_file in python_files:
    if "__pycache__" in str(py_file):
      logger.debug("Skipping cache directory: {}", py_file)
      continue

    if py_file.name.startswith("test_"):
      logger.debug("Skipping test file: {}", py_file)
      continue

    try:
      if process_file(py_file):
        generated_count += 1
    except SyntaxError as e:
      logger.error("Syntax error in {}: {}", py_file, e)
    except OSError as e:
      logger.error("I/O error reading {}: {}", py_file, e)
    except ValueError as e:
      logger.error("Invalid constraint in {}: {}", py_file, e)
    except Exception as e:
      # Catch any unexpected errors
      logger.error("Unexpected error processing {}: {}", py_file, e)

  logger.info("Generated {} stub file(s)", generated_count)
  return 0


if __name__ == "__main__":
  sys.exit(main())
