SPARK_CONF_DIR=/etc/spark3/conf \
SPARK_HOME=venv/lib/python3.10/site-packages/pyspark \
spark3-submit \
    --driver-cores 8 \
    --driver-memory 32G \
    --executor-cores 1 \
    --executor-memory 16g \
    --conf spark.executor.memoryOverhead=3GB \
    --conf spark.dynamicAllocation.maxExecutors=160 \
    --conf spark.executorEnv.SPARK_HOME=venv/lib/python3.10/site-packages/pyspark \
    --conf spark.executorEnv.SPARK_CONF_DIR=/etc/spark3/conf \
    --master yarn \
    --conf spark.dynamicAllocation.enabled=true \
    --conf spark.shuffle.service.enabled=true \
    --conf spark.yarn.maxAppAttempts=1 \
    --conf spark.shuffle.service.name=spark_shuffle_3_3 \
    --conf spark.shuffle.service.port=7338 \
    --conf spark.sql.shuffle.partitions=10240 \
    --conf write.spark.accept-any-schema=true \
    --conf spark.yarn.archive=hdfs:///user/spark/share/lib/spark-3.3.2-assembly.zip \
    --conf spark.jars.packages=org.apache.iceberg:iceberg-spark-runtime-3.3_2.12:1.6.1 \
    --conf spark.jars.ivySettings=/etc/maven/ivysettings.xml \
    --conf 'spark.driver.extraJavaOptions=-Divy.cache.dir=/tmp/table_maintenance_iceberg_monthly/ivy_spark3/cache -Divy.home=/tmp/table_maintenance_iceberg_monthly/ivy_spark3/home -Dlog4j.configuration=file:quieter_spark_log4j.properties' \
    --conf spark.sql.autoBroadcastJoinThreshold=-1 \
    --conf spark.sql.iceberg.locality.enabled=false \
    --conf spark.hadoop.mapreduce.fileoutputcommitter.algorithm.version=2 \
    --conf spark.yarn.appMasterEnv.SPARK_CONF_DIR=/etc/spark3/conf \
    --conf spark.yarn.appMasterEnv.SPARK_HOME=venv/lib/python3.10/site-packages/pyspark \
    --archives 'hdfs:///wmf/cache/artifacts/airflow/analytics/spark-3.3.2-1.0.0.conda.tgz#venv' \
    --name cp_mediawiki_content_history_v2 \
    --deploy-mode cluster \
    --class org.wikimedia.analytics.refinery.job.MediawikiContentHistoryMigration \
    ~/analytics_refinery_source/refinery-job/target/refinery-job-0.3.3-SNAPSHOT-shaded.jar \
    --source-table wmf_content.mediawiki_content_history_v1 \
    --target-table aqu.mediawiki_content_history_v1

package org.wikimedia.analytics.refinery.job

import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions.{col, struct}
import org.wikimedia.analytics.refinery.tools.LogHelper
import scopt.OptionParser

object MediawikiContentHistoryMigration extends LogHelper {

  case class Config(
      sourceTable: String = "wmf_content.mediawiki_content_history_v1",
      targetTable: String = "aqu.mediawiki_content_history_v1",
      dryRun: Boolean = false,
      appName: String = "MediawikiContentHistoryMigration"
  )

  val argsParser = new OptionParser[Config]("MediawikiContentHistoryMigration") {
    head("MediawikiContentHistoryMigration", "1.0")
    note("Migrate mediawiki_content_history_v1 table to add origin_rev_id field")
    
    opt[String]("source-table")
      .required()
      .action((x, c) => c.copy(sourceTable = x))
      .text("Source table name (default: wmf_content.mediawiki_content_history_v1)")
    
    opt[String]("target-table")
      .required()
      .action((x, c) => c.copy(targetTable = x))
      .text("Target table name (default: wmf_content.mediawiki_content_history_v1_tmp)")
    
    opt[Unit]("dry-run")
      .action((_, c) => c.copy(dryRun = true))
      .text("Dry run mode - only show what would be done")
    
    opt[String]("app-name")
      .action((x, c) => c.copy(appName = x))
      .text("Application name for Spark job (default: MediawikiContentHistoryMigration)")
    
    help("help").text("Print this usage text")
  }

  def main(args: Array[String]): Unit = {
    argsParser.parse(args, Config()) match {
      case Some(config) =>
        val spark = SparkSession.builder()
          .appName(config.appName)
          .enableHiveSupport()
          .getOrCreate()
        
        try {
          runMigration(spark, config)
        } finally {
          spark.stop()
        }
        
      case None =>
        sys.exit(1)
    }
  }

