DROP PROCEDURE IF EXISTS fix_files_duplicates;
PRINT 'CREATE PROCEDURE fix_files_duplicates...';
GO
CREATE PROCEDURE fix_files_duplicates(@new_collation VARCHAR(100)) AS
BEGIN
  DECLARE @statement VARCHAR(MAX);

  DROP TABLE IF EXISTS files2;
  CREATE TABLE files2(
    fileId BIGINT NOT NULL IDENTITY(10000000001,1),
    fileName NVARCHAR(260) NOT NULL,
    f_last_used DATE NOT NULL);
  SET @statement = CONCAT('ALTER TABLE files2 ALTER COLUMN fileName NVARCHAR(260) COLLATE ', @new_collation, ' NOT NULL');
  EXEC(@statement);
  ALTER TABLE files2 ADD CONSTRAINT PK_files2 PRIMARY KEY NONCLUSTERED (fileName) WITH (IGNORE_DUP_KEY = ON);
  CREATE CLUSTERED index idx_files_fileId ON files2(fileId);
  SET IDENTITY_INSERT files2 ON;
  INSERT INTO files2(fileId, fileName, f_last_used)
    SELECT fileId, fileName, f_last_used FROM files;
  SET IDENTITY_INSERT files2 OFF;

  DROP TABLE IF EXISTS #problematic_files;
  CREATE TABLE #problematic_files(fileId BIGINT, fileName NVARCHAR(260), idEq BIGINT DEFAULT 0);
  SET @statement = CONCAT('ALTER TABLE #problematic_files ALTER COLUMN fileName NVARCHAR(260) COLLATE ', @new_collation);
  EXEC(@statement);
  SET @statement = CONCAT('INSERT INTO #problematic_files(fileId, fileName) SELECT fileId, fileName COLLATE ', @new_collation,
    ' FROM files WHERE NOT EXISTS(SELECT 1 FROM files2 WHERE files2.fileId = files.fileId)');
  EXEC(@statement);

  UPDATE #problematic_files SET #problematic_files.idEq = files2.fileId
    FROM #problematic_files
    INNER JOIN files2 ON files2.fileName = #problematic_files.fileName;
  DELETE FROM #problematic_files WHERE idEq = 0;

  DROP TABLE files;
  EXEC sp_rename 'files2', 'files';
  EXEC sp_rename 'PK_files2', 'PK_files', 'OBJECT';

  DECLARE c_files CURSOR FOR SELECT fileId, idEq FROM #problematic_files;
  DECLARE @fileId BIGINT, @idEq BIGINT;
  OPEN c_files;
  FETCH NEXT FROM c_files INTO @fileId, @idEq;
  WHILE @@FETCH_STATUS = 0
  BEGIN
    UPDATE events SET fileId = @idEq WHERE fileId = @fileId;
    FETCH NEXT FROM c_files INTO @fileId, @idEq;
  END;
  CLOSE c_files;
  DEALLOCATE c_files;
  DROP TABLE IF EXISTS #problematic_files;
END;
GO

