from __future__ import annotations

import ast
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_functions,
  filter_imports_by_usage,
  has_default_usage,
)
from hipr.stubs.utils import to_cap_words

if TYPE_CHECKING:
  from pathlib import Path

  from hipr.stubs.models import FunctionInfo, HyperParam


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

  # Use 'type' alias syntax instead of variable annotation to avoid
  # "Variable not allowed in type expression" errors in stub files
  return f"""class {class_name}:
    Config: type[{config_class}]
    {call_sig}

type {func_name} = {class_name}"""


def generate_stub_header(
  *,
  future_imports: list[str],
  other_imports: list[str],
  has_default: bool,
) -> 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

  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
  # We check if they are already imported to avoid duplication,
  # but since we don't parse the imports deeply, we'll just add them.
  # Python handles duplicate imports fine.
  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:  # noqa: PLR0914
  """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
  config_classes = [
    generate_config_class(func.name, func.params, func.return_type)
    for func in functions
  ]
  stub_body = "\n" + "\n\n".join(config_classes)

  configurable_classes = [generate_configurable_class(func) for func in functions]
  stub_body += "\n\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 = extract_local_type_definitions(
    source_tree, type_refs, configurable_names
  )

  # 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)

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

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

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

  # 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
