"""Tests for enum handling in stub generation.

These tests verify that:
- Enums used as Hyper type annotations are included in stubs
"""

from __future__ import annotations

from enum import Enum, StrEnum

from hipr import Hyper, configurable


class Mode(StrEnum):
  """Example enum for testing."""

  FAST = "fast"
  SLOW = "slow"
  MEDIUM = "medium"


class Priority(int, Enum):
  """Integer enum for testing."""

  LOW = 1
  MEDIUM = 2
  HIGH = 3


class TestEnumRuntime:
  """Runtime tests for enum handling."""

  def test_enum_hyper_param(self) -> None:
    """Enum can be used as Hyper type annotation."""

    @configurable
    def process(
      mode: Hyper[Mode] = Mode.FAST,
    ) -> str:
      return f"Processing in {mode.value} mode"

    config = process.Config(mode=Mode.SLOW)
    fn = config.make()
    result = fn()
    assert result == "Processing in slow mode"

  def test_enum_string_coercion(self) -> None:
    """Pydantic should coerce string to enum value."""

    @configurable
    def process(
      mode: Hyper[Mode] = Mode.FAST,
    ) -> str:
      return mode.value

    # Pass string value instead of enum
    config = process.Config(mode="slow")  # type: ignore[arg-type]
    fn = config.make()
    result = fn()
    assert result == "slow"

  def test_int_enum(self) -> None:
    """Integer enums should work with Hyper."""

    @configurable
    def prioritize(
      priority: Hyper[Priority] = Priority.MEDIUM,
    ) -> int:
      return priority.value

    config = prioritize.Config(priority=Priority.HIGH)
    fn = config.make()
    assert fn() == 3

  def test_enum_in_class(self) -> None:
    """Enum Hyper param in class __init__."""

    @configurable
    class Processor:
      def __init__(self, mode: Hyper[Mode] = Mode.FAST) -> None:
        self.mode = mode

    config = Processor.Config(mode=Mode.SLOW)
    instance = config.make()
    assert instance.mode == Mode.SLOW


class TestEnumStubGeneration:
  """Tests for enum handling in stub generation."""

  def test_enum_type_reference_collected(self) -> None:
    """Enum types should be collected as type references."""
    from hipr.stubs.imports import collect_type_references
    from hipr.stubs.models import FunctionInfo, HyperParam

    func = FunctionInfo(
      name="process",
      params=[HyperParam("mode", "Mode", "Mode.FAST")],
      all_params=[("mode", "Mode", "Mode.FAST")],
      return_type="str",
      has_data_param=False,
      is_function=True,
    )

    refs = collect_type_references([func])
    assert "Mode" in refs

  def test_local_enum_extracted(self) -> None:
    """Local enum definitions should be extracted for stubs."""
    import ast

    from hipr.stubs.imports import extract_local_type_definitions

    source = """
class Mode(str, Enum):
    FAST = "fast"
    SLOW = "slow"

@configurable
def process(mode: Hyper[Mode] = Mode.FAST) -> str:
    return mode.value
"""
    tree = ast.parse(source)
    type_names = {"Mode"}
    configurable_names = {"process"}

    definitions, base_refs = extract_local_type_definitions(
      tree, type_names, configurable_names
    )

    # Mode should be extracted
    assert any("class Mode" in d for d in definitions)
    # Enum should be in base class refs
    assert "Enum" in base_refs

  def test_enum_import_preserved(self) -> None:
    """Enum import should be preserved in stub."""
    import ast

    from hipr.stubs.imports import extract_imports

    source = """
from enum import Enum

class Mode(str, Enum):
    FAST = "fast"

@configurable
def process(mode: Hyper[Mode]) -> str:
    pass
"""
    tree = ast.parse(source)
    _future_imports, other_imports = extract_imports(tree)

    # Enum import should be extracted
    assert any("Enum" in imp for imp in other_imports)


class TestPublicClassStubInclusion:
  """Tests verifying all public classes are included in stubs."""

  def test_public_class_extracted(self) -> None:
    """Non-configurable public classes should be extracted for stubs."""
    import ast

    from hipr.stubs.imports import extract_public_classes

    source = """
from enum import Enum

class Mode(str, Enum):
    FAST = "fast"
    SLOW = "slow"

class Helper:
    def process(self, data: list) -> list:
        return data

@configurable
class Processor:
    def __init__(self, mode: Hyper[Mode] = Mode.FAST) -> None:
        self.mode = mode
"""
    tree = ast.parse(source)
    configurable_names = {"Processor"}

    class_defs, base_refs = extract_public_classes(tree, configurable_names)

    # Extract class names from definitions
    def get_class_name(def_str: str) -> str:
      # Parse "class Name(Base):" or "class Name:"
      import re

      match = re.match(r"class (\w+)", def_str)
      return match.group(1) if match else ""

    class_names = [get_class_name(d) for d in class_defs]
    assert "Mode" in class_names
    assert "Helper" in class_names
    assert "Processor" not in class_names

    # Enum should be in base refs
    assert "Enum" in base_refs

  def test_private_classes_excluded(self) -> None:
    """Private classes (starting with _) should not be extracted."""
    import ast

    from hipr.stubs.imports import extract_public_classes

    source = """
class PublicClass:
    pass

class _PrivateClass:
    pass

class __DunderPrivate:
    pass
"""
    tree = ast.parse(source)
    class_defs, _ = extract_public_classes(tree, set())

    class_names = [d.split(":")[0].replace("class ", "") for d in class_defs]
    assert "PublicClass" in class_names
    assert "_PrivateClass" not in class_names
    assert "__DunderPrivate" not in class_names

  def test_class_methods_stubbed(self) -> None:
    """Class methods should have ... body in stubs."""
    import ast

    from hipr.stubs.imports import extract_public_classes

    source = """
class MyClass:
    value: int = 0

    def process(self, data: list) -> list:
        return [x * 2 for x in data]

    def reset(self) -> None:
        self.value = 0
"""
    tree = ast.parse(source)
    class_defs, _ = extract_public_classes(tree, set())

    assert len(class_defs) == 1
    stub = class_defs[0]

    # Methods should be present with ... body
    assert "def process" in stub
    assert "def reset" in stub
    # Implementation should not be present
    assert "x * 2" not in stub