DROP PROCEDURE IF EXISTS fix_directories_duplicates;
PRINT 'CREATE PROCEDURE fix_directories_duplicates...';
GO
CREATE PROCEDURE fix_directories_duplicates(@new_collation VARCHAR(100)) AS
BEGIN
  DECLARE @statement VARCHAR(MAX);

  DROP TABLE IF EXISTS directories2;
  CREATE TABLE directories2(
    dirId BIGINT NOT NULL IDENTITY(10000000001,1),
    dirName NVARCHAR(330) NOT NULL,
    d_last_used DATE NOT NULL);
  SET @statement = CONCAT('ALTER TABLE directories2 ALTER COLUMN dirName NVARCHAR(330) COLLATE ', @new_collation, ' NOT NULL');
  EXEC(@statement);
  ALTER TABLE directories2 ADD CONSTRAINT PK_directories2 PRIMARY KEY NONCLUSTERED (dirName) WITH (IGNORE_DUP_KEY = ON);
  CREATE CLUSTERED index idx_directories_dirId ON directories2(dirId);
  SET IDENTITY_INSERT directories2 ON;
  INSERT INTO directories2(dirId, dirName, d_last_used)
    SELECT dirId, dirName, d_last_used FROM directories;
  SET IDENTITY_INSERT directories2 OFF;

  DROP TABLE IF EXISTS #problematic_directories;
  CREATE TABLE #problematic_directories(dirId BIGINT, dirName NVARCHAR(330), idEq BIGINT DEFAULT 0);
  SET @statement = CONCAT('ALTER TABLE #problematic_directories ALTER COLUMN dirName NVARCHAR(330) COLLATE ', @new_collation);
  EXEC(@statement);
  SET @statement = CONCAT('INSERT INTO #problematic_directories(dirId, dirName) SELECT dirId, dirName COLLATE ', @new_collation,
    ' FROM directories WHERE NOT EXISTS(SELECT 1 FROM directories2 WHERE directories2.dirId = directories.dirId)');
  EXEC(@statement);

  UPDATE #problematic_directories SET #problematic_directories.idEq = directories2.dirId
    FROM #problematic_directories
    INNER JOIN directories2 ON directories2.dirName = #problematic_directories.dirName;
  DELETE FROM #problematic_directories WHERE idEq = 0;

  DROP TABLE directories;
  EXEC sp_rename 'directories2', 'directories';
  EXEC sp_rename 'PK_directories2', 'PK_directories', 'OBJECT';

  DECLARE c_directories CURSOR FOR SELECT dirId, idEq FROM #problematic_directories;
  DECLARE @dirId BIGINT, @idEq BIGINT;
  OPEN c_directories;
  FETCH NEXT FROM c_directories INTO @dirId, @idEq;
  WHILE @@FETCH_STATUS = 0
  BEGIN
    UPDATE events SET dirId = @idEq WHERE dirId = @dirId;
    FETCH NEXT FROM c_directories INTO @dirId, @idEq;
  END;
  CLOSE c_directories;
  DEALLOCATE c_directories;
  DROP TABLE IF EXISTS #problematic_directories;
END;
GO

DROP PROCEDURE IF EXISTS fix_modules_file_duplicates;
PRINT 'CREATE PROCEDURE fix_modules_file_duplicates...';
GO
CREATE PROCEDURE fix_modules_file_duplicates(@new_collation VARCHAR(100)) AS
BEGIN
  DECLARE @statement VARCHAR(MAX);

  DROP TABLE IF EXISTS modules_file2;
  CREATE TABLE modules_file2(
    mfId INT NOT NULL IDENTITY(1,1),
    mfName NVARCHAR(260) NOT NULL,
    CONSTRAINT PK_modules_file2 PRIMARY KEY CLUSTERED (mfId));
  SET @statement = CONCAT('ALTER TABLE modules_file2 ALTER COLUMN mfName NVARCHAR(260) COLLATE ', @new_collation, ' NOT NULL');
  EXEC(@statement);
  CREATE UNIQUE NONCLUSTERED INDEX idx_unique_modulesFile ON modules_file2(mfName) WITH (IGNORE_DUP_KEY = ON);
  SET IDENTITY_INSERT modules_file2 ON;
  INSERT INTO modules_file2(mfId, mfName)
    SELECT mfId, mfName FROM modules_file;
  SET IDENTITY_INSERT modules_file2 OFF;

  DROP TABLE IF EXISTS #problematic_modules_files;
  CREATE TABLE #problematic_modules_files(mfId INT, mfName NVARCHAR(260), idEq INT DEFAULT 0);
  SET @statement = CONCAT('ALTER TABLE #problematic_modules_files ALTER COLUMN mfName NVARCHAR(260) COLLATE ', @new_collation);
  EXEC(@statement);
  SET @statement = CONCAT('INSERT INTO #problematic_modules_files(mfId, mfName) SELECT mfId, mfName COLLATE ', @new_collation,
    ' FROM modules_file WHERE NOT EXISTS(SELECT 1 FROM modules_file2 WHERE modules_file2.mfId = modules_file.mfId)');
  EXEC(@statement);

  UPDATE #problematic_modules_files SET #problematic_modules_files.idEq = modules_file2.mfId
    FROM #problematic_modules_files
    INNER JOIN modules_file2 ON modules_file2.mfName = #problematic_modules_files.mfName;
  DELETE FROM #problematic_modules_files WHERE idEq = 0;

  DROP TABLE modules_file;
  EXEC sp_rename 'modules_file2', 'modules_file';
  EXEC sp_rename 'PK_modules_file2', 'PK_modules_file', 'OBJECT';

  DECLARE c_module_files CURSOR FOR SELECT mfId, idEq FROM #problematic_modules_files;
  DECLARE @mfId INT, @idEq INT;
  OPEN c_module_files;
    FETCH NEXT FROM c_module_files INTO @mfId, @idEq;
    WHILE @@FETCH_STATUS = 0
    BEGIN
      UPDATE modules SET mainFileId = @idEq WHERE mainFileId = @mfId;
      BEGIN TRY -- it may violite unique key in module_paths, but it's veeery unlikely, so let's ignore it
        UPDATE module_paths SET fileId = @idEq WHERE fileId = @mfId;
      END TRY
      BEGIN CATCH
      END CATCH;
      UPDATE processes SET fileId = @idEq WHERE fileId = @mfId;
      FETCH NEXT FROM c_module_files INTO @mfId, @idEq;
    END;
  CLOSE c_module_files;
  DEALLOCATE c_module_files;
  DROP TABLE IF EXISTS #problematic_modules_files;
