defmodule CrucibleDatasets.Loader.Rubric do
  @moduledoc """
  Loader for rubric-based evaluation datasets.

  Supports:
    - Feedback-Collection (prometheus-eval/Feedback-Collection)

  These datasets contain instructions with scoring rubrics for training
  evaluator/grader models.

  ## Examples

      # Load Feedback-Collection
      {:ok, dataset} = CrucibleDatasets.Loader.Rubric.load(:feedback_collection)

      # Load with sample size
      {:ok, dataset} = CrucibleDatasets.Loader.Rubric.load(:feedback_collection, sample_size: 500)

  """

  alias CrucibleDatasets.Dataset
  alias CrucibleDatasets.Fetcher.HuggingFace

  require Logger

  @datasets %{
    feedback_collection: %{
      repo_id: "prometheus-eval/Feedback-Collection",
      description: "Prometheus Feedback Collection for rubric-based evaluation"
    }
  }

  @doc """
  Load a rubric-based evaluation dataset.

  ## Arguments
    * `dataset_name` - Currently only `:feedback_collection`
    * `opts` - Options (see below)

  ## Options
    * `:split` - Dataset split (default: "train")
    * `:sample_size` - Limit number of items
    * `:synthetic` - Use synthetic data for testing (default: false)
    * `:token` - HuggingFace API token

  """
  @spec load(atom(), keyword()) :: {:ok, Dataset.t()} | {:error, term()}
  def load(dataset_name, opts \\ [])

  def load(dataset_name, opts) when is_atom(dataset_name) do
    case Map.get(@datasets, dataset_name) do
      nil ->
        {:error, {:unknown_dataset, dataset_name, Map.keys(@datasets)}}

      dataset_info ->
        synthetic = Keyword.get(opts, :synthetic, false)

        if synthetic do
          load_synthetic(dataset_name, opts)
        else
          load_from_huggingface(dataset_name, dataset_info, opts)
        end
    end
  end

  defp load_from_huggingface(dataset_name, %{repo_id: repo_id}, opts) do
    split = Keyword.get(opts, :split, "train") |> to_string()
    sample_size = Keyword.get(opts, :sample_size)
    token = Keyword.get(opts, :token)

    # Feedback-Collection uses a single JSON file, not standard train/test splits
    # Download the specific file directly
    file_path = "new_feedback_collection.json"

    case HuggingFace.download_file(repo_id, file_path, token: token) do
      {:ok, data} ->
        case Jason.decode(data) do
          {:ok, raw_data} when is_list(raw_data) ->
            items = parse_rubric_data(raw_data)
            items = if sample_size, do: Enum.take(items, sample_size), else: items

            dataset =
              Dataset.new(
                to_string(dataset_name),
                "1.0",
                items,
                %{
                  source: "huggingface:#{repo_id}",
                  split: split,
                  license: "apache-2.0",
                  domain: "rubric_evaluation"
                }
              )

            {:ok, dataset}

          {:ok, _} ->
            {:error, {:parse_error, :expected_array}}

          {:error, reason} ->
            {:error, {:json_parse_error, reason}}
        end

      {:error, reason} ->
        if Application.get_env(:crucible_datasets, :fallback_to_synthetic, false) do
          Logger.warning(
            "Failed to load #{dataset_name} from HuggingFace: #{inspect(reason)}, falling back to synthetic"
          )

          load_synthetic(dataset_name, opts)
        else
          {:error, {:huggingface_fetch_failed, reason}}
        end
    end
  end

  defp parse_rubric_data(raw_data) do
    raw_data
    |> Enum.with_index()
    |> Enum.map(fn {item, idx} ->
      parse_feedback_item(item, idx)
    end)
    |> Enum.reject(&is_nil/1)
  end

  defp parse_feedback_item(item, idx) do
    instruction = item["orig_instruction"] || item["instruction"] || ""
    criteria = item["orig_criteria"] || item["criteria"] || ""
    reference = item["orig_reference_answer"] || item["reference_answer"] || ""

    # Build rubric from score descriptions
    rubric = build_rubric(item)

    %{
      id: "feedback_#{idx}",
      input: %{
        instruction: instruction,
        criteria: criteria
      },
      expected: %{
        reference_answer: reference,
        rubric: rubric
      },
      metadata: %{
        source: "feedback_collection",
        has_rubric: map_size(rubric) > 0
      }
    }
  end

  defp build_rubric(item) do
    # Feedback-Collection has score1_description through score5_description
    1..5
    |> Enum.reduce(%{}, fn i, acc ->
      key = "orig_score#{i}_description"
      alt_key = "score#{i}_description"

      case item[key] || item[alt_key] do
        nil -> acc
        description -> Map.put(acc, i, description)
      end
    end)
  end

  defp load_synthetic(dataset_name, opts) do
    sample_size = Keyword.get(opts, :sample_size, 20)

    items = generate_synthetic_items(sample_size)

    dataset =
      Dataset.new(
        to_string(dataset_name),
        "1.0",
        items,
        %{
          source: "synthetic",
          license: "apache-2.0",
          domain: "rubric_evaluation"
        }
      )

    {:ok, dataset}
  end

  defp generate_synthetic_items(count) do
    examples = [
      {
        "Explain the water cycle in simple terms.",
        "Clarity and completeness of explanation",
        "The water cycle consists of evaporation, condensation, precipitation, and collection. " <>
          "Water evaporates from bodies of water, rises as vapor, condenses into clouds, " <>
          "falls as precipitation, and collects in bodies of water to repeat the cycle.",
        %{
          1 => "Response is unclear or completely incorrect",
          2 => "Response shows some understanding but has major gaps",
          3 => "Response is partially correct with minor gaps",
          4 => "Response is mostly complete and accurate",
          5 => "Response is comprehensive, clear, and fully accurate"
        }
      },
      {
        "What are the benefits of regular exercise?",
        "Comprehensiveness and accuracy of health information",
        "Regular exercise improves cardiovascular health, strengthens muscles and bones, " <>
          "helps maintain healthy weight, reduces stress and anxiety, improves sleep quality, " <>
          "and boosts overall energy levels.",
        %{
          1 => "Response contains incorrect or harmful information",
          2 => "Response mentions 1-2 benefits with limited accuracy",
          3 => "Response covers some benefits accurately",
          4 => "Response covers multiple benefits accurately",
          5 => "Response is comprehensive, accurate, and well-organized"
        }
      },
      {
        "How do you make a peanut butter and jelly sandwich?",
        "Completeness and clarity of instructions",
        "1. Gather bread, peanut butter, jelly, and a knife. " <>
          "2. Spread peanut butter on one slice of bread. " <>
          "3. Spread jelly on the other slice. " <>
          "4. Press the two slices together.",
        %{
          1 => "Instructions are missing critical steps or incorrect",
          2 => "Instructions are incomplete or confusing",
          3 => "Instructions cover main steps but lack detail",
          4 => "Instructions are clear and mostly complete",
          5 => "Instructions are clear, complete, and easy to follow"
        }
      }
    ]

    for i <- 0..(count - 1) do
      {instruction, criteria, reference, rubric} = Enum.at(examples, rem(i, length(examples)))

      %{
        id: "rubric_#{i}",
        input: %{
          instruction: instruction,
          criteria: criteria
        },
        expected: %{
          reference_answer: reference,
          rubric: rubric
        },
        metadata: %{
          source: "synthetic",
          has_rubric: true
        }
      }
    end
  end

  @doc """
  List available rubric datasets.
  """
  @spec available_datasets() :: [atom()]
  def available_datasets, do: Map.keys(@datasets)
end