import { create } from "@bufbuild/protobuf";
import * as wire from "../gen/camino/schema_pb.js";
import type { QueryArgument, QuerySelection } from "./types.js";
import type {
  QueryAggregatePredicate,
  QueryExpression,
  QueryPathStep,
  QueryPredicate,
  QueryReduction,
  QueryRelationalPlan,
  QueryRow,
} from "./relational.js";

export function relationalWire(
  argument: (v: QueryArgument) => wire.QueryArgument,
  selection: (v: QuerySelection) => wire.QuerySelection,
) {
  const path = (p: QueryPathStep): wire.QueryPathStep =>
    create(wire.QueryPathStepSchema, {
      step: p.kind === "relation" ? { case: "relation", value: p } : { case: p.kind, value: true },
    });
  const expression = (e: QueryExpression) =>
    create(wire.QueryExpressionSchema, {
      path: e.path.map(path),
      valueTypeJson: JSON.stringify(e.type),
      leaf:
        e.leaf.kind === "field"
          ? { case: "field", value: e.leaf }
          : e.leaf.kind === "ref"
            ? { case: "ref", value: e.leaf.interfaceRevisionId }
            : { case: e.leaf.kind, value: true },
    });
  const reduction = (r: QueryReduction) =>
    create(wire.QueryReductionSchema, {
      operator: {
        count: wire.QueryAggregateOperator.COUNT,
        countPresent: wire.QueryAggregateOperator.COUNT_PRESENT,
        sum: wire.QueryAggregateOperator.SUM,
        avg: wire.QueryAggregateOperator.AVG,
        min: wire.QueryAggregateOperator.MIN,
        max: wire.QueryAggregateOperator.MAX,
      }[r.operator],
      operand: r.operand && expression(r.operand),
      valueTypeJson: JSON.stringify(r.type),
    });
  const operand = (v: QueryExpression | QueryReduction) =>
    create(wire.QueryOperandSchema, {
      operand:
        "operator" in v ? { case: "reduction", value: reduction(v) } : { case: "expression", value: expression(v) },
    });
  const predicate = (p: QueryPredicate | QueryAggregatePredicate): wire.QueryPredicate => {
    switch (p.kind) {
      case "and":
      case "or":
        return create(wire.QueryPredicateSchema, {
          predicate: { case: p.kind, value: { children: p.children.map(predicate) } },
        });
      case "not":
        return create(wire.QueryPredicateSchema, { predicate: { case: "not", value: predicate(p.child) } });
      case "compare":
        return create(wire.QueryPredicateSchema, {
          predicate: {
            case: "compare",
            value: {
              operand: operand(p.expression),
              value: argument(p.value),
              operator: {
                eq: wire.QueryComparisonOperator.EQ,
                in: wire.QueryComparisonOperator.IN,
                isNull: wire.QueryComparisonOperator.IS_NULL,
                lt: wire.QueryComparisonOperator.LT,
                lte: wire.QueryComparisonOperator.LTE,
                gt: wire.QueryComparisonOperator.GT,
                gte: wire.QueryComparisonOperator.GTE,
              }[p.operator],
            },
          },
        });
      case "relation":
        return create(wire.QueryPredicateSchema, {
          predicate: {
            case: "relation",
            value: {
              path: p.path.map(path),
              relation: p,
              operator: {
                some: wire.QueryRelationPredicateOperator.SOME,
                none: wire.QueryRelationPredicateOperator.NONE,
                is: wire.QueryRelationPredicateOperator.IS,
                isNull: wire.QueryRelationPredicateOperator.IS_NULL,
              }[p.operator],
              predicate: p.predicate && predicate(p.predicate),
              value: p.value && argument(p.value),
            },
          },
        });
      case "reduce":
        return create(wire.QueryPredicateSchema, {
          predicate: {
            case: "reduce",
            value: {
              path: p.path.map(path),
              relation: p,
              where: p.where && predicate(p.where),
              having: predicate(p.having),
            },
          },
        });
    }
  };
  const row = (r: QueryRow): wire.QueryRow =>
    create(wire.QueryRowSchema, {
      row:
        r.kind === "object"
          ? {
              case: "object",
              value: {
                interfaceRevisionId: r.interfaceRevisionId,
                membership: !!r.membership,
                keyType: r.membership?.keyType,
              },
            }
          : { case: "pair", value: { source: row(r.source), target: row(r.target) } },
    });
  const plan = (p: QueryRelationalPlan) =>
    create(wire.QueryRelationalPlanSchema, {
      stages: p.stages.map((s) => {
        const op = s.operation;
        const operation: wire.QueryRelationalStage["operation"] =
          op.kind === "source"
            ? { case: "source", value: create(wire.QueryRelationPathSchema, op) }
            : op.kind === "filter"
              ? { case: "filter", value: predicate(op.predicate) }
              : op.kind === "expand"
                ? {
                    case: "expand",
                    value: create(wire.QueryExpansionSchema, { path: op.path.map(path), relation: op }),
                  }
                : {
                    case: "distinct",
                    value: create(wire.QueryDistinctSchema, {
                      keys: op.keys.map((k) => ({ ...k, expression: expression(k.expression) })),
                    }),
                  };
        return create(wire.QueryRelationalStageSchema, { id: s.id, input: s.input, row: row(s.row), operation });
      }),
      terminals: p.terminals.map((t) => ({
        stage: t.stage,
        path: t.path,
        groups: t.kind === "groups",
        keys: t.keys.map((k) => ({ ...k, expression: expression(k.expression) })),
        reductions: t.reductions.map((r) => ({ ...r, reduction: reduction(r.reduction) })),
        where: t.where && predicate(t.where),
        having: t.having && predicate(t.having),
        order: t.order.map((o) => ({ operand: operand(o.expression), descending: o.descending })),
        first: t.first && argument(t.first),
        all: t.all && argument(t.all),
        after: t.after && argument(t.after),
        selection: t.selection.map(selection),
        residual: t.residual,
      })),
    });
  return { predicate, plan };
}