from __future__ import annotations

import dataclasses
import inspect
import types
from typing import (
  TYPE_CHECKING,
  Annotated,
  Any,
  cast,
  get_args,
  get_origin,
  get_type_hints,
)

from annotated_types import (
  BaseMetadata,
  Ge as AtGe,
  Gt as AtGt,
  Le as AtLe,
  Lt as AtLt,
  MaxLen as AtMaxLen,
  MinLen as AtMinLen,
  MultipleOf as AtMultipleOf,
)
from loguru import logger
from pydantic import Field
from pydantic.fields import FieldInfo

from hipr.models import MakeableModel
from hipr.registry import _config_creation_stack  # pyright: ignore[reportPrivateUsage]
from hipr.typedefs import (
  DefaultSentinel,
  Hyper,
  HyperMarker,
)

if TYPE_CHECKING:
  from collections.abc import Callable

  from pydantic import BaseModel

__all__ = [
  "_extract_annotated_constraint",
  "_extract_constraints_from_metadata",
  "_extract_dataclass_params",
  "_extract_hyper_field_info",
  "_extract_init_params",
  "_extract_public_fields",
  "_get_type_hints_with_fallback",
]

from hipr.constraints import validate_constraint_conflicts

# Pattern constraints use tuple format (key, value)
_PATTERN_CONSTRAINT_TUPLE_LENGTH = 2


def _extract_public_fields(model: BaseModel) -> dict[str, object]:
  """Extract non-private fields from a Pydantic model.

  Args:
    model: The Pydantic model instance

  Returns:
    Dictionary of field names to values, excluding private fields (starting with _)
  """
  # Filter out private fields from model's __dict__
  model_dict = cast("dict[str, object]", model.__dict__)
  data: dict[str, object] = {
    k: v for k, v in model_dict.items() if not k.startswith("_")
  }

  return data


def _get_type_hints_with_fallback(fn: Callable[..., object]) -> dict[str, object]:
  """Get type hints with fallback for TYPE_CHECKING imports.

  Args:
    fn: Function or class to extract type hints from

  Returns:
    Dictionary mapping parameter names to their type annotations

  Note:
    When get_type_hints() fails due to NameError (from TYPE_CHECKING imports),
    retries with an augmented namespace containing the Hyper type.
  """
  # Try normal resolution first
  try:
    # Classes don't have __globals__, so use simpler call for them
    if inspect.isclass(fn):
      return get_type_hints(fn, include_extras=True)
    return get_type_hints(fn, globalns=fn.__globals__, include_extras=True)
  except NameError:
    # Fallback: augment namespace with our types
    if inspect.isclass(fn):
      # For classes, get the module's globals
      module = inspect.getmodule(fn)
      augmented_globals = {} if module is None else vars(module).copy()
    else:
      augmented_globals = fn.__globals__.copy()

    # Add Hyper to globals for TYPE_CHECKING imports
    augmented_globals["Hyper"] = Hyper
    return get_type_hints(fn, globalns=augmented_globals, include_extras=True)


def _extract_annotated_constraint(item: BaseMetadata) -> dict[str, object]:
  """Extract constraint from an annotated_types metadata object.

  Args:
    item: annotated_types constraint object (Ge, Le, Gt, Lt, etc.)

  Returns:
    Dictionary with constraint field names and values for Pydantic Field
  """
  constraints: dict[str, object] = {}
  if isinstance(item, AtGe):
    constraints["ge"] = item.ge
  elif isinstance(item, AtLe):
    constraints["le"] = item.le
  elif isinstance(item, AtGt):
    constraints["gt"] = item.gt
  elif isinstance(item, AtLt):
    constraints["lt"] = item.lt
  elif isinstance(item, AtMinLen):
    constraints["min_length"] = item.min_length
  elif isinstance(item, AtMaxLen):
    constraints["max_length"] = item.max_length
  elif isinstance(item, AtMultipleOf):
    constraints["multiple_of"] = item.multiple_of
  return constraints


