@callback forward(program(), inputs()) :: {:ok, outputs()} | {:error, term()}
@callback forward(program(), inputs(), options()) :: {:ok, outputs()} | {:error, term()}
@optional_callbacks forward: 3

@callback compile(
  student :: program(),
  teacher :: program(),
  trainset :: trainset(),
  metric_fn :: metric_fn(),
  opts :: opts()
) :: compilation_result()

@callback instructions() :: String.t()
@callback input_fields() :: [atom()]
@callback output_fields() :: [atom()]
@callback fields() :: [atom()]

@callback apply(Bucket.t(), struct(), map()) :: {:ok, Bucket.t()} | {:error, term()}
@callback applicable?(Bucket.t(), map()) :: boolean()

defmodule DSPEx.Client do
  @callback request(messages :: [message()], opts :: request_options()) :: 
    {:ok, response()} | {:error, error_reason()}
  
  @callback request(client_id :: atom(), messages :: [message()], opts :: request_options()) :: 
    {:ok, response()} | {:error, error_reason()}
end

defmodule DSPEx.Adapter do
  @callback format_messages(signature :: signature(), inputs :: inputs()) :: 
    {:ok, messages()} | {:error, atom()}
  
  @callback parse_response(signature :: signature(), response :: api_response()) :: 
    {:ok, outputs()} | {:error, atom()}
end

defmodule DSPEx.Evaluator do
  @callback run(program :: DSPEx.Program.t(), examples :: [DSPEx.Example.t()], 
                metric_fn :: function(), opts :: evaluation_options()) :: 
    {:ok, evaluation_result()} | {:error, term()}
  
  @callback run_local(program :: DSPEx.Program.t(), examples :: [DSPEx.Example.t()], 
                      metric_fn :: function(), opts :: evaluation_options()) :: 
    {:ok, evaluation_result()} | {:error, term()}
  
  @callback run_distributed(program :: DSPEx.Program.t(), examples :: [DSPEx.Example.t()], 
                            metric_fn :: function(), opts :: evaluation_options()) :: 
    {:ok, evaluation_result()} | {:error, term()}
end

# lib/dspex/client.ex - Add behaviour definition
defmodule DSPEx.Client do
  @type message :: %{role: String.t(), content: String.t()}
  @type request_options :: %{
    optional(:provider) => atom(),
    optional(:model) => String.t(),
    optional(:temperature) => float(),
    optional(:max_tokens) => pos_integer(),
    optional(:correlation_id) => String.t(),
    optional(:timeout) => pos_integer()
  }
  @type response :: %{choices: [%{message: message()}]}
  @type error_reason :: :timeout | :network_error | :invalid_response | :api_error
  
  @callback request(messages :: [message()]) :: {:ok, response()} | {:error, error_reason()}
  @callback request(messages :: [message()], opts :: request_options()) :: 
    {:ok, response()} | {:error, error_reason()}
  @callback request(client_id :: atom(), messages :: [message()], opts :: request_options()) :: 
    {:ok, response()} | {:error, error_reason()}

  # Keep existing implementation but add @behaviour DSPEx.Client
  # to concrete client implementations
end

# lib/dspex/adapter.ex - Add behaviour definition
defmodule DSPEx.Adapter do
  @type signature :: module()
  @type inputs :: map()
  @type outputs :: map()
  @type messages :: [DSPEx.Client.message()]
  @type api_response :: DSPEx.Client.response()

  @callback format_messages(signature :: signature(), inputs :: inputs()) :: 
    {:ok, messages()} | {:error, atom()}
  @callback parse_response(signature :: signature(), response :: api_response()) :: 
    {:ok, outputs()} | {:error, atom()}

  # Keep existing implementation as default behaviour
end

# lib/dspex/evaluate.ex - Add behaviour definition
defmodule DSPEx.Evaluator do
  @type evaluation_result :: %{
    score: float(),
    stats: %{
      total_examples: non_neg_integer(),
      successful: non_neg_integer(),
      failed: non_neg_integer(),
      duration_ms: non_neg_integer(),
      success_rate: float(),
      throughput: float(),
      errors: [term()]
    }
  }
  @type evaluation_options :: keyword()

  @callback run(program :: DSPEx.Program.t(), examples :: [DSPEx.Example.t()], 
                metric_fn :: function(), opts :: evaluation_options()) :: 
    {:ok, evaluation_result()} | {:error, term()}

  # Add @behaviour DSPEx.Evaluator to the module itself
