require 'net/http'
require 'json'

class Llm::Chat
  OPENAI_API_URL = "https://api.openai.com/v1/chat/completions"
  API_KEY = ENV['OPENAI_API_KEY']

  def initialize(payload:)
    @payload = payload
  end

  def call
    response = request_openai
    parse_response(response)
  end

  private

  def request_openai
    uri = URI(OPENAI_API_URL)
    http = Net::HTTP.new(uri.host, uri.port)
    http.use_ssl = true
    request = Net::HTTP::Post.new(uri.path, headers)
    request.body = @payload.to_json
    http.request(request)
  end

  def headers
    {
      "Content-Type" => "application/json",
      "Authorization" => "Bearer #{API_KEY}"
    }
  end

  def parse_response(response)
    JSON.parse(response.body)["choices"]&.first["message"]["content"]&.strip
  rescue
    nil
  end
end

class TableIdentifier
  def initialize(query:)
    @query = query
  end

  def call
    chat
  end

  private

  def chat
    Llm::Chat.new(payload: chat_payload).call
  end

  def chat_payload
    {
      "model": "gpt-4",
      "messages": [
        { "role": "system", "content": Prompts::TABLE_IDENTIFICATION },
        { "role": "user", "content": @query }
      ],
      "max_tokens": 100
    }
  end
end

module Prompts
  TABLE_IDENTIFICATION = "Given a user query, determine the most relevant table or tables from [users, departments, tickets]. If the query involves multiple tables (e.g., grouping users by department), return a comma-separated list of table names. Only return the table name(s) with no extra text."
end

class ReportGenerator
  require "#{Rails.root}/lib/llm"

  def initialize(query:)
    @query = query
  end

  def call
    report
  end

  private

  def report
    [
      { type: "text", data: "Here is your report" },
      { type: "table", data: ActiveRecord::Base.connection.select_all(query).to_a }
    ]
  end

  def table_structure
    ActiveRecord::Base.connection.execute("SHOW CREATE TABLE #{table_name}").first[1]
  end

  def table_name
    TableIdentifier.new(query: @query).call
  end

  def query
    Llm::Chat.new(payload: query_payload).call
  end

  def query_payload
    {
      "model": "gpt-4",
      "messages": [
        { "role": "system", "content": Prompts::SQL_GENERATION % { table_structure: table_structure } },
        { "role": "user", "content": @query }
      ],
      "max_tokens": 1000
    }
  end
end

module Prompts
  SQL_GENERATION = "Generate a MySQL query based on the table structure: %{table_structure}. Support queries involving multiple tables where applicable (e.g., grouping users by department). Only return the SQL query as plain text with no formatting, explanations, or markdown."
end

CREATE TABLE users (
  id INT PRIMARY KEY,
  name VARCHAR(255),
  email VARCHAR(255),
  status VARCHAR(50),
  department_id INT,
  created_at TIMESTAMP,
  updated_at TIMESTAMP
);

ReportGenerator.new(query: "count of inactive users").call
ReportGenerator.new(query: "list of active users").call
ReportGenerator.new(query: "number of users per department").call

SELECT COUNT(*) FROM users WHERE status = 'inactive';

SELECT d.name, COUNT(u.id)
FROM users u
JOIN departments d ON u.department_id = d.id
GROUP BY d.name;