# Direct enhancement of existing DSPEx modules with ElixirML foundation
defmodule DSPEx.Program do
  @moduledoc """
  Enhanced DSPEx Program using ElixirML foundation components.
  Leverages Variable System for automatic optimization and Schema Engine for validation.
  """
  
  use ElixirML.Resource
  alias ElixirML.{Schema, Variable, Process}
  
  defstruct [
    :signature,
    :variable_space,
    :config,
    :metadata
  ]
  
  def new(signature, opts \\ []) do
    # Create program with variable space for automatic optimization
    variable_space = Variable.Space.from_signature(signature, opts)
    
    %__MODULE__{
      signature: Schema.from_signature(signature),
      variable_space: variable_space,
      config: Keyword.get(opts, :config, %{}),
      metadata: %{created_at: DateTime.utc_now()}
    }
  end
  
  def forward(program, inputs, opts \\ []) do
    # Use Process Orchestrator for execution
    Process.Pipeline.execute(program, inputs, opts)
  end
end

# mix.exs updates
defp deps do
  [
    # Existing dependencies preserved
    {:finch, "~> 0.16"},
    {:jason, "~> 1.4"},
    {:telemetry, "~> 1.2"},
    
    # New foundation dependencies
    {:ash, "~> 3.0"},           # Resource Framework
    {:sinter, "~> 0.1"},       # Enhanced schema validation
    {:nx, "~> 0.7"},           # Numerical computing
    {:gen_stage, "~> 1.2"},    # Process orchestration
    
    # Development/testing
    {:stream_data, "~> 0.6", only: [:test]},
    {:mox, "~> 1.0", only: [:test]}
  ]
end

defmodule ElixirML.Schema do
  @moduledoc """
  Enhanced schema system for ElixirML/DSPEx with ML-specific types,
  compile-time optimization, and comprehensive validation.
  """

  defmacro defschema(name, do: block) do
    quote do
      defmodule unquote(name) do
        use ElixirML.Schema.Definition
        import ElixirML.Schema.DSL
        
        unquote(block)
        
        @before_compile ElixirML.Schema.Compiler
      end
    end
  end

  # Core DSL macros
  defmacro field(name, type, opts \\ []) do
    quote do
      @fields {unquote(name), unquote(type), unquote(opts)}
    end
  end

  defmacro validation(name, do: block) do
    quote do
      @validations {unquote(name), unquote(block)}
    end
  end

  defmacro transform(name, do: block) do
    quote do
      @transforms {unquote(name), unquote(block)}
    end
  end
end

defmodule ElixirML.Schema.Types do
  @type ml_type :: 
    :embedding |
    :tensor |
    :token_list |
    :probability |
    :confidence_score |
    :model_response |
    :variable_config

  def validate_type(value, :embedding) do
    case value do
      list when is_list(list) and length(list) > 0 ->
        if Enum.all?(list, &is_number/1) do
          {:ok, value}
        else
          {:error, "Embedding must be list of numbers"}
        end
      _ -> {:error, "Invalid embedding format"}
    end
  end

  def validate_type(value, :probability) do
    case value do
      num when is_number(num) and num >= 0.0 and num <= 1.0 ->
        {:ok, value}
      _ -> {:error, "Probability must be between 0.0 and 1.0"}
    end
  end
end

defmodule DSPEx.Signature do
  @moduledoc """
  Enhanced signature system using ElixirML.Schema foundation.
  Provides ML-specific validation and automatic variable extraction.
  """

  use ElixirML.Schema
  alias ElixirML.{Schema, Variable}

  defmacro __using__(opts) do
    quote do
      use ElixirML.Schema
      import DSPEx.Signature.DSL
      
      @before_compile DSPEx.Signature.Compiler
    end
  end

  defmacro input(field, type, opts \\ []) do
    quote do
      field unquote(:"input_#{field}"), unquote(type), unquote([role: :input] ++ opts)
    end
  end

  defmacro output(field, type, opts \\ []) do
    quote do
      field unquote(:"output_#{field}"), unquote(type), unquote([role: :output] ++ opts)
  end