def _extract_constraints_from_metadata(
  metadata: tuple[object, ...],
) -> dict[str, object]:
  """Extract constraints from Annotated metadata.

  Args:
    metadata: Tuple of metadata items from Annotated type

  Returns:
    Dictionary of constraint field names to values for Pydantic Field

  Note:
    Handles both annotated_types objects (Ge, Le, etc.) and tuple format
    used by Pattern constraints.
  """
  constraints: dict[str, object] = {}
  for item in metadata:
    if isinstance(item, BaseMetadata):
      constraints.update(_extract_annotated_constraint(item))
    elif isinstance(item, tuple) and len(item) == _PATTERN_CONSTRAINT_TUPLE_LENGTH:  # pyright: ignore[reportUnknownArgumentType]
      # Handle tuple constraints (Pattern uses this format)
      # Cast is safe after length check - Pattern constraints always have (str, value) format
      key, value = cast("tuple[str, object]", item)
      constraints[key] = value

  return constraints


def _extract_hyper_field_info(  # noqa: C901, PLR0912, PLR0914, PLR0915
  annotation: object,
  default_val: object,
  param_name: str,
  *,
  is_class: bool = False,
) -> tuple[type, FieldInfo] | None:
  """Extract field information from a parameter annotation.

  Args:
    annotation: Type annotation to analyze
    default_val: Default value for the parameter
    param_name: Name of the parameter (for error messages)
    is_class: If True, treat all params as hyperparams (for classes/dataclasses)

  Returns:
    Tuple of (inner_type, FieldInfo) if param should be a hyperparam, None otherwise

  Raises:
    TypeError: If Hyper parameter lacks default value or DEFAULT is misused
  """
  inner_type: type
  metadata: tuple[object, ...] = ()
  has_hyper_marker = False

  # Check for Annotated[T, ...] pattern (includes Hyper[T])
  if get_origin(annotation) is Annotated:
    args = get_args(annotation)
    if not args:
      return None  # Empty Annotated, skip
    inner_type = cast("type", args[0])
    metadata = cast("tuple[object, ...]", args[1:])
    has_hyper_marker = any(
      isinstance(arg, type) and issubclass(arg, HyperMarker) for arg in metadata
    )
  else:
    # Plain type (not Annotated)
    inner_type = cast("type", annotation)

  # For functions: only Hyper[T] marked params become hyperparams
  # For classes: all params become hyperparams
  if not is_class and not has_hyper_marker:
    return None

  # Extract constraints from metadata (if any) - do this early for all code paths
  constraints = _extract_constraints_from_metadata(metadata)

  # Validate constraints for conflicts
  validate_constraint_conflicts(constraints, param_name)

  if constraints:
    logger.trace(
      "Parameter '{}' has constraints: {}",
      param_name,
      constraints,
    )

  # Determine field type - use Config type for nested configurables
  field_type = inner_type
  if hasattr(inner_type, "Config") and issubclass(inner_type.Config, MakeableModel):  # pyright: ignore[reportUnknownMemberType]
    field_type = inner_type.Config  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]

  # Handle required fields (no default value)
  if default_val is inspect.Parameter.empty:
    return (field_type, Field(..., **cast("Any", constraints)))  # pyright: ignore[reportExplicitAny, reportAny, reportUnknownVariableType]

  # Check if there's an existing FieldInfo in the metadata
  existing_field_info = next(
    (arg for arg in metadata if isinstance(arg, FieldInfo)),
    None,
  )

  # Handle DEFAULT sentinel for nested configs
  if isinstance(default_val, DefaultSentinel):
    # Check if inner_type is a MakeableModel subclass (nested config)
    # or a union type that contains a MakeableModel subclass
    config_type: type[MakeableModel[object]] | None = None

    if inspect.isclass(inner_type) and issubclass(inner_type, MakeableModel):
      config_type = inner_type  # pyright: ignore[reportUnknownVariableType]
    # Check if inner_type is a configurable class (has .Config)
    elif hasattr(inner_type, "Config") and issubclass(inner_type.Config, MakeableModel):  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
      config_type = inner_type.Config  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
    elif isinstance(inner_type, types.UnionType):  # pyright: ignore[reportUnnecessaryIsInstance]
      # Handle union types
      # Find the MakeableModel subclass in the union
      for arg in get_args(inner_type):  # pyright: ignore[reportAny]
        if inspect.isclass(arg) and issubclass(arg, MakeableModel):  # pyright: ignore[reportAny]
          config_type = arg  # pyright: ignore[reportUnknownVariableType]
          break

    if config_type is not None:
      # Check for circular dependency
      stack = _config_creation_stack.get() or []
      config_name = f"{config_type.__module__}.{config_type.__qualname__}"

      if config_name in stack:
        cycle_path = " -> ".join([*stack, config_name])
        msg = (
          f"Circular dependency detected: {cycle_path}. "
          f"Parameter '{param_name}' creates a cycle."
        )
        raise ValueError(msg)

      # Track this config creation
      new_stack = [*stack, config_name]
      token = _config_creation_stack.set(new_stack)

      try:
        instance = config_type()
        if not hasattr(instance, "make") or not callable(instance.make):
          msg = (
            f"Nested config type '{config_type.__name__}' for parameter "
            f"'{param_name}' does not have a callable 'make' method"
          )
          raise TypeError(msg)
        default_val = instance
      except ValueError:
        raise
      except Exception as e:
        msg = (
          f"Failed to instantiate nested config type '{config_type.__name__}' "
          f"for parameter '{param_name}': {e}"
        )
        raise TypeError(msg) from e
      finally:
        _config_creation_stack.reset(token)
    else:
      type_name = getattr(inner_type, "__name__", str(inner_type))  # pyright: ignore[reportUnknownArgumentType]
      msg = (
        f"DEFAULT can only be used with nested Config types, "
        f"but '{param_name}' has type '{type_name}'"
      )
      raise TypeError(msg)

  if existing_field_info:
    return (inner_type, existing_field_info)  # pyright: ignore[reportUnknownVariableType]

  # If inner_type is a configurable class, use its Config type
  # This enforces strict Config types in the Pydantic model
  field_type = inner_type  # pyright: ignore[reportUnknownVariableType]
  if hasattr(inner_type, "Config") and issubclass(inner_type.Config, MakeableModel):  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
    field_type = inner_type.Config  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportAttributeAccessIssue]

  # Create Field with constraints
  return (  # pyright: ignore[reportUnknownVariableType]
    field_type,
    Field(default=default_val, **cast("Any", constraints)),  # pyright: ignore[reportExplicitAny, reportAny]
  )


