from __future__ import annotations

import ast
from typing import TYPE_CHECKING

from loguru import logger

if TYPE_CHECKING:
  from pathlib import Path

from hipr.constraints import validate_constraint_conflicts
from hipr.stubs.constraints import extract_constraints_from_annotation
from hipr.stubs.models import FunctionInfo, HyperParam
from hipr.stubs.utils import (
  extract_default_value,
  extract_inner_type,
  is_configurable_decorated,
  is_dataclass_decorated,
  is_hyper_annotated,
)

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


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:
        # Required parameter (no default)
        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 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 all fields from dataclass as hyperparameters.

  For @configurable dataclasses, ALL fields (except field(init=False) and 'data')
  are treated as hyperparameters, regardless of whether they're wrapped in Hyper[].

  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

      # Skip fields with init=False
      if _is_init_false_field(item.value):
        continue

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

      # For Hyper[] annotated fields, validate constraints
      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)
      else:
        # Non-Hyper fields are also hyperparams in dataclasses
        type_str = ast.unparse(item.annotation)

      default_str = extract_default_value(item.value)
      params.append(HyperParam(field_name, type_str, default_str))

  return params, has_data_param


def _is_init_false_field(value: ast.expr | None) -> bool:
  """Check if a field value is field(..., init=False).

  Args:
    value: AST expression for the field default value

  Returns:
    True if the field has init=False, False otherwise
  """
  if value is None:
    return False

  # Check for field(..., init=False)
  if not (isinstance(value, ast.Call) and isinstance(value.func, ast.Name)):
    return False

  if value.func.id != "field":
    return False

  for keyword in value.keywords:
    if (
      keyword.arg == "init"
      and isinstance(keyword.value, ast.Constant)
      and keyword.value.value is False
    ):
      return True
  return False


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

      # Skip fields with init=False
      if _is_init_false_field(item.value):
        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,
          is_function=True,
        )
      )

    # 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,
          is_function=False,
        )
      )

  return functions