end

  defmacro instructions(text) do
    quote do
      @instructions unquote(text)
    end
  end

  # Enhanced with variable extraction
  def extract_variables(signature_module) do
    Schema.Variable.Generator.extract_variables(signature_module)
  end

  # Enhanced with ML-specific validation
  def validate_io(signature_module, inputs, outputs) do
    with {:ok, _} <- validate_inputs(signature_module, inputs),
         {:ok, _} <- validate_outputs(signature_module, outputs) do
      {:ok, {inputs, outputs}}
    end
  end
end

# Temperature that adapts based on model capabilities
ElixirML.Variable.MLTypes.conditional_temperature(:temperature,
  conditions: [
    {provider: :openai, model: "gpt-4"} => {range: {0.0, 2.0}, default: 0.7},
    {provider: :groq} => {range: {0.1, 0.9}, default: 0.5}
  ],
  fallback: {range: {0.5, 0.8}, default: 0.7}
)

defmodule ElixirML.Variable.MLTypes do
  @moduledoc """
  ML-specific variable types with built-in compatibility and optimization logic.
  """

  alias ElixirML.Variable

  @doc "Create a provider variable with model compatibility"
  def provider(name, opts \\ []) do
    providers = Keyword.get(opts, :providers, [:openai, :anthropic, :groq, :google])
    
    Variable.choice(name, providers, 
      description: "LLM Provider selection",
      hints: [
        optimization_priority: :cost_performance,
        compatibility_aware: true
      ],
      metadata: %{
        cost_weights: %{
          openai: 1.0,
          anthropic: 1.2,
          groq: 0.3,
          google: 0.8
        },
        performance_weights: %{
          openai: 0.9,
          anthropic: 0.95,
          groq: 0.7,
          google: 0.85
        }
      }
    )
  end

  @doc "Create an adapter variable with capability awareness"
  def adapter(name, opts \\ []) do
    adapters = Keyword.get(opts, :adapters, [
      ElixirML.Adapter.JSON,
      ElixirML.Adapter.Markdown,
      ElixirML.Adapter.Chat
    ])
    
    Variable.module(name,
      modules: adapters,
      behavior: ElixirML.Adapter,
      description: "Response format adapter",
      hints: [
        performance_impact: :medium,
        model_compatibility: true
      ]
    )
  end

  @doc "Create a reasoning strategy variable"
  def reasoning_strategy(name, opts \\ []) do
    strategies = Keyword.get(opts, :strategies, [
      ElixirML.Reasoning.Predict,
      ElixirML.Reasoning.ChainOfThought,
      ElixirML.Reasoning.ProgramOfThought,
      ElixirML.Reasoning.ReAct
    ])
    
    Variable.module(name,
      modules: strategies,
      behavior: ElixirML.Reasoning.Strategy,
      description: "Reasoning strategy selection",
      hints: [
        performance_impact: :high,
        accuracy_impact: :high,
        cost_impact: :high
      ]
    )
  end
end

defmodule DSPEx.Program do
  @moduledoc """
  Enhanced DSPEx Program with automatic variable extraction and optimization.
  """

  use ElixirML.Resource
  alias ElixirML.{Variable, Schema}

  defstruct [
    :signature,
    :variable_space,
    :config,
    :optimizations,
    :metadata
  ]

  def new(signature, opts \\ []) do
    # Extract variables from signature and configuration
    variables = Variable.MLTypes.extract_from_signature(signature)
    
    # Add common ML variables
    enhanced_variables = variables
    |> Variable.Space.add_variable(Variable.MLTypes.provider(:provider))
    |> Variable.Space.add_variable(Variable.MLTypes.adapter(:adapter))
    |> Variable.Space.add_variable(Variable.MLTypes.reasoning_strategy(:reasoning))
    |> Variable.Space.add_variable(Variable.float(:temperature, range: {0.0, 2.0}, default: 0.7))

    %__MODULE__{
      signature: signature,
      variable_space: enhanced_variables,
      config: Keyword.get(opts, :config, %{}),
      optimizations: [],
      metadata: %{created_at: DateTime.utc_now()}
    }
  end

  def optimize(program, training_data, opts \\ []) do
    # Enhanced SIMBA integration with Variable System
    DSPEx.Teleprompter.SIMBA.optimize(program, training_data, opts)
  end