end

# test/test_helper.exs - Add centralized mock definitions
ExUnit.start()

# Define all behaviour-based mocks
Mox.defmock(DSPEx.MockClient, for: DSPEx.Client)
Mox.defmock(DSPEx.MockAdapter, for: DSPEx.Adapter) 
Mox.defmock(DSPEx.MockEvaluator, for: DSPEx.Evaluator)
Mox.defmock(DSPEx.MockProgram, for: DSPEx.Program)
Mox.defmock(DSPEx.MockTeleprompter, for: DSPEx.Teleprompter)

# Load existing test support files
Code.require_file("support/mock_helpers.exs", __DIR__)
Code.require_file("support/test_helpers.exs", __DIR__)
Code.require_file("support/beacon_test_mocks.exs", __DIR__)
Code.require_file("support/mock_provider.ex", __DIR__)

# New centralized mock support
Code.require_file("support/behaviour_mocks.exs", __DIR__)

# Configure existing settings...

# test/support/behaviour_mocks.exs
defmodule DSPEx.BehaviourMocks do
  @moduledoc """
  Centralized mock behaviors and helpers for contract-first testing.
  
  This module provides pre-configured mock implementations that follow
  the established behaviours, making tests more reliable and reducing
  boilerplate.
  """
  
  import Mox
  
  def setup_standard_client_mock(provider \\ :openai) do
    expect(DSPEx.MockClient, :request, fn messages, _opts ->
      # Standard success response
      content = generate_mock_response(messages)
      {:ok, %{choices: [%{message: %{role: "assistant", content: content}}]}}
    end)
  end
  
  def setup_standard_adapter_mock() do
    expect(DSPEx.MockAdapter, :format_messages, fn _signature, inputs ->
      # Convert inputs to simple message format
      content = Map.values(inputs) |> Enum.join(", ")
      {:ok, [%{role: "user", content: content}]}
    end)
    
    expect(DSPEx.MockAdapter, :parse_response, fn signature, response ->
      # Extract first output field and populate with response content
      output_fields = signature.output_fields()
      first_field = List.first(output_fields)
      content = get_in(response, [:choices, Access.at(0), :message, :content])
      {:ok, %{first_field => content}}
    end)
  end
  
  def setup_standard_evaluator_mock(score \\ 0.8) do
    expect(DSPEx.MockEvaluator, :run, fn _program, examples, _metric_fn, _opts ->
      count = length(examples)
      {:ok, %{
        score: score,
        stats: %{
          total_examples: count,
          successful: count,
          failed: 0,
          duration_ms: 1000,
          success_rate: score,
          throughput: count / 1.0,
          errors: []
        }
      }}
    end)
  end
  
  def setup_standard_program_mock(outputs \\ %{answer: "mocked response"}) do
    expect(DSPEx.MockProgram, :forward, fn _program, _inputs, _opts ->
      {:ok, outputs}
    end)
  end
  
  def setup_failing_client_mock(error \\ :network_error) do
    expect(DSPEx.MockClient, :request, fn _messages, _opts ->
      {:error, error}
    end)
  end
  
  # Helper to generate contextual mock responses
  defp generate_mock_response(messages) do
    last_message = List.last(messages)
    content = Map.get(last_message, :content, "")
    
    cond do
      String.contains?(content, "2+2") -> "4"
      String.contains?(content, "capital") -> "The capital is..."
      String.contains?(content, "summarize") -> "Summary: ..."
      true -> "Mock response for: #{String.slice(content, 0, 50)}"
    end
  end
end

# Scattered across tests
defmodule MockClient do
  def request(messages) do
    {:ok, %{choices: [%{message: %{content: "4"}}]}}
  end
end

defmodule MockProgram do
  def forward(_program, _inputs, _opts) do
    {:ok, %{answer: "test response"}}
  end
end

test "program evaluation with mocked components" do
  # Setup all mocks consistently
  DSPEx.BehaviourMocks.setup_standard_client_mock()
  DSPEx.BehaviourMocks.setup_standard_adapter_mock()
  DSPEx.BehaviourMocks.setup_standard_program_mock()
  
  # Test with dependency injection
  program = %DSPEx.Predict{
    signature: TestSignature,
    client: DSPEx.MockClient,  # Injected mock
    adapter: DSPEx.MockAdapter # Injected mock
  }
  
  # Test execution
  {:ok, result} = DSPEx.Program.forward(program, %{question: "2+2?"})
  assert result.answer == "mocked response"
  
  # Verify all mocks were called as expected
  verify!(DSPEx.MockClient)
  verify!(DSPEx.MockAdapter)
end

# lib/dspex/program.ex - Enhanced for dependency injection
defmodule DSPEx.Program do
  # Add configuration resolution for testability
  def resolve_client(program, opts) do
    cond do
      # Explicit injection (highest priority - for tests)
      client = Keyword.get(opts, :client) -> client
      # Program-specific client
      Map.has_key?(program, :client) -> program.client
      # Global test configuration
      Application.get_env(:dspex, :test_client) -> Application.get_env(:dspex, :test_client)
      # Default production behavior
      true -> DSPEx.Client
    end
  end
  
  def resolve_adapter(program, opts) do
    cond do
      adapter = Keyword.get(opts, :adapter) -> adapter
      Map.has_key?(program, :adapter) -> program.adapter
      Application.get_env(:dspex, :test_adapter) -> Application.get_env(:dspex, :test_adapter)
      true -> DSPEx.Adapter
    end
  end
  
  def resolve_evaluator(opts) do
    cond do
      evaluator = Keyword.get(opts, :evaluator) -> evaluator
      Application.get_env(:dspex, :test_evaluator) -> Application.get_env(:dspex, :test_evaluator)
      true -> DSPEx.Evaluate
    end
  end
end

# test/support/dependency_injection.exs
defmodule DSPEx.TestDependencies do
  @moduledoc """
  Helpers for injecting mock dependencies into programs and tests.
  """
  
  def configure_test_environment() do
    Application.put_env(:dspex, :test_client, DSPEx.MockClient)
    Application.put_env(:dspex, :test_adapter, DSPEx.MockAdapter)
    Application.put_env(:dspex, :test_evaluator, DSPEx.MockEvaluator)
  end
  
  def reset_test_environment() do
    Application.delete_env(:dspex, :test_client)
    Application.delete_env(:dspex, :test_adapter)
    Application.delete_env(:dspex, :test_evaluator)
  end
  
  def with_mocked_dependencies(test_fn) do
    configure_test_environment()
    
    try do
      test_fn.()
    after
      reset_test_environment()
    end
  end
end

# test/integration/behaviour_contracts_test.exs
defmodule DSPEx.BehaviourContractsTest do
  use ExUnit.Case, async: true
  
  describe "Client behaviour contract verification" do
    test "all client implementations follow DSPEx.Client behaviour" do
      implementations = [
        DSPEx.ClientManager,  # If it implements the behaviour
        DSPEx.MockClientManager  # Should be updated to implement behaviour
      ]
      
      for impl <- implementations do
        assert behaviour_implemented?(impl, DSPEx.Client),
               "#{impl} does not implement DSPEx.Client behaviour"
      end
    end
  end
  
  describe "Adapter behaviour contract verification" do
    test "all adapter implementations follow DSPEx.Adapter behaviour" do
      implementations = [
        DSPEx.Adapters.InstructorLiteGemini,  # Should implement behaviour
        # Add others as they're converted
      ]
      
      for impl <- implementations do
        assert behaviour_implemented?(impl, DSPEx.Adapter),
               "#{impl} does not implement DSPEx.Adapter behaviour"
      end
    end
  end
  
  defp behaviour_implemented?(module, behaviour) do
    case Code.ensure_loaded(module) do
      {:module, ^module} ->
        behaviours = module.module_info(:attributes) |> Keyword.get(:behaviour, [])
        behaviour in behaviours
      _ ->
        false
    end
  end
end

# test/integration/full_stack_mock_test.exs
defmodule DSPEx.FullStackMockTest do
  use ExUnit.Case, async: true
  import Mox
  
  setup :verify_on_exit!
  setup :set_mox_from_context
  
  test "complete teleprompter workflow with fully mocked stack" do
    # This demonstrates the power of the behaviour-based approach
    
    # Setup teacher program mock - generates "good" demonstrations
    expect(DSPEx.MockProgram, :forward, 5, fn _program, inputs, _opts ->
      case Map.get(inputs, :question) do
        q when q =~ "easy" -> {:ok, %{answer: "easy correct answer"}}
        q when q =~ "hard" -> {:ok, %{answer: "hard correct answer"}}
        _ -> {:ok, %{answer: "generic correct answer"}}
      end
    end)
    
    # Setup evaluator mock - scores demonstrations
    expect(DSPEx.MockEvaluator, :run, 3, fn _program, _examples, _metric_fn, _opts ->
      {:ok, %{
        score: 0.9,  # High score for "good" demonstrations
        stats: %{total_examples: 10, successful: 9, failed: 1, 
                 duration_ms: 500, success_rate: 0.9, throughput: 20.0, errors: []}
      }}
    end)
    
    # Setup client mock for any internal requests
    expect(DSPEx.MockClient, :request, fn _messages, _opts ->
      {:ok, %{choices: [%{message: %{role: "assistant", content: "Generated instruction"}}]}}
    end)
    
    # Setup adapter mock
    DSPEx.BehaviourMocks.setup_standard_adapter_mock()
    
    # Execute teleprompter with fully mocked dependencies
    student = %DSPEx.Predict{signature: TestSignature, client: DSPEx.MockClient}
    teacher = %DSPEx.Predict{signature: TestSignature, client: DSPEx.MockClient}
    trainset = [
      DSPEx.Example.new(%{question: "easy question"}, %{answer: "easy answer"}),
      DSPEx.Example.new(%{question: "hard question"}, %{answer: "hard answer"})
    ]
    
    teleprompter = DSPEx.Teleprompter.BootstrapFewShot.new(evaluator: DSPEx.MockEvaluator)
    
    {:ok, optimized_student} = DSPEx.Teleprompter.compile(
      teleprompter, student, teacher, trainset, 
      &DSPEx.Teleprompter.exact_match(:answer)
    )
    
    # Verify the optimization worked
    assert optimized_student != student
    assert length(optimized_student.demos) > 0
    
    # Verify all mocks were called as expected
    verify!(DSPEx.MockProgram)
    verify!(DSPEx.MockEvaluator)
    verify!(DSPEx.MockClient)
    verify!(DSPEx.MockAdapter)
  end
end