  def runMigration(spark: SparkSession, config: Config): Unit = {
    log.info(s"Starting migration from ${config.sourceTable} to ${config.targetTable}")
    log.info(s"Dry run: ${config.dryRun}")

    if (config.dryRun) {
      log.info("DRY RUN: Would migrate data using 4-batch strategy")
      return
    }

    // Batch 1: All wikis except the big 3
    log.info("Processing Batch 1: All wikis except enwiki, wikidatawiki, commonswiki")
    migrateWikiBatch(spark, config, notWiki = List("enwiki", "wikidatawiki", "commonswiki"))

    // Batch 2: enwiki
    log.info("Processing Batch 2: enwiki")
    migrateWikiBatch(spark, config, wiki = "enwiki")

    // Batch 3: wikidatawiki
    log.info("Processing Batch 3: wikidatawiki")
    migrateWikiBatch(spark, config, wiki = "wikidatawiki")

    // Batch 4: commonswiki
    log.info("Processing Batch 4: commonswiki")
    migrateWikiBatch(spark, config, wiki = "commonswiki")
    
    // Verification: Compare source and target table stats
    log.info("Verifying migration by comparing table statistics...")
    verifyMigration(spark, config)
    
    log.info("Migration completed successfully!")
  }

  def getWikiStats(spark: SparkSession, sourceTable: String): Array[(String, Long)] = {
    log.info("Getting wiki statistics...")
    
    val wikiStatsDf = spark.sql(s"""SELECT wiki_id, COUNT(*) as row_count
                                FROM $sourceTable
                                GROUP BY wiki_id
                                ORDER BY row_count DESC""")
    
    wikiStatsDf.collect().map(row => (row.getString(0), row.getLong(1)))
  }

  def migrateWikiBatch(spark: SparkSession, config: Config, wiki: String = null, notWiki: List[String] = List("enwiki", "wikidatawiki", "commonswiki")): Unit = {
    val sourceDf = if (wiki != null) {
      // Process specific wiki
      log.info(s"Processing single wiki: $wiki")
      spark.table(config.sourceTable).filter(col("wiki_id") === wiki)
    } else {
      // Process all wikis except the ones in notWiki list
      log.info(s"Processing all wikis except: ${notWiki.mkString(", ")}")
      spark.table(config.sourceTable).filter(!col("wiki_id").isin(notWiki: _*))
    }
    
    val migratedDf = addOriginRevIdToContentSlots(sourceDf)
    
    migratedDf.writeTo(config.targetTable).append()
    
    log.info(s"Completed batch migration")
  }

  def verifyMigration(spark: SparkSession, config: Config): Unit = {
    val sourceStats = getWikiStats(spark, config.sourceTable)
    val targetStats = getWikiStats(spark, config.targetTable)
    
    val sourceMap = sourceStats.toMap
    val targetMap = targetStats.toMap
    
    log.info("=== MIGRATION VERIFICATION ===")
    log.info(s"Source table: ${sourceStats.length} wikis, ${sourceStats.map(_._2).sum} total rows")
    log.info(s"Target table: ${targetStats.length} wikis, ${targetStats.map(_._2).sum} total rows")
    
    // Check for differences
    val allWikiIds = (sourceMap.keys ++ targetMap.keys).toSet
    
    var differences = 0
    allWikiIds.foreach { wikiId =>
      val sourceCount = sourceMap.getOrElse(wikiId, 0L)
      val targetCount = targetMap.getOrElse(wikiId, 0L)
      
      if (sourceCount != targetCount) {
        log.warn(s"DIFFERENCE: Wiki $wikiId - Source: $sourceCount, Target: $targetCount, Diff: ${targetCount - sourceCount}")
        differences += 1
      }
    }
    
    if (differences == 0) {
      log.info("✓ All wiki row counts match between source and target tables")
    } else {
      log.warn(s"⚠ Found $differences wikis with row count differences")
    }
    
    log.info("=== END VERIFICATION ===")
  }

  def addOriginRevIdToContentSlots(df: DataFrame): DataFrame = {
    df.select(
      df.schema.fields.map { field =>
        if (field.name == "revision_content_slots") {
          transformContentSlotsMap(col(field.name), col("revision_id")).as("revision_content_slots")
        } else {
          col(field.name)
        }
      }: _*
    )
  }

  def transformContentSlotsMap(contentSlotsCol: org.apache.spark.sql.Column, revisionIdCol: org.apache.spark.sql.Column): org.apache.spark.sql.Column = {
    org.apache.spark.sql.functions.transform_values(
      contentSlotsCol,
      (key: org.apache.spark.sql.Column, valueStruct: org.apache.spark.sql.Column) => {
        struct(
          valueStruct.getField("content_body").as("content_body"),
          valueStruct.getField("content_format").as("content_format"),
          valueStruct.getField("content_model").as("content_model"),
          valueStruct.getField("content_sha1").as("content_sha1"),
          valueStruct.getField("content_size").as("content_size"),
          revisionIdCol.as("origin_rev_id")
        )
      }
    )
  }
}