END;
GO

DROP PROCEDURE IF EXISTS fix_modules_directory_duplicates;
PRINT 'CREATE PROCEDURE fix_modules_directory_duplicates...';
GO
CREATE PROCEDURE fix_modules_directory_duplicates(@new_collation VARCHAR(100)) AS
BEGIN
  DECLARE @statement VARCHAR(MAX);

  DROP TABLE IF EXISTS modules_directory2;
  CREATE TABLE modules_directory2(
    mdId INT NOT NULL IDENTITY(1,1),
    mdName NVARCHAR(330) NOT NULL,
    CONSTRAINT PK_modules_directory2 PRIMARY KEY CLUSTERED (mdId));
  SET @statement = CONCAT('ALTER TABLE modules_directory2 ALTER COLUMN mdName NVARCHAR(330) COLLATE ', @new_collation, ' NOT NULL');
  EXEC(@statement);
  CREATE UNIQUE NONCLUSTERED INDEX idx_unique_modulesDirectory ON modules_directory2(mdName) WITH (IGNORE_DUP_KEY = ON);
  SET IDENTITY_INSERT modules_directory2 ON;
  INSERT INTO modules_directory2(mdId, mdName)
    SELECT mdId, mdName FROM modules_directory;
  SET IDENTITY_INSERT modules_directory2 OFF;

  DROP TABLE IF EXISTS #problematic_modules_directories;
  CREATE TABLE #problematic_modules_directories(mdId INT, mdName NVARCHAR(330), idEq INT DEFAULT 0);
  SET @statement = CONCAT('ALTER TABLE #problematic_modules_directories ALTER COLUMN mdName NVARCHAR(330) COLLATE ', @new_collation);
  EXEC(@statement);
  SET @statement = CONCAT('INSERT INTO #problematic_modules_directories(mdId, mdName) SELECT mdId, mdName COLLATE ', @new_collation,
    ' FROM modules_directory WHERE NOT EXISTS(SELECT 1 FROM modules_directory2 WHERE modules_directory2.mdId = modules_directory.mdId)');
  EXEC(@statement);

  UPDATE #problematic_modules_directories SET #problematic_modules_directories.idEq = modules_directory2.mdId
    FROM #problematic_modules_directories
    INNER JOIN modules_directory2 ON modules_directory2.mdName = #problematic_modules_directories.mdName;
  DELETE FROM #problematic_modules_directories WHERE idEq = 0;

  DROP TABLE modules_directory;
  EXEC sp_rename 'modules_directory2', 'modules_directory';
  EXEC sp_rename 'PK_modules_directory2', 'PK_modules_directory', 'OBJECT';

  DECLARE c_module_directories CURSOR FOR SELECT mdId, idEq FROM #problematic_modules_directories;
  DECLARE @mdId INT, @idEq INT;
  OPEN c_module_directories;
    FETCH NEXT FROM c_module_directories INTO @mdId, @idEq;
    WHILE @@FETCH_STATUS = 0
    BEGIN
      BEGIN TRY -- it may violite unique key in module_paths, but it's veeery unlikely, so let's ignore it
        UPDATE module_paths SET dirId = @idEq WHERE dirId = @mdId;
      END TRY
      BEGIN CATCH
      END CATCH;
      UPDATE processes SET dirId = @idEq WHERE dirId = @mdId;
      FETCH NEXT FROM c_module_directories INTO @mdId, @idEq;
    END;
  CLOSE c_module_directories;
  DEALLOCATE c_module_directories;
  DROP TABLE IF EXISTS #problematic_modules_directories;