# lib/dspex/behaviours/nx_test_helpers.ex
defmodule DSPEx.Behaviours.NxTestHelpers do
  @moduledoc """
  Nx-powered testing utilities for behaviour validation.
  
  Provides numerical analysis of behaviour implementations to ensure
  performance consistency, output quality, and statistical reliability.
  """
  
  import Nx.Defn
  
  @doc """
  Analyze behaviour implementation performance characteristics.
  """
  def analyze_behaviour_performance(implementation_results) do
    times = Enum.map(implementation_results, & &1.time)
    qualities = Enum.map(implementation_results, & &1.quality)
    
    time_tensor = Nx.tensor(times)
    quality_tensor = Nx.tensor(qualities)
    
    analyze_performance_impl(time_tensor, quality_tensor)
  end
  
  defn analyze_performance_impl(times, qualities) do
    # Performance metrics
    mean_time = Nx.mean(times)
    time_stability = 1.0 / (1.0 + Nx.standard_deviation(times) / mean_time)
    
    # Quality metrics
    mean_quality = Nx.mean(qualities)
    quality_consistency = 1.0 / (1.0 + Nx.standard_deviation(qualities))
    
    # Combined performance score
    performance_score = (time_stability + quality_consistency + mean_quality) / 3.0
    
    %{
      performance_score: performance_score,
      mean_time: mean_time,
      time_stability: time_stability,
      mean_quality: mean_quality,
      quality_consistency: quality_consistency
    }
  end
  
  @doc """
  Property-based testing for behaviour contracts using Nx.
  """
  def validate_behaviour_contract(behaviour_module, implementation, test_data) do
    results = 
      test_data
      |> Enum.map(fn test_case ->
        validate_single_contract(behaviour_module, implementation, test_case)
      end)
    
    analyze_contract_validation(results)
  end
  
  defp validate_single_contract(behaviour_module, implementation, test_case) do
    try do
      # Check if implementation satisfies behaviour
      case apply(implementation, :module_info, [:attributes]) do
        attributes ->
          behaviours = Keyword.get_values(attributes, :behaviour)
          
          if behaviour_module in behaviours do
            # Run test case
            result = run_test_case(implementation, test_case)
            validate_result_properties(result, test_case)
          else
            {:error, :behaviour_not_implemented}
          end
      end
    rescue
      error -> {:error, {:exception, error}}
    end
  end
  
  defp analyze_contract_validation(validation_results) do
    success_count = Enum.count(validation_results, fn result -> 
      case result do
        %{all_valid: true} -> true
        _ -> false
      end
    end)
    
    total_count = length(validation_results)
    success_rate = success_count / total_count
    
    # Extract numerical metrics where available
    numerical_results = extract_numerical_validation_metrics(validation_results)
    
    case numerical_results do
      [] ->
        %{success_rate: success_rate, numerical_analysis: :unavailable}
      
      metrics ->
        tensor = Nx.tensor(metrics)
        numerical_analysis = analyze_numerical_validation(tensor)
        
        Map.merge(%{success_rate: success_rate}, numerical_analysis)
    end
  end
  
  defn analyze_numerical_validation(metrics) do
    %{
      mean_metric: Nx.mean(metrics),
      metric_stability: 1.0 / (1.0 + Nx.standard_deviation(metrics)),
      min_metric: Nx.reduce_min(metrics),
      max_metric: Nx.reduce_max(metrics)
    }
  end
  
  # Helper functions
  
  defp run_test_case(implementation, test_case) do
    function_name = Map.get(test_case, :function)
    args = Map.get(test_case, :args, [])
    
    apply(implementation, function_name, args)
  end
  
  defp validate_result_properties(result, test_case) do
    expected_properties = Map.get(test_case, :expected_properties, %{})
    
    property_validations = 
      Enum.map(expected_properties, fn {property, expected_value} ->
        actual_value = extract_property(result, property)
        validate_property_match(property, expected_value, actual_value)
      end)
    
    %{
      result: result,
      property_validations: property_validations,
      all_valid: Enum.all?(property_validations, & &1.valid)
    }
  end
  
  defp extract_property(result, property) do
    case result do
      map when is_map(map) -> Map.get(map, property)
      _ -> nil
    end
  end
  
  defp validate_property_match(property, expected, actual) do
    %{
      property: property,
      expected: expected,
      actual: actual,
      valid: expected == actual
    }
  end
  
  defp extract_numerical_validation_metrics(validation_results) do
    validation_results
    |> Enum.flat_map(fn result ->
      case result do
        %{result: %{score: score}} when is_number(score) -> [score]
        %{result: %{quality: quality}} when is_number(quality) -> [quality]
        _ -> []
      end
    end)
  end
end

# config/config.exs - Nx Configuration for Behaviour Testing
config :dspex, :behaviours,
  # Nx backend configuration
  nx_backend: {Nx.BinaryBackend, []},
  
  # Testing and validation settings
  testing: %{
    performance_analysis: true,
    statistical_comparison: true,
    property_validation: true,
    numerical_precision: 1.0e-6
  },
  
  # Mock behaviour validation
  mock_validation: %{
    enabled: true,
    pattern_analysis: true,
    consistency_threshold: 0.8,
    response_tracking: true
  },
  
  # Implementation comparison
  comparison: %{
    min_samples: 10,
    confidence_threshold: 0.8,
    effect_size_threshold: 0.5
  }

# mix.exs - Add Nx dependency for behaviour testing
defp deps do
  [
    # ... existing dependencies ...
    {:nx, "~> 0.6"},              # Numerical computing for behaviour testing
    {:mox, "~> 1.0", only: :test}, # Already present - Mock testing
    # ... other dependencies ...
  ]
end