end

defmodule ElixirML.Resource do
  @moduledoc """
  Ash-inspired resource framework treating programs, optimizations, 
  and configurations as first-class resources with relationships,
  validations, and lifecycle hooks.
  """

  defmacro __using__(opts) do
    quote do
      use Ash.Resource, unquote(opts)
      
      # ElixirML-specific extensions
      import ElixirML.Resource.Attributes
      import ElixirML.Resource.Relationships
      import ElixirML.Resource.Actions
      
      # Automatic variable tracking
      Module.register_attribute(__MODULE__, :variables, accumulate: true)
      
      # Lifecycle callbacks
      @before_compile ElixirML.Resource.Compiler
    end
  end

  # Resource-specific attributes
  defmacro ml_attribute(name, type, opts \\ []) do
    quote do
      attribute unquote(name), unquote(type), unquote(opts)
      
      if unquote(opts[:variable]) do
        @variables {unquote(name), unquote(type), unquote(opts)}
      end
    end
  end
end

defmodule ElixirML.Resources.Program do
  use Ash.Resource,
    domain: ElixirML.Domain,
    data_layer: AshPostgres.DataLayer

  use ElixirML.Schema.ResourceIntegration

  attributes do
    uuid_primary_key :id

    attribute :name, :string do
      allow_nil? false
      constraints max_length: 255
    end

    attribute :type, :atom do
      allow_nil? false
      constraints one_of: [:predict, :chain_of_thought, :react, :custom]
    end

    schema_attribute :signature_config, ElixirML.Schemas.Program.SignatureConfig do
      allow_nil? false
    end

    schema_attribute :program_config, ElixirML.Schemas.Program.ProgramConfig do
      default %{}
    end

    attribute :status, :atom do
      default :draft
      constraints one_of: [:draft, :active, :optimizing, :archived]
    end

    schema_attribute :performance_metrics, ElixirML.Schemas.Program.PerformanceMetrics do
      default %{}
    end
  end

  relationships do
    belongs_to :variable_space, ElixirML.Resources.VariableSpace
    has_many :optimization_runs, ElixirML.Resources.OptimizationRun
    has_many :executions, ElixirML.Resources.Execution
  end

  actions do
    defaults [:create, :read, :update, :destroy]

    action :execute, ElixirML.Resources.Execution do
      argument :inputs, :map, allow_nil?: false
      argument :variable_configuration, :map, default: %{}
      run ElixirML.Actions.ExecuteProgram
    end

    action :optimize, ElixirML.Resources.OptimizationRun do
      argument :training_data, {:array, :map}, allow_nil?: false
      argument :optimization_strategy, :atom, default: :simba
      run ElixirML.Actions.OptimizeProgram
    end
  end

  calculations do
    calculate :current_performance_score, :float, ElixirML.Calculations.CurrentPerformanceScore
    calculate :optimization_count, :integer, ElixirML.Calculations.OptimizationCount
    calculate :variable_importance_scores, {:array, :map}, ElixirML.Calculations.VariableImportanceScores
  end
end