END;
GO

DROP PROCEDURE IF EXISTS full_unicode_support;
PRINT 'CREATE PROCEDURE full_unicode_support...';
GO
CREATE PROCEDURE full_unicode_support AS
BEGIN
  DECLARE @statement VARCHAR(1000);
  DECLARE @progress VARCHAR(1000);

  -- update amsi tables (now they are joined by id)
  EXEC procAddColumn 'events_amsi', 'fragment_id', 'INT';
  IF dbo.fnIsColumnThere('events_amsi', 'fragmentHash') = 1
    EXEC('UPDATE events_amsi SET fragment_id = amsi_scripts_fragments.id
      FROM events_amsi JOIN amsi_scripts_fragments ON amsi_scripts_fragments.sha1 = events_amsi.fragmentHash');
  EXEC procDropIndex 'events_amsi', 'idx_events_amsi_fragmentHash';
  EXEC procAddIndex 'events_amsi', 'fragment_id', 'idx_events_amsi_fragment_id';
  EXEC procDropColumn 'events_amsi', 'fragmentHash';
  EXEC procAddIndex 'amsi_scripts_fragments', 'sha1', 'idx_amsi_scripts_fragments_sha1';
  EXEC procDropPrimaryKey 'amsi_scripts_fragments';
  EXEC procAddPrimaryKey 'amsi_scripts_fragments', 'id', 'PK_amsi_scripts_fragments';
  EXEC procDropIndex 'amsi_scripts_fragments', 'idx_amsi_scripts_fragments_id';

  -- remove duplicated index
  SELECT @statement = 'ALTER TABLE tags DROP CONSTRAINT ' + name FROM sys.indexes WHERE object_id = OBJECT_ID('tags') AND is_unique_constraint = 1;
  EXEC(@statement);

  DECLARE @current_collation VARCHAR(100) = (SELECT collation_name FROM sys.databases WHERE name = DB_NAME());
  DECLARE @new_collation VARCHAR(100) = (SELECT TOP(1) name FROM sys.fn_helpcollations()
    WHERE COLLATIONPROPERTY(name, 'LCID') = COLLATIONPROPERTY(@current_collation, 'LCID') -- same collation family
    AND COLLATIONPROPERTY(name, 'ComparisonStyle') = COLLATIONPROPERTY(@current_collation, 'ComparisonStyle') -- same comparision rules
    AND COLLATIONPROPERTY(name, 'CodePage') > 0
    AND (COLLATIONPROPERTY(name, 'Version') >= 3 -- collations with _140_ in name - they supports supplementary characters by default
      OR CHARINDEX('_SC', name) > 0) -- support for supplementary characters
    AND CHARINDEX('_BIN', name) = 0
    AND CHARINDEX('SQL_', name) = 0 -- unadvised collations
    ORDER BY COLLATIONPROPERTY(name, 'Version') DESC);
  IF @new_collation IS NOT NULL -- there may not be a proper collation for some regions
  BEGIN
    DROP TABLE IF EXISTS #columns_to_update;
    DROP TABLE IF EXISTS #indexes_to_update;
    DROP TABLE IF EXISTS #statements;

    SELECT objects.object_id, columns.column_id, objects.name AS t_name, columns.name AS c_name, columns.max_length, columns.is_nullable, types.name AS type
    INTO #columns_to_update
    FROM sys.columns
    JOIN sys.objects ON objects.object_id = columns.object_id
    JOIN sys.types ON types.user_type_id = columns.user_type_id
    WHERE objects.schema_id = SCHEMA_ID() AND columns.collation_name <> @new_collation AND (types.name = 'nvarchar' OR types.name = 'varchar');

    -- indexes that contains selected columns (collation of column cannot be changed if that column is contained in some index)
    SELECT DISTINCT index_columns.object_id, index_id
    INTO #indexes_to_update
    FROM #columns_to_update
    INNER JOIN sys.index_columns ON index_columns.object_id = #columns_to_update.object_id AND index_columns.column_id = #columns_to_update.column_id;

    -- drop primary keys and unique constraints
    SELECT
      'ALTER TABLE ' + objects.name +
      ' DROP CONSTRAINT ' + indexes.name AS statement, 'Removing constraints in ' + objects.name AS progress
    INTO #statements
    FROM #indexes_to_update
    JOIN sys.indexes ON indexes.object_id = #indexes_to_update.object_id AND indexes.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.objects ON objects.object_id = #indexes_to_update.object_id
    WHERE is_primary_key = 1 OR is_unique_constraint = 1;

    -- drop other indexes
    INSERT INTO #statements
    SELECT
      'DROP INDEX ' + indexes.name +
      ' ON ' + objects.name, 'Dropping indexes in ' + objects.name
    FROM #indexes_to_update
    JOIN sys.indexes ON indexes.object_id = #indexes_to_update.object_id AND indexes.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.objects ON objects.object_id = #indexes_to_update.object_id
    WHERE is_primary_key = 0 AND is_unique_constraint = 0;

    -- modify column collation
    INSERT INTO #statements
    SELECT
      'ALTER TABLE ' + #columns_to_update.t_name +
      ' ALTER COLUMN ' + #columns_to_update.c_name +
      ' ' + type + '(' + (CASE
        WHEN max_length = -1 THEN 'MAX'
        WHEN varchar_lengths.c_length IS NOT NULL THEN CONVERT(VARCHAR(10), c_length) -- from SqlColumnsLength.h
        WHEN type = 'nvarchar' THEN CONVERT(VARCHAR(10), max_length / 2)
        ELSE CONVERT(VARCHAR(10), max_length)
      END) + ') COLLATE ' + @new_collation +
      (CASE is_nullable WHEN 0 THEN ' NOT ' ELSE ' ' END) + 'NULL', 'Changing collation in ' + #columns_to_update.t_name
    FROM #columns_to_update
    LEFT JOIN varchar_lengths ON
      varchar_lengths.t_name COLLATE DATABASE_DEFAULT = #columns_to_update.t_name
      AND varchar_lengths.c_name COLLATE DATABASE_DEFAULT = #columns_to_update.c_name;

    -- recreate primary keys
    INSERT INTO #statements
    SELECT
      'ALTER TABLE ' + objects.name +
      ' ADD CONSTRAINT ' +
      (CASE CHARINDEX('__', indexes.name) WHEN 0 THEN indexes.name ELSE 'PK_' + objects.name END) + -- change default name for more readable
      ' PRIMARY KEY ' + (CASE indexes.type WHEN 2 THEN 'NON' ELSE '' END) + 'CLUSTERED ' +
      ' (' + STRING_AGG(columns.name + (CASE is_descending_key WHEN 0 THEN ' ASC' ELSE ' DESC' END), ', ') WITHIN GROUP (ORDER BY index_column_id ASC) + ')' +
      (CASE ignore_dup_key WHEN 1 THEN ' WITH (IGNORE_DUP_KEY = ON)' ELSE '' END) AS statement, 'Recreating primary key in ' + objects.name
    FROM #indexes_to_update
    JOIN sys.indexes ON indexes.object_id = #indexes_to_update.object_id AND indexes.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.index_columns ON index_columns.object_id = #indexes_to_update.object_id AND index_columns.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.objects ON objects.object_id = #indexes_to_update.object_id
    LEFT JOIN sys.columns ON columns.object_id = #indexes_to_update.object_id AND columns.column_id = index_columns.column_id
    WHERE is_primary_key = 1
    GROUP BY indexes.object_id, indexes.index_id, indexes.name, is_unique, indexes.type, objects.name, ignore_dup_key;

    -- recreate other indexes
    INSERT INTO #statements
    SELECT
      'CREATE ' +
      (CASE is_unique WHEN 1 THEN 'UNIQUE ' ELSE '' END) +
      (CASE indexes.type WHEN 2 THEN 'NON' ELSE '' END) +
      'CLUSTERED INDEX ' +
      (CASE CHARINDEX('__', indexes.name)
        WHEN 0 THEN indexes.name
        ELSE 'idx_' +
          (CASE is_unique WHEN 1 THEN 'unique_' ELSE '' END) +
          objects.name + '_' +
          STRING_AGG(columns.name, '_') WITHIN GROUP (ORDER BY index_column_id ASC)
      END) +
      ' ON ' + objects.name +
      ' (' + STRING_AGG(columns.name + (CASE is_descending_key WHEN 0 THEN ' ASC' ELSE ' DESC' END), ', ') WITHIN GROUP (ORDER BY index_column_id ASC) + ')' +
      (CASE ignore_dup_key WHEN 1 THEN ' WITH (IGNORE_DUP_KEY = ON)' ELSE '' END), 'Recreating indexes in ' + objects.name
    FROM #indexes_to_update
    JOIN sys.indexes ON indexes.object_id = #indexes_to_update.object_id AND indexes.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.index_columns ON index_columns.object_id = #indexes_to_update.object_id AND index_columns.index_id = #indexes_to_update.index_id
    LEFT JOIN sys.objects ON objects.object_id = #indexes_to_update.object_id
    LEFT JOIN sys.columns ON columns.object_id = #indexes_to_update.object_id AND columns.column_id = index_columns.column_id
    WHERE is_primary_key = 0
    GROUP BY indexes.object_id, indexes.index_id, indexes.name, is_unique, indexes.type, objects.name, ignore_dup_key;

    -- run generated statements
    DECLARE c_statements CURSOR FOR SELECT * FROM #statements;
    OPEN c_statements;
    FETCH NEXT FROM c_statements INTO @statement, @progress;
    WHILE @@FETCH_STATUS = 0
    BEGIN
      BEGIN TRY
        EXEC(@statement);
      END TRY
      BEGIN CATCH
        IF ERROR_NUMBER() IN(1750, 1505) -- duplicate key in primary key/unique index
        BEGIN
          IF CHARINDEX(' files ', @statement) > 0
            EXEC fix_files_duplicates @new_collation;
          ELSE IF CHARINDEX(' directories ',  @statement) > 0
            EXEC fix_directories_duplicates @new_collation;
          ELSE IF CHARINDEX(' modules_file ',  @statement) > 0
            EXEC fix_modules_file_duplicates @new_collation;
          ELSE IF CHARINDEX(' modules_directory ',  @statement) > 0
            EXEC fix_modules_directory_duplicates @new_collation;
          ELSE
            THROW;
        END
        ELSE
          THROW;
      END CATCH;
      FETCH NEXT FROM c_statements INTO @statement, @progress;
    END;
    CLOSE c_statements;
    DEALLOCATE c_statements;

    DROP TABLE IF EXISTS #columns_to_update;
    DROP TABLE IF EXISTS #indexes_to_update;
    DROP TABLE IF EXISTS #statements;

    -- change database collation
    IF @current_collation <> @new_collation
    BEGIN
      EXEC('ALTER DATABASE CURRENT SET SINGLE_USER WITH ROLLBACK IMMEDIATE');
      EXEC('ALTER DATABASE CURRENT COLLATE ' + @new_collation);
      EXEC('ALTER DATABASE CURRENT SET MULTI_USER WITH ROLLBACK IMMEDIATE');
    END;
  END;
END
GO

EXEC full_unicode_support;

DROP PROCEDURE IF EXISTS fix_files_duplicates;
DROP PROCEDURE IF EXISTS fix_directories_duplicates;
DROP PROCEDURE IF EXISTS fix_modules_file_duplicates;
DROP PROCEDURE IF EXISTS fix_modules_directory_duplicates;
DROP PROCEDURE IF EXISTS full_unicode_support;