defmodule Superintelligence.AIAgent.Brain do
  use GenServer
  require Logger

  @moduledoc """
  The Brain module handles AI reasoning and decision making using GPT-4.
  It manages the conversation context and tool usage.
  """

  @openai_url "https://api.openai.com/v1/chat/completions"
  @model "gpt-4-turbo-preview"
  @max_tokens 4096
  @temperature 0.7

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

  @impl true
  def init(_opts) do
    # Initialize built-in tools
    Superintelligence.AIAgent.ToolRegistry.init_builtin_tools()
    
    {:ok, %{
      api_key: System.get_env("OPENAI_API_KEY"),
      conversation_history: [],
      max_iterations: 50,
      system_prompt: build_system_prompt()
    }}
  end

  # Client API

  def process_task(task_description) do
    GenServer.call(__MODULE__, {:process_task, task_description}, :infinity)
  end

  def get_conversation_history do
    GenServer.call(__MODULE__, :get_conversation_history)
  end

  def clear_history do
    GenServer.call(__MODULE__, :clear_history)
  end

  # Server Callbacks

  @impl true
  def handle_call({:process_task, task_description}, _from, state) do
    Logger.info("Processing task: #{task_description}")
    
    # Add user message to history
    user_message = %{role: "user", content: task_description}
    messages = [%{role: "system", content: state.system_prompt} | state.conversation_history] ++ [user_message]
    
    # Process with GPT-4 in a loop until task is complete
    result = process_with_iterations(messages, state, 0)
    
    {:reply, result, state}
  end

  @impl true
  def handle_call(:get_conversation_history, _from, state) do
    {:reply, state.conversation_history, state}
  end

  @impl true
  def handle_call(:clear_history, _from, state) do
    {:reply, :ok, %{state | conversation_history: []}}
  end

  # Private Functions

  defp process_with_iterations(messages, state, iteration) when iteration >= state.max_iterations do
    {:error, "Max iterations reached"}
  end

  defp process_with_iterations(messages, state, iteration) do
    Logger.info("Iteration #{iteration + 1} running...")
    
    # Get available tools
    tools = Superintelligence.AIAgent.ToolRegistry.list_tools()
    
    # Make API call to GPT-4
    case call_gpt4(messages, tools, state.api_key) do
      {:ok, response} ->
        # Handle the response
        handle_gpt_response(response, messages, state, iteration)
      
      {:error, reason} ->
        {:error, reason}
    end
  end

  defp handle_gpt_response(response, messages, state, iteration) do
    # Extract message and tool calls from response
    message = response["choices"] |> List.first() |> Map.get("message")
    
    # Add assistant message to history
    new_messages = messages ++ [message]
    
    # Log the response
    if message["content"] do
      Logger.info("GPT-4 Response: #{message["content"]}")
    end
    
    # Check for tool calls
    case message["tool_calls"] do
      nil ->
        # No tool calls, return the response
        {:ok, message["content"]}
      
      tool_calls ->
        # Execute tool calls
        tool_results = execute_tool_calls(tool_calls)
        
        # Add tool results to messages
        tool_messages = Enum.map(tool_results, fn {tool_call_id, result} ->
          %{
            role: "tool",
            tool_call_id: tool_call_id,
            content: Jason.encode!(result)
          }
        end)
        
        final_messages = new_messages ++ tool_messages
        
        # Check if task is completed
        if task_completed?(tool_calls) do
          {:ok, "Task completed successfully"}
        else
          # Continue processing
          process_with_iterations(final_messages, state, iteration + 1)
        end
    end
  end

  defp execute_tool_calls(tool_calls) do
    Enum.map(tool_calls, fn tool_call ->
      function_name = tool_call["function"]["name"]
      args = Jason.decode!(tool_call["function"]["arguments"])
      
      Logger.info("Calling tool: #{function_name} with args: #{inspect(args)}")
      
      result = case Superintelligence.AIAgent.ToolRegistry.execute_tool(function_name, args) do
        {:ok, res} -> res
        {:error, err} -> err
      end
      
      Logger.info("Result of #{function_name}: #{inspect(result)}")
      
      {tool_call["id"], result}
    end)
  end

  defp task_completed?(tool_calls) do
    Enum.any?(tool_calls, fn tc -> tc["function"]["name"] == "task_completed" end)
  end

  defp call_gpt4(messages, tools, api_key) do
    headers = [
      {"Authorization", "Bearer #{api_key}"},
      {"Content-Type", "application/json"}
    ]
    
    body = %{
      model: @model,
      messages: messages,
      tools: tools,
      tool_choice: "auto",
      max_tokens: @max_tokens,
      temperature: @temperature
    }
    
    case HTTPoison.post(@openai_url, Jason.encode!(body), headers, recv_timeout: 60_000) do
      {:ok, %HTTPoison.Response{status_code: 200, body: body}} ->
        {:ok, Jason.decode!(body)}
      
      {:ok, %HTTPoison.Response{status_code: status, body: body}} ->
        {:error, "API call failed with status #{status}: #{body}"}
      
      {:error, %HTTPoison.Error{reason: reason}} ->
        {:error, "HTTP request failed: #{reason}"}
    end
  end

  defp build_system_prompt do
    api_keys = list_available_api_keys()
    
    """
    You are an AI assistant designed to iteratively build and execute Elixir functions using tools provided to you.
    Your task is to complete the requested task by creating and using tools in a loop until the task is fully done.
    Do not ask for user input until you find it absolutely necessary. If you need required information that is likely available online, create the required tools to find this information.
    
    You have the following tools available to start with:
    
    1. **create_or_update_tool**: This tool allows you to create new functions or update existing ones.
       You must provide the function name, code, description, and parameters.
       The code should be valid Elixir code that will be executed within a function.
       Example of 'parameters': %{
         "param1" => %{"type" => "string", "description" => "Description of param1"},
         "param2" => %{"type" => "integer", "description" => "Description of param2"}
       }
    
    2. **install_package**: Installs an Elixir package using mix.
    
    3. **execute_bash**: Executes a bash command.
    
    4. **read_file**: Reads a file from the filesystem.
    
    5. **write_file**: Writes content to a file.
    
    6. **task_completed**: This tool should be used to signal when you believe the requested task is fully completed.
    
    Available API keys in the environment:
    #{api_keys}
    
    Your workflow should include:
    - Creating or updating tools with all required arguments
    - Using 'install_package' when a required library is missing
    - Using created tools to progress towards completing the task
    - When creating or updating tools, provide the complete Elixir code
    - Handling any errors by adjusting your tools or arguments as necessary
    - Being token-efficient: avoid returning excessively long outputs
    - Prioritize using tools that you have access to via the available API keys
    - Signaling task completion with 'task_completed' when done
    
    Please ensure that all function calls include all required parameters.
    """
  end

  defp list_available_api_keys do
    api_key_patterns = ["API_KEY", "ACCESS_TOKEN", "SECRET_KEY", "TOKEN", "APISECRET"]
    
    available_keys = System.get_env()
      |> Enum.filter(fn {key, _value} ->
        Enum.any?(api_key_patterns, &String.contains?(String.upcase(key), &1))
      end)
      |> Enum.map(fn {key, _value} -> "- #{key}" end)
      |> Enum.join("\n")
    
    if available_keys == "" do
      "No API keys are available."
    else
      available_keys
    end
  end
end