-- Engine-independent query rendering. Each engine adapter defines
--   bench_pred(mode, terms) -> boolean SQL over docs (NULL = mode unsupported as a filter)
--   bench_rank(mode, terms) -> ORDER BY expression incl. direction (NULL = no top-k)
--   bench_source(mode, terms) -> FROM item of ranked statements, aliased docs, when the engine needs another shape
--                                than `docs WHERE bench_pred` (NULL = that one)
-- Modes: and | or | phrase | not (all terms but the last, minus the last term) | prefix (term*) |
-- fuzzy (edit distance 1). Unsupported modes must render NULL in both functions.

CREATE OR REPLACE FUNCTION tsq(mode text, terms text[]) RETURNS text LANGUAGE sql IMMUTABLE AS $$
  SELECT CASE mode
    WHEN 'and'    THEN string_agg(quote_literal(t), ' & ')
    WHEN 'or'     THEN string_agg(quote_literal(t), ' | ')
    WHEN 'phrase' THEN string_agg(quote_literal(t), ' <-> ')
    WHEN 'not'    THEN string_agg(CASE WHEN o = cardinality(terms) THEN '!' ELSE '' END || quote_literal(t), ' & ')
    WHEN 'prefix' THEN string_agg(quote_literal(t) || ':*', ' & ')
  END
  FROM unnest(terms) WITH ORDINALITY AS u(t, o)
$$;

CREATE OR REPLACE FUNCTION bench_stmt(kind text, mode text, terms text[]) RETURNS text LANGUAGE plpgsql STABLE AS $$
DECLARE
  pred text := bench_pred(mode, terms);
  rank text := bench_rank(mode, terms);
  ranked boolean := kind IN ('top10', 'filter', 'join');
  extra text := CASE kind
    WHEN 'filter' THEN format('docs.created >= %L', (SELECT v FROM params WHERE k = 'recent_cutoff'))
    WHEN 'join'   THEN 'categories.region = 3'
  END;
  src text := CASE WHEN ranked THEN bench_source(mode, terms) END;
  conds text := concat_ws(' AND ', CASE WHEN src IS NULL THEN pred END, extra);
BEGIN
  IF (ranked AND rank IS NULL) OR (NOT ranked AND pred IS NULL) THEN
    RETURN NULL;
  END IF;
  RETURN CASE kind
    WHEN 'count' THEN 'SELECT count(*) FROM docs WHERE ' || pred
    WHEN 'match' THEN 'SELECT docs.id FROM docs WHERE ' || pred
    ELSE 'SELECT docs.id FROM ' || coalesce(src, 'docs')
      || CASE kind WHEN 'join' THEN ' JOIN categories ON categories.id = docs.category_id' ELSE '' END
      || CASE WHEN conds <> '' THEN ' WHERE ' || conds ELSE '' END
      || ' ORDER BY ' || rank || ' LIMIT 10'
  END;
END $$;

-- The classes (workloads) of the articles dataset: the statement kind and the query modes. `trace` classes take the sampled
-- substrings of `queries` (each as and / or / phrase); `extra` classes (term, prefix, fuzzy) are derived from the first
-- substrings, see extra_queries. A class runs only if every mode renders.
CREATE TABLE IF NOT EXISTS workloads (wl text PRIMARY KEY, kind text, modes text[], src text);
INSERT INTO workloads VALUES
  ('and', 'top10', '{and}', 'trace'), ('or', 'top10', '{or}', 'trace'), ('phrase', 'top10', '{phrase}', 'trace'),
  ('mixed', 'top10', '{and,or,phrase}', 'trace'), ('count', 'count', '{and,or,phrase}', 'trace'),
  ('filter', 'filter', '{and,or,phrase}', 'trace'), ('join', 'join', '{and,or,phrase}', 'trace'),
  ('term', 'top10', '{or}', 'extra'), ('prefix', 'top10', '{prefix}', 'extra'), ('fuzzy', 'top10', '{fuzzy}', 'extra')
ON CONFLICT DO NOTHING;

-- 36 derived queries from the first six 2-4 term substrings: term = their longest word as a one-term OR, prefix = the first
-- length-3 chars (min 3) of the longest word + *, fuzzy = it minus its second-to-last char (edit distance 1); prefix and fuzzy take
-- the first six substrings whose longest word is lowercase ASCII (a CJK run is tokenized differently by every engine and matches nothing).
CREATE OR REPLACE VIEW extra_queries AS
WITH subs AS (
  SELECT sub, terms FROM queries WHERE mode = 'and' AND cardinality(terms) BETWEEN 2 AND 4 ORDER BY sub LIMIT 6
), longest AS (
  SELECT sub, terms, (SELECT t FROM unnest(terms) t ORDER BY length(t) DESC, t LIMIT 1) AS w
  FROM queries WHERE mode = 'and' AND cardinality(terms) BETWEEN 2 AND 4
), words AS (
  SELECT sub, w FROM longest WHERE sub IN (SELECT sub FROM subs)
), pwords AS (
  SELECT sub, w FROM longest WHERE w ~ '^[a-z]+$' ORDER BY sub LIMIT 6
)
SELECT 'term' AS wl, sub AS qid, 'or' AS mode, ARRAY[w] AS terms FROM words
UNION ALL SELECT 'prefix', sub, 'prefix', ARRAY[left(w, greatest(3, length(w) - 3))] FROM pwords
UNION ALL SELECT 'fuzzy', sub, 'fuzzy', ARRAY[overlay(w PLACING '' FROM length(w) - 1 FOR 1)] FROM pwords;

CREATE OR REPLACE VIEW class_queries AS
SELECT w.wl, q.id AS qid, q.mode, q.terms FROM workloads w JOIN queries q ON q.mode = ANY (w.modes) WHERE w.src = 'trace'
UNION ALL SELECT e.wl, e.qid, e.mode, e.terms FROM extra_queries e;

CREATE TABLE IF NOT EXISTS stmts (wl text, n int, qid int, stmt text, PRIMARY KEY (wl, n));

CREATE OR REPLACE FUNCTION bench_build_stmts() RETURNS SETOF text LANGUAGE plpgsql AS $$
DECLARE w workloads;
BEGIN
  TRUNCATE stmts;
  FOR w IN SELECT * FROM workloads LOOP
    CONTINUE WHEN EXISTS (
      SELECT FROM unnest(w.modes) m WHERE bench_stmt(w.kind, m, '{x,y}') IS NULL);
    INSERT INTO stmts
    SELECT w.wl, row_number() OVER (ORDER BY hashtext(c.qid::text || c.mode), c.qid, c.mode), c.qid, bench_stmt(w.kind, c.mode, c.terms)
    FROM class_queries c WHERE c.wl = w.wl;
    RETURN NEXT w.wl;
  END LOOP;
END $$;

-- Runs a statement returning ids and gives them back as a sorted array.
CREATE OR REPLACE FUNCTION bench_ids(stmt text) RETURNS bigint[] LANGUAGE plpgsql AS $$
DECLARE r bigint[];
BEGIN
  EXECUTE format('SELECT coalesce(array_agg(id ORDER BY id), ''{}'') FROM (%s) s', stmt) INTO r;
  RETURN r;
END $$;