defmodule DSPex.Python.SnakepitConfig do
  @moduledoc """
  Snakepit pool configurations for different workload types
  """
  
  def pools do
    [
      # General purpose pool for lightweight operations
      general: [
        adapter_module: Snakepit.Adapters.GenericPythonV2,
        pool_size: System.schedulers_online() * 2,
        pool_config: %{
          memory_limit: "512MB",
          timeout: 30_000,
          health_check_interval: 30_000
        }
      ],
      
      # Optimizer pool for heavy ML tasks
      optimizer: [
        adapter_module: DSPex.Adapters.OptimizerPython,
        pool_size: 2,
        pool_config: %{
          memory_limit: "4GB",
          timeout: 300_000,  # 5 minutes
          health_check_interval: 60_000
        }
      ],
      
      # Neural pool for GPU operations
      neural: [
        adapter_module: DSPex.Adapters.NeuralPython,
        pool_size: 4,
        pool_config: %{
          memory_limit: "8GB",
          gpu_enabled: true,
          timeout: 600_000,  # 10 minutes
          health_check_interval: 60_000
        }
      ]
    ]
  end
end

defmodule DSPex.Python.Snakepit do
  @moduledoc """
  Snakepit integration for DSPex Python operations
  """
  
  use Supervisor
  
  def start_link(opts) do
    Supervisor.start_link(__MODULE__, opts, name: __MODULE__)
  end
  
  def init(_opts) do
    children = build_pool_specs()
    Supervisor.init(children, strategy: :one_for_one)
  end
  
  defp build_pool_specs do
    DSPex.Python.SnakepitConfig.pools()
    |> Enum.map(fn {name, config} ->
      %{
        id: :"snakepit_pool_#{name}",
        start: {Snakepit.Pool, :start_link, [
          Keyword.put(config, :name, pool_name(name))
        ]}
      }
    end)
  end
  
  def pool_name(type), do: :"dspex_python_#{type}"
  
  # Public API
  def execute(pool_type, command, args, opts \\ []) do
    pool = pool_name(pool_type)
    timeout = opts[:timeout] || default_timeout(pool_type)
    
    Snakepit.execute(pool, command, args, timeout: timeout)
  end
  
  def execute_in_session(session_id, pool_type, command, args, opts \\ []) do
    pool = pool_name(pool_type)
    Snakepit.execute_in_session(session_id, command, args, opts)
  end
end

defmodule DSPex.Python.HealthMonitor do
  use GenServer
  require Logger
  
  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end
  
  def init(_opts) do
    schedule_health_check()
    {:ok, %{pools: DSPex.Python.SnakepitConfig.pools()}}
  end
  
  def handle_info(:health_check, state) do
    Enum.each(state.pools, fn {name, _config} ->
      check_pool_health(name)
    end)
    
    schedule_health_check()
    {:noreply, state}
  end
  
  defp check_pool_health(pool_type) do
    pool = DSPex.Python.Snakepit.pool_name(pool_type)
    
    case Snakepit.get_pool_stats(pool) do
      {:ok, stats} ->
        :telemetry.execute(
          [:dspex, :python, :pool_health],
          %{
            available: stats.available,
            busy: stats.busy,
            total: stats.size
          },
          %{pool: pool_type}
        )
        
      {:error, reason} ->
        Logger.error("Pool health check failed for #{pool_type}: #{inspect(reason)}")
    end
  end
end

defmodule DSPex.Python.SessionManager do
  @moduledoc """
  Manages DSPex sessions with Snakepit session support
  """
  
  def create_session(opts \\ []) do
    session_id = generate_session_id()
    pool_type = opts[:pool_type] || :general
    
    metadata = %{
      created_at: DateTime.utc_now(),
      pool_type: pool_type,
      dspex_context: opts[:context] || %{}
    }
    
    case Snakepit.Bridge.SessionStore.create_session(session_id, 
      ttl: opts[:ttl] || 3600,
      metadata: metadata
    ) do
      {:ok, _} -> {:ok, session_id}
      error -> error
    end
  end
  
  def execute_in_session(session_id, command, args) do
    case get_session_metadata(session_id) do
      {:ok, %{pool_type: pool_type}} ->
        DSPex.Python.Snakepit.execute_in_session(
          session_id,
          pool_type,
          command,
          args
        )
        
      {:error, :not_found} ->
        {:error, :session_not_found}
    end
  end
end

def handle_execute_error({:error, :pool_saturated}) do
  {:error, %{
    type: :resource_exhausted,
    message: "Python pool is at capacity",
    retry_after: 1000
  }}
end

def handle_execute_error({:error, :worker_timeout}) do
  {:error, %{
    type: :timeout,
    message: "Python operation timed out",
    suggestion: "Consider using optimizer pool for long operations"
  }}
end

# Simple execution
{:ok, result} = DSPex.Python.Snakepit.execute(
  :general,
  "echo",
  %{message: "Hello from Python"}
)

# Long-running optimization
{:ok, result} = DSPex.Python.Snakepit.execute(
  :optimizer,
  "optimize_model",
  %{iterations: 1000, data: data},
  timeout: 300_000
)

# Session-based execution
{:ok, session_id} = DSPex.Python.SessionManager.create_session(
  pool_type: :neural,
  ttl: 7200
)

{:ok, result} = DSPex.Python.SessionManager.execute_in_session(
  session_id,
  "train_model",
  %{epochs: 10, batch_size: 32}
)