using System;

using SqlKata;

using SqlKata.Compilers;

​

enum Materialization {

    None,

    Materialized,

    NotMaterialized

}

​

class MaterializableCteClause : AbstractFrom {

    public string Expression { get; set; }

    public Materialization Materialization { get; set; } = Materialization.None;

​

    public object[] Bindings { set; get; }

​

    /// <inheritdoc />

    public override AbstractClause Clone()

    {

      MaterializableCteClause mCte = new MaterializableCteClause();

      mCte.Engine = this.Engine;

      mCte.Alias = this.Alias;

      mCte.Expression = this.Expression;

      mCte.Bindings = this.Bindings;

      mCte.Component = this.Component;

      return (AbstractClause) mCte;

    }

}

​

class RecursiveCteClause : AbstractFrom {

    public string Expression { get; set; }

    public object[] Bindings { set; get; }

​

    /// <inheritdoc />

    public override AbstractClause Clone()

    {

      RecursiveCteClause recursiveCte = new RecursiveCteClause();

      recursiveCte.Engine = this.Engine;

      recursiveCte.Alias = this.Alias;

      recursiveCte.Expression = this.Expression;

      recursiveCte.Bindings = this.Bindings;

      recursiveCte.Component = this.Component;

      return (AbstractClause) recursiveCte;

    }

}

​

class NotMaterializedQueryFromClause : QueryFromClause {

}

class ImprovedCtePostgresCompiler : PostgresCompiler {

    public ImprovedCtePostgresCompiler() : base() {}

    private string GetMaterializationString(Materialization materialization) {

        switch (materialization) {

            case Materialization.Materialized:

                return "MATERIALIZED";

            case Materialization.NotMaterialized:

                return "NOT MATERIALIZED";

            default: return "";

        }

    }

    public override SqlResult CompileCte(AbstractFrom cte) {

        var result = new SqlResult();

        switch (cte) {

            case MaterializableCteClause mCteClause:

                result.Bindings.AddRange(mCteClause.Bindings);

                result.RawSql = this.WrapValue(mCteClause.Alias) + " AS " + GetMaterializationString(mCteClause.Materialization) + " (" + this.WrapIdentifiers(mCteClause.Expression) + ")";

                return result;

            case NotMaterializedQueryFromClause queryClause:

                Console.WriteLine("not mat query from");

                SqlResult compiledQuery = this.CompileSelectQuery(queryClause.Query);

                result.Bindings.AddRange(compiledQuery.Bindings);

                result.RawSql = this.WrapValue(queryClause.Alias) + " AS NOT MATERIALIZED (" + compiledQuery.RawSql + ")";

                return result;

            case RecursiveCteClause rCteClause:

                result.Bindings.AddRange(rCteClause.Bindings);

                result.RawSql = " RECURSIVE " + this.WrapValue(rCteClause.Alias) + " AS (" + this.WrapIdentifiers(rCteClause.Expression) + ")";

                return result;

            default:

                return base.CompileCte(cte);

        }

    }

}

​

static class QueryExtensions {

    public static Query WithMaterializationRaw(this Query query, string alias, string expression, Materialization materialization = default, object[] bindings = null) {

        var clause = new MaterializableCteClause();

        clause.Alias = alias;

        clause.Expression = expression;

        clause.Materialization = materialization;

        clause.Bindings = bindings ?? [];

        query.AddComponent("cte", clause);

        return query;

    }

    public static Query WithNotMaterialized(this Query root, string alias, Query query) {

        var clause = new NotMaterializedQueryFromClause();

        clause.Query = query;

        clause.Alias = alias;

        return root.AddComponent("cte", (AbstractClause) clause);

    }

    public static Query WithNotMaterialized(this Query query, string alias, Func<Query, Query> fn) {

        return query.WithNotMaterialized(alias, fn(new Query()));

    }

    public static Query WithRecursiveRaw(this Query query, string alias, string expression, object[] bindings = null) {

        var clause = new RecursiveCteClause();

        clause.Alias = alias;

        clause.Expression = expression;

        clause.Bindings = bindings ?? [];

        query.AddComponent("cte", clause);

        return query;

    }

}

​

public class Program

{

    public static void Main()

    {

        var compiler = new ImprovedCtePostgresCompiler();

        var query1 = new Query();

        query1

            .WithNotMaterialized("userHierarchyMembers", q => 

                                 q.From("UserHierarchyMembers as uhm")

                                 .Join("OrganizationStructureMembers as osm", "uhm.OrganizationStructureMemberId", "osm.Id")

                                .Join("OrganizationStructureLevels as osl", "osm.LevelId", "osl.Id")

                                .Where("uhm.Permission", "View"))

            .WithMaterializationRaw("amd", "select 1", Materialization.NotMaterialized)

            .WithMaterializationRaw("abc", "select 2", Materialization.Materialized);

        var compiled1 = compiler.Compile(query1);

        Console.WriteLine(compiled1.ToString());

        /*

        var query2 = new Query();

        query2.WithRecursiveRaw("amd", "select 1");

        var compiled2 = compiler.Compile(query2);

        Console.WriteLine(compiled2.ToString());

        */

    }

}