defmodule ElixirML.Resources.VariableSpace do
  use Ash.Resource,
    domain: ElixirML.Domain,
    data_layer: AshPostgres.DataLayer

  attributes do
    uuid_primary_key :id
    attribute :name, :string, allow_nil?: false

    schema_attribute :variable_definitions, ElixirML.Schemas.Variable.VariableDefinitions do
      allow_nil? false
    end

    schema_attribute :constraints, ElixirML.Schemas.Variable.Constraints do
      default %{dependencies: [], validations: []}
    end

    attribute :discrete_space_size, :integer, default: 1
    attribute :continuous_dimensions, :integer, default: 0
  end

  relationships do
    has_many :programs, ElixirML.Resources.Program
    has_many :optimization_runs, ElixirML.Resources.OptimizationRun
    has_many :configurations, ElixirML.Resources.VariableConfiguration
  end

  actions do
    action :generate_configuration, ElixirML.Resources.VariableConfiguration do
      argument :strategy, :atom, default: :random
      run ElixirML.Actions.GenerateVariableConfiguration
    end

    action :validate_configuration, :boolean do
      argument :configuration, :map, allow_nil?: false
      run ElixirML.Actions.ValidateConfiguration
    end
  end

  calculations do
    calculate :variable_count, :integer, ElixirML.Calculations.VariableCount
    calculate :complexity_score, :float, ElixirML.Calculations.ComplexityScore
    calculate :optimization_difficulty, :atom, ElixirML.Calculations.OptimizationDifficulty
  end
end

defmodule ElixirML.Resources.OptimizationRun do
  use Ash.Resource,
    domain: ElixirML.Domain,
    data_layer: AshPostgres.DataLayer

  attributes do
    uuid_primary_key :id
    attribute :name, :string, allow_nil?: false
    attribute :strategy, :atom, allow_nil?: false
    attribute :status, :atom, default: :pending

    schema_attribute :configuration, ElixirML.Schemas.Optimization.Configuration do
      allow_nil? false
    end

    schema_attribute :results, ElixirML.Schemas.Optimization.Results do
      default %{}
    end

    attribute :iterations_completed, :integer, default: 0
    attribute :best_score, :float
    attribute :started_at, :utc_datetime
    attribute :completed_at, :utc_datetime
  end

  relationships do
    belongs_to :program, ElixirML.Resources.Program
    belongs_to :variable_space, ElixirML.Resources.VariableSpace
    has_many :evaluations, ElixirML.Resources.Evaluation
  end

  actions do
    action :start_optimization do
      argument :program_id, :uuid, allow_nil?: false
      run ElixirML.Actions.StartOptimization
    end

    action :update_progress do
      argument :new_evaluation, :map
      run ElixirML.Actions.UpdateOptimizationProgress
    end

    action :complete do
      argument :final_results, :map, allow_nil?: false
      run ElixirML.Actions.CompleteOptimization
    end
  end

  calculations do
    calculate :duration_seconds, :integer, ElixirML.Calculations.OptimizationDuration
    calculate :progress_percentage, :float, ElixirML.Calculations.OptimizationProgress
    calculate :convergence_status, :atom, ElixirML.Calculations.ConvergenceStatus
  end
end

defmodule ElixirML.Process.Orchestrator do
  @moduledoc """
  Advanced supervision and process management for ElixirML.
  Every major component runs in its own supervised process for fault tolerance.
  """

  use Supervisor

  def start_link(opts) do
    Supervisor.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @impl true
  def init(_opts) do
    children = [
      # Core services
      {ElixirML.Process.SchemaRegistry, []},
      {ElixirML.Process.VariableRegistry, []},
      {ElixirML.Process.ResourceManager, []},
      
      # Execution services
      {ElixirML.Process.ProgramSupervisor, []},
      {ElixirML.Process.PipelinePool, []},
      {ElixirML.Process.ClientPool, []},
      
      # Intelligence services
      {ElixirML.Process.TeleprompterSupervisor, []},
      {ElixirML.Process.EvaluationWorkers, []},
      
      # Integration services
      {ElixirML.Process.ToolRegistry, []},
      {ElixirML.Process.DatasetManager, []}
    ]
    
    Supervisor.init(children, strategy: :one_for_one)
  end
end

