using Microsoft.AnalysisServices.AdomdClient;
using Microsoft.AnalysisServices.Tabular;
using System.Diagnostics;

using (var server = new Server())
{
    /// Connect to a dataset in a workspace.
    /// 
    server.Connect($"Data Source=\"{TestModel.WorkspaceUrl}\";"
        + $"Initial Catalog={TestModel.DatabaseName};");

    Database database = server.Databases.FindByName(TestModel.DatabaseName);
    if (database == null)
    {
        throw new ApplicationException(
            $"Database '{TestModel.DatabaseName}' cannot be found!");
    }

    /// Make sure the dataset has the minimum compat level 
    /// for auto aggs
    /// 
    if (database.CompatibilityLevel < 1564)
    {
        database.CompatibilityLevel = 1564;
        database.Update();
    }

    var model = database.Model;

    /// Check if there are any auto aggs tables, 
    /// i.e. system-managed tables with an inferred partition source.
    if (model.Tables.Where(t => t.SystemManaged
        && t.Partitions.All(x => x.SourceType == PartitionSourceType.Inferred)).Any())
    {
        Console.WriteLine("This dataset has auto aggs tables!");
    }
    else
    {
        Console.WriteLine("This dataset does not have any auto aggs tables!");
    }

    /// Display the current automatic aggregations configuration.
    /// 
    var aggsConfig = model.AutomaticAggregationOptions;
    if (aggsConfig.IsEmpty)
    {
        Console.WriteLine("This dataset has no auto aggs settings.");
    }
    else
    {
        Console.WriteLine(
            $"Percentage of queries to use aggs: {aggsConfig.QueryCoverage}\n" +
            $"Min size of detail tables in rows: {aggsConfig.DetailTableMinRows}\n" +
            $"Max size of aggs tables in rows:   {aggsConfig.AggregationTableMaxRows}\n" +
            $"Rel. size limit of any aggs table: {aggsConfig.AggregationTableSizeLimit}"
            );
    }

    /// Add a trace to capture the ProgressReportEnd event.
    /// 
    const string traceName = "AutoAggsTrainingTrace";
    var trace = server.Traces.OfType<Microsoft.AnalysisServices.Tabular.Trace>()
        .Where(t => t.Name.Equals(traceName)).FirstOrDefault();
    if(trace != null)
    {
        if (trace.IsStarted) trace.Stop();
        trace.Drop();
    }

    trace = server.Traces.Add(traceName);
    trace.AutoRestart = true;

    TraceEvent traceEvent = trace.Events.Add(
        Microsoft.AnalysisServices.TraceEventClass.ProgressReportEnd);

    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.EventSubclass);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.EventClass);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.TextData);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.Duration);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.Error);

    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.StartTime);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.EndTime);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.DatabaseName);

    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.ActivityID);
    traceEvent.Columns.Add(Microsoft.AnalysisServices.TraceColumn.RequestID);

    trace.OnEvent += (object sender, TraceEventArgs e) =>
    {
        var defaultForeground = Console.ForegroundColor;

        if (e.EventClass == Microsoft.AnalysisServices.TraceEventClass.ProgressReportEnd &&
            e.EventSubclass == Microsoft.AnalysisServices.TraceEventSubclass.AutoAggsTraining)
        {

            var trainingResults = System.Text.Json.Nodes.JsonNode.Parse(e.TextData)!;
            if (trainingResults["statistics"]!["aggregationTableCount"]!.GetValue<int>() == 0)
            {
                Console.ForegroundColor = ConsoleColor.Red;
                Console.Write("No auto aggs tables generated! ");
            }
            else
            {
                Console.ForegroundColor = ConsoleColor.Green;
                Console.Write($"Success! {trainingResults["statistics"]!["aggregationTableCount"]!.GetValue<int>()} auto aggs table(s) generated! ");
            }

            Console.WriteLine($"The number of eligible query shapes was {trainingResults["statistics"]!["queryShapes"]!["eligible"]!.GetValue<int>()}.\n");
            Console.ForegroundColor = defaultForeground;
        }

        Console.WriteLine($"EventClass: {e.EventClass}");
        Console.WriteLine($"EventSubclass: {e.EventSubclass}");

        Console.ForegroundColor = ConsoleColor.Yellow;
        Console.WriteLine($"TextData:{e.TextData}");
        Console.ForegroundColor = defaultForeground;

        Console.WriteLine($"Success: {string.IsNullOrEmpty(e.Error) || e.Error == "0"}");
        Console.WriteLine($"StartTime: {e.StartTime}");
        Console.WriteLine($"EndTime: {e.EndTime}ms");
        Console.WriteLine($"Duration: {e.Duration}ms");
        Console.WriteLine();
        Console.WriteLine($"Database: {e.DatabaseName}");
        Console.WriteLine(
            $"ActivityID: {e[Microsoft.AnalysisServices.TraceColumn.ActivityID]}");
        Console.WriteLine(
            $"RequestID: {e[Microsoft.AnalysisServices.TraceColumn.RequestID]}");
        Console.WriteLine("--------------------");

        if (e.EventClass == Microsoft.AnalysisServices.TraceEventClass.ProgressReportEnd &&
        e.EventSubclass == Microsoft.AnalysisServices.TraceEventSubclass.AutoAggsTraining)
        {
            Console.WriteLine("Auto aggs training done. Press [Enter] to continue!");
        }
    };


    trace.Update(Microsoft.AnalysisServices.UpdateOptions.Default, 
        Microsoft.AnalysisServices.UpdateMode.CreateOrReplace);

    trace.Start();

    /// Perform auto aggs training using a one-off configuration.
    /// 
    model.ApplyAutomaticAggregations(
        new AutomaticAggregationOptions
        {
            QueryCoverage = 0.5
        });

    /// Wait for the user to press the Enter key
    /// before continuing.
    /// 
    Console.ReadLine();

    /// Stop and remove the trace
    /// 
    trace.Stop();
    trace.Drop();

    /// Persist any changes.
    /// 
    database.Update(Microsoft.AnalysisServices.UpdateOptions.ExpandFull);

    /// Refresh the model to import data into the aggs tables.
    /// 
    model.RequestRefresh(RefreshType.Full);
    model.SaveChanges();

    /// Count the rows in every aggs table.
    /// 
    using (var adomdConnection = new AdomdConnection(server.ConnectionString))
    {
        adomdConnection.Open();

        foreach (var table in model.Tables.Where(t => t.SystemManaged &&
            t.Partitions.All(x => x.SourceType == PartitionSourceType.Inferred)))
        {
            string query = "DEFINE\n" +
                $"   var tblRows = COUNTROWS('{table.Name}')\n" +
                "   var rowCount = SUMMARIZECOLUMNS(\"Row Count\", IF(tblRows > 0, tblRows, 0))\n\n" +
                "Evaluate rowCount";

            using (var cmd = new AdomdCommand(query, adomdConnection))
            using (var reader = cmd.ExecuteReader())
            {
                while (reader.Read())
                {
                    var rowCount = reader.GetInt64(0);

                    var defaultForeground = Console.ForegroundColor;
                    Console.ForegroundColor = rowCount > 0 ? ConsoleColor.Green : ConsoleColor.Red;
                    Console.WriteLine($"Auto aggs table {table.Name} has {rowCount} rows.");
                    Console.ForegroundColor = defaultForeground;
                    break;
                }
            }
        }
    }
}

internal class TestModel
{
    public const string WorkspaceUrl = "powerbi://api.powerbi.com/v1.0/myorg/AdventureWorksTest";
    public const string DatabaseName = "AdventureWorksDQ";
}