defmodule DSPex.LLM.Adapters.Mock do
  @behaviour DSPex.LLM.Adapter
  
  defstruct [
    :responses,           # Predefined responses
    :response_pattern,    # :sequential | :random | :function
    :delay,              # Simulate network delay
    :failure_rate,       # Simulate failures
    :call_history,       # Track calls for assertions
    :behavior_overrides  # Per-test customization
  ]
end

# Static responses
%Mock{
  responses: %{
    "What is 2+2?" => "4",
    "default" => "This is a mock response"
  }
}

# Pattern-based responses
%Mock{
  response_pattern: :function,
  response_fn: fn prompt, opts ->
    cond do
      String.contains?(prompt, "error") ->
        {:error, :simulated_error}
        
      opts[:schema] ->
        generate_mock_structured(opts[:schema])
        
      true ->
        "Mock response for: #{prompt}"
    end
  end
}

# Sequential responses
%Mock{
  response_pattern: :sequential,
  responses: ["First", "Second", "Third"],
  current_index: 0
}

defp generate_mock_structured(schema) do
  schema
  |> Enum.map(fn {key, type} ->
    {key, generate_mock_value(type)}
  end)
  |> Map.new()
end

defp generate_mock_value(:string), do: "mock_string"
defp generate_mock_value(:integer), do: 42
defp generate_mock_value(:float), do: 3.14
defp generate_mock_value(:boolean), do: true
defp generate_mock_value({:array, type}), do: [generate_mock_value(type)]
defp generate_mock_value({:optional, type}), do: generate_mock_value(type)
defp generate_mock_value(map) when is_map(map) do
  generate_mock_structured(map)
end

def generate(adapter, prompt, opts) do
  # Record call
  record_call(adapter, {:generate, prompt, opts})
  
  # Simulate delay
  if adapter.delay do
    Process.sleep(adapter.delay)
  end
  
  # Simulate failures
  if should_fail?(adapter) do
    {:error, select_error_type(adapter)}
  else
    {:ok, get_mock_response(adapter, prompt, opts)}
  end
end

defp should_fail?(%{failure_rate: rate}) when is_number(rate) do
  :rand.uniform() < rate
end

defp select_error_type(adapter) do
  Enum.random([
    :timeout,
    {:api_error, 500, "Internal server error"},
    {:rate_limited, 60},
    :network_error
  ])
end

defmodule MockHistory do
  use Agent
  
  def start_link do
    Agent.start_link(fn -> [] end, name: __MODULE__)
  end
  
  def record_call(call) do
    Agent.update(__MODULE__, &[call | &1])
  end
  
  def get_calls do
    Agent.get(__MODULE__, &Enum.reverse(&1))
  end
  
  def clear do
    Agent.update(__MODULE__, fn _ -> [] end)
  end
  
  def assert_called(expected) do
    calls = get_calls()
    unless Enum.any?(calls, &match?(^expected, &1)) do
      raise "Expected call not found: #{inspect(expected)}"
    end
  end
end

defmodule DSPex.LLM.MockHelpers do
  def with_mock_responses(responses, fun) do
    adapter = %Mock{responses: responses}
    Process.put(:mock_adapter, adapter)
    
    try do
      fun.(adapter)
    after
      Process.delete(:mock_adapter)
      MockHistory.clear()
    end
  end
  
  def assert_llm_called(prompt) do
    MockHistory.assert_called({:generate, prompt, _})
  end
  
  def stub_structured_response(schema, response) do
    adapter = Process.get(:mock_adapter)
    updated = %{adapter | 
      behavior_overrides: Map.put(
        adapter.behavior_overrides || %{},
        schema,
        response
      )
    }
    Process.put(:mock_adapter, updated)
  end
end

# In tests
defmodule MyTest do
  use ExUnit.Case
  import DSPex.LLM.MockHelpers
  
  test "processes user query" do
    with_mock_responses(%{
      "Hello" => "Hi there!",
      "default" => "Mock response"
    }, fn adapter ->
      # Your test code
      result = MyModule.process("Hello")
      assert result == "Hi there!"
      
      # Verify LLM was called
      assert_llm_called("Hello")
    end)
  end
  
  test "handles structured output" do
    adapter = %Mock{
      response_pattern: :function,
      response_fn: fn _prompt, opts ->
        if opts[:schema] do
          %{
            "name" => "Test User",
            "age" => 25,
            "active" => true
          }
        end
      end
    }
    
    {:ok, user} = DSPex.LLM.Adapter.generate(
      adapter,
      "Generate user",
      schema: %{name: :string, age: :integer, active: :boolean}
    )
    
    assert user["name"] == "Test User"
  end
end