defmodule ElixirML.Process.SchemaRegistry do
  @moduledoc """
  High-performance registry for schema validation and caching.
  Uses ETS tables for fast lookup and LRU eviction.
  """

  use GenServer

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  def init(opts) do
    table_name = Keyword.get(opts, :table_name, :schema_cache)
    max_size = Keyword.get(opts, :max_size, 10_000)
    
    :ets.new(table_name, [:named_table, :public, :set])
    
    {:ok, %{
      table: table_name,
      max_size: max_size,
      current_size: 0,
      access_order: :queue.new()
    }}
  end

  def get_cached_validation(schema_module, data_hash) do
    case :ets.lookup(:schema_cache, {schema_module, data_hash}) do
      [{_, result, _timestamp}] -> 
        GenServer.cast(__MODULE__, {:access, schema_module, data_hash})
        {:hit, result}
      [] -> 
        :miss
    end
  end

  def cache_validation_result(schema_module, data_hash, result) do
    GenServer.cast(__MODULE__, {:cache, schema_module, data_hash, result})
  end
end

defmodule DSPEx.Teleprompter.SIMBA do
  @moduledoc """
  Enhanced SIMBA teleprompter with ElixirML Variable System integration.
  Enables automatic module selection and multi-objective optimization.
  """

  use ElixirML.Resource
  alias ElixirML.{Variable, Process}

  def optimize(program, training_data, opts \\ []) do
    # Enhanced optimization with Variable System
    variable_space = program.variable_space
    optimization_config = build_optimization_config(opts)
    
    # Start optimization process
    {:ok, optimization_run} = ElixirML.Resources.OptimizationRun.create(%{
      program_id: program.id,
      variable_space_id: variable_space.id,
      strategy: :simba,
      configuration: optimization_config
    })

    # Run SIMBA with variable-aware evaluation
    run_variable_aware_simba(optimization_run, training_data, opts)
  end

  defp run_variable_aware_simba(optimization_run, training_data, opts) do
    # Enhanced SIMBA algorithm with automatic module selection
    max_iterations = Keyword.get(opts, :max_iterations, 100)
    
    # Initialize with variable space sampling
    initial_configurations = Variable.Space.sample_configurations(
      optimization_run.variable_space, 
      count: 10
    )

    # Run optimization loop with multi-objective evaluation
    Enum.reduce_while(1..max_iterations, initial_configurations, fn iteration, configurations ->
      # Evaluate configurations with multi-objective scoring
      evaluated_configs = evaluate_configurations(configurations, training_data)
      
      # Check convergence
      if converged?(evaluated_configs, iteration) do
        {:halt, select_best_configuration(evaluated_configs)}
      else
        # Generate new configurations using SIMBA strategy
        new_configs = generate_next_configurations(evaluated_configs)
        {:cont, new_configs}
      end
    end)
  end

  defp evaluate_configurations(configurations, training_data) do
    # Multi-objective evaluation: accuracy, cost, latency
    Enum.map(configurations, fn config ->
      results = evaluate_single_configuration(config, training_data)
      
      %{
        configuration: config,
        accuracy: results.accuracy,
        cost: results.cost,
        latency: results.latency,
        composite_score: calculate_composite_score(results)
      }
    end)
  end
end

defmodule DSPEx do
  @moduledoc """
  Enhanced DSPEx with ElixirML foundation integration.
  Provides the complete DSPEx API with automatic optimization capabilities.
  """

  alias ElixirML.{Schema, Variable, Resource, Process}

  @doc """
  Create a new DSPEx program with automatic variable extraction.
  """
  def program(signature, opts \\ []) do
    DSPEx.Program.new(signature, opts)
  end

  @doc """
  Execute a program with automatic optimization if configured.
  """
  def forward(program, inputs, opts \\ []) do
    # Use Process Orchestrator for execution
    Process.Pipeline.execute(program, inputs, opts)
  end

  @doc """
  Optimize a program using the enhanced SIMBA teleprompter.
  """
  def optimize(program, training_data, opts \\ []) do
    DSPEx.Teleprompter.SIMBA.optimize(program, training_data, opts)
  end
end