from __future__ import annotations

import ast
import textwrap
from typing import TYPE_CHECKING

from loguru import logger

from hipr.stubs.extraction import scan_module
from hipr.stubs.imports import (
  collect_constant_references,
  collect_type_references,
  collect_type_references_from_ast,
  collect_used_names_from_stub,
  extract_imports,
  extract_local_type_definitions,
  extract_module_constants,
  extract_public_classes,
  extract_public_functions,
  filter_imports_by_usage,
  has_default_usage,
)
from hipr.stubs.models import FunctionInfo
from hipr.stubs.utils import to_cap_words

if TYPE_CHECKING:
  from pathlib import Path

  from hipr.stubs.models import HyperParam


def generate_config_class(
  func_name: str,
  params: list[HyperParam],
  return_type: str,
  data_params: list[tuple[str, str, str]] | None = None,
  class_name: str | None = None,
) -> str:
  """Generate Pydantic config class stub.

  Args:
    func_name: Name of the function/class
    params: List of hyperparameters
    return_type: Return type annotation
    data_params: For functions, list of (name, type, default) tuples for data params.
                 If present, make() returns a Callable that takes these params.
    class_name: Optional override for the class name (default: {FuncName}Config)

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

  # Generate make() override
  # For functions with data params: make() -> Callable[[DataTypes], ReturnType]
  # For classes (no data_params): make() -> ReturnType (instance directly)
  if data_params is not None:
    # Build the callable signature
    param_types = ", ".join(t for _, t, _ in data_params)
    make_return = f"Callable[[{param_types}], {return_type}]"
  else:
    make_return = return_type

  make_sig = f"    @override\n    def make(self) -> {make_return}: ..."

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

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

  # Generate __init__ signature
  init_params_list: list[str] = []
  for p in params:
    if p.default_value:
      init_params_list.append(f"{p.name}: {p.type_annotation} = {p.default_value}")
    else:
      init_params_list.append(f"{p.name}: {p.type_annotation}")

  init_params = ", ".join(init_params_list)
  init_sig = f"    def __init__(self, *, {init_params}) -> None: ..."

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


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

  Args:
    func_info: Information about the configurable function/class
    nested_config_body: Body of the nested Config class

  Returns:
    String containing the ConfigurableIndicator stub definition
  """
  func_name = func_info.name

  # 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 __new__(cls, {params_sig}) -> {func_info.return_type}: ..."
  if len(one_line_call) > max_line_length:
    call_sig = (
      "    def __new__(\n"
      f"        cls, {params_sig}\n"
      f"    ) -> {func_info.return_type}: ..."
    )
  else:
    call_sig = one_line_call

  # Add __call__ to make instances callable (matches BoundFunction behavior)
  instance_call_sig = call_sig.replace("def __new__(cls,", "def __call__(self,")
  if "def __new__(\n        cls," in call_sig:
    instance_call_sig = call_sig.replace(
      "def __new__(\n        cls,", "def __call__(\n        self,"
    )

  # Add hyper parameter attributes so they're accessible on the bound function
  # This matches the __getattr__ behavior of BoundFunction at runtime
  hyper_attrs = "\n".join(
    f"    {p.name}: {p.type_annotation}" for p in func_info.params
  )

  # Add to_config method stub
  to_config_sig = "    @classmethod\n    def to_config(cls) -> Config: ..."

  # Add Type alias class to support my_func.Type syntax
  # Use type annotation with forward reference to avoid self-reference/scoping issues
  type_alias = f"    class Type:\n        Config: type['{func_name}.Config']"

  # Add noqa: N801 since function names may be lowercase (e.g., train_model)
  body_elements = [nested_config_body]
  if hyper_attrs:
    body_elements.append(hyper_attrs)
  body_elements.append(type_alias)
  body_elements.append(call_sig)
  body_elements.append(instance_call_sig)
  body_elements.append(to_config_sig)

  body = "\n".join(body_elements)

  return f"class {func_name}:  # noqa: N801\n{body}"


def generate_stub_header(
  *,
  future_imports: list[str],
  other_imports: list[str],
  has_default: bool,
  has_callable: bool = False,
) -> str:
  """Generate stub file header with imports.

  Args:
    future_imports: List of from __future__ import ... statements
    other_imports: List of other import statements
    has_default: Whether to import DEFAULT sentinel
    has_callable: Whether Callable is used in return types

  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.
"""