def _extract_dataclass_params(cls: type) -> dict[str, tuple[type, FieldInfo]]:
  """Extract hyperparameters from a dataclass.

  Args:
    cls: Dataclass to extract parameters from

  Returns:
    Dictionary mapping parameter names to (type, FieldInfo) tuples
  """
  hyper_params: dict[str, tuple[type, FieldInfo]] = {}
  type_hints = _get_type_hints_with_fallback(cls)

  for field in dataclasses.fields(cls):
    # Skip fields that are not in __init__
    if not field.init:
      continue

    if field.name not in type_hints:
      continue

    field_type = type_hints[field.name]

    default_val: object
    if field.default is not dataclasses.MISSING:
      default_val = cast("object", field.default)
    elif field.default_factory is not dataclasses.MISSING:
      default_val = cast("object", field.default_factory())
    else:
      default_val = inspect.Parameter.empty

    field_info = _extract_hyper_field_info(
      field_type, default_val, field.name, is_class=True
    )
    if field_info:
      hyper_params[field.name] = field_info
  return hyper_params


def _extract_init_params(cls: type) -> dict[str, tuple[type, FieldInfo]]:
  """Extract hyperparameters from class __init__ signature.

  Args:
    cls: Class to extract __init__ parameters from

  Returns:
    Dictionary mapping parameter names to (type, FieldInfo) tuples
  """
  hyper_params: dict[str, tuple[type, FieldInfo]] = {}
  init_fn = cls.__init__
  sig = inspect.signature(init_fn)
  type_hints = _get_type_hints_with_fallback(init_fn)

  for param in sig.parameters.values():
    if param.name == "self":
      continue

    if param.name not in type_hints:
      continue

    annotation = type_hints[param.name]
    default_val: object = cast("object", param.default)

    field_info = _extract_hyper_field_info(
      annotation, default_val, param.name, is_class=True
    )
    if field_info:
      hyper_params[param.name] = field_info
  return hyper_params