'''

  # Add future imports first
  if future_imports:
    header += "\n".join(future_imports) + "\n\n"

  # Add other extracted imports
  if other_imports:
    header += "\n".join(other_imports) + "\n\n"

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

  # Add typing imports
  header += "from typing import override\n"

  # Add Callable import if used (for function configs with data params)
  if has_callable:
    header += "from collections.abc import Callable\n"

  return header


def generate_stub_content(functions: list[FunctionInfo], source_path: Path) -> str:  # noqa: C901, PLR0912, PLR0914, PLR0915
  """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)

  future_imports, other_imports = extract_imports(source_tree)
  has_default = has_default_usage(source_tree)

  # Generate stub body first (without header) so we can analyze what's used
  configurable_classes: list[str] = []

  # Primitive/builtin types that should NOT get .Config appended
  primitive_types = {
    "int",
    "float",
    "str",
    "bool",
    "bytes",
    "None",
    "list",
    "dict",
    "set",
    "tuple",
    "Any",
    "Optional",
    "Union",
    "Callable",
    "Sequence",
    "Mapping",
    "Iterable",
  }

  for func in functions:
    # Map Hyper[T] to TConfig if T is a local configurable class
    # Or to T.Config if T is an external configurable (non-primitive, non-local)
    configurable_names = {f.name for f in functions}

    # Transform params for Config class (use .Config types)
    config_params: list[HyperParam] = []
    for p in func.params:
      type_ann = p.type_annotation.strip("'\"")

      # If type already ends with .Config or Config, it's already handled
      if ".Config" in type_ann or type_ann.endswith("Config"):
        pass
      # Handle .Type suffix (e.g., my_func.Type -> my_func.Config)
      elif type_ann.endswith(".Type") and type_ann[:-5] in configurable_names:
        type_ann = f"{type_ann[:-5]}.Config"
      # If type matches a local configurable, append .Config suffix
      elif type_ann in configurable_names or (
        type_ann.isidentifier() and type_ann not in primitive_types
      ):
        type_ann = f"{type_ann}.Config"

      config_params.append(p._replace(type_annotation=type_ann))

    # Transform params for Instance attributes (use Component types)
    instance_params: list[HyperParam] = []
    for p in func.params:
      type_ann = p.type_annotation.strip("'\"")

      # Handle .Type suffix (e.g., my_func.Type -> my_func)
      if type_ann.endswith(".Type") and type_ann[:-5] in configurable_names:
        type_ann = type_ann[:-5]
      # For standard configurables, keep them as is (don't append .Config)
      # This ensures instance fields are typed as the Component, not Config

      instance_params.append(p._replace(type_annotation=type_ann))

    # Also transform all_params types for __call__ signature
    mapped_all_params: list[tuple[str, str, str]] = []
    for name, ptype, default in func.all_params:
      # Apply same transformation logic
      resolved_type = ptype
      if ".Config" in ptype or ptype.endswith("Config"):
        pass
      elif ptype.endswith(".Type") and ptype[:-5] in configurable_names:
        resolved_type = f"{ptype[:-5]}.Config"
      elif ptype in configurable_names or (
        ptype.isidentifier() and ptype not in primitive_types
      ):
        resolved_type = f"{ptype}.Config"
      mapped_all_params.append((name, resolved_type, default))

    # Create mapped func with transformed params
    # Use instance_params for the FunctionInfo passed to generate_configurable_class
    mapped_func = FunctionInfo(
      name=func.name,
      params=instance_params,
      all_params=mapped_all_params,
      return_type=func.return_type,
      has_data_param=func.has_data_param,
      is_function=func.is_function,
    )

    # Extract data params (non-Hyper params) for functions
    # These are the params that make() returns a Callable for
    hyper_names = {p.name for p in func.params}
    data_params = (
      [
        (name, ptype, default)
        for name, ptype, default in func.all_params
        if name not in hyper_names
      ]
      if func.is_function
      else None
    )

    # Generate Config class content (nested logic)
    # Use config_params here
    config_cls_str = generate_config_class(
      func.name, config_params, func.return_type, data_params, class_name="Config"
    )

    # Indent the config class for nesting
    nested_config_body = textwrap.indent(config_cls_str, "    ")

    # Use mapped_func for configurable class to get correct types
    configurable_classes.append(
      generate_configurable_class(mapped_func, nested_config_body)
    )

  stub_body = "\n" + "\n\n".join(configurable_classes) + "\n"

  # Collect names used in the generated stub
  used_names = collect_used_names_from_stub(stub_body)

  # Always include MakeableModel since it's used in all config classes
  used_names.add("MakeableModel")
  if has_default:
    used_names.add("DEFAULT")

  # Filter imports to only those actually used
  filtered_imports = filter_imports_by_usage(other_imports, used_names)

  # Collect type references from function signatures
  type_refs = collect_type_references(functions)

  # Get configurable names to exclude from local type definitions
  configurable_names = {func.name for func in functions}
  local_types, base_class_refs = extract_local_type_definitions(
    source_tree, type_refs, configurable_names
  )

  # Add base class references to used_names for import filtering
  used_names.update(base_class_refs)

  # Collect constants referenced in default values
  constant_refs = collect_constant_references(functions)
  module_constants = extract_module_constants(source_tree, constant_refs)

  # Extract public (non-configurable) functions
  public_functions, public_func_nodes = extract_public_functions(
    source_tree, configurable_names
  )

  # Collect type references from public functions for import filtering
  public_type_refs = collect_type_references_from_ast(public_func_nodes)
  used_names.update(public_type_refs)

  # Extract all public classes (not just those in type annotations)
  public_classes, public_class_bases = extract_public_classes(
    source_tree, configurable_names
  )
  used_names.update(public_class_bases)

  # Re-filter imports now that we have public function and class types
  filtered_imports = filter_imports_by_usage(other_imports, used_names)

  # Check if any function has data params (needs Callable import)
  has_callable = any(f.has_data_param or f.is_function for f in functions)

  # Generate header with filtered imports
  stub = generate_stub_header(
    future_imports=future_imports,
    other_imports=filtered_imports,
    has_default=has_default,
    has_callable=has_callable,
  )

  # Add local type definitions after imports
  if local_types:
    stub += "\n# Local type definitions\n"
    stub += "\n\n".join(local_types) + "\n"

  # Add module constants after type definitions
  if module_constants:
    stub += "\n# Module constants\n"
    stub += "\n".join(module_constants) + "\n"

  # Add public functions after constants
  if public_functions:
    stub += "\n# Public functions\n"
    stub += "\n\n".join(public_functions) + "\n"

  # Add public classes (non-configurable)
  if public_classes:
    stub += "\n# Public classes\n"
    stub += "\n\n".join(public_classes) + "\n"

  # Append the pre-generated body
  stub += stub_body

  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
