require 'json'
require 'fileutils'
require 'set'
require_relative '../helpers/github'
require_relative '../helpers/pretty_table'


module Fastlane
  module Actions
    class BenchmarkWorkflowAction < Action
      def self.run(params)
        repo_name = params[:repo].include?("/") ? Github.repo(url: params[:repo]) : params[:repo]
        pr_number = params[:pr]
        workflow_head = params[:workflow]
        workflow_base = params[:workflow_base] || params[:workflow]
        run_total = params[:runs]
        timeout_minutes = params[:timeout]
        interval_seconds = params[:interval]
        threshold = params[:threshold]

        benchmark_url = nil
        benchmark_context = "CI Benchmark / #{workflow_head.split(".").first}"
        if params[:benchmark_run_id] && params[:benchmark_repo]
          benchmark_url = "https://github.com/#{params[:benchmark_repo]}/actions/runs/#{params[:benchmark_run_id]}"
        end

        pr = Github.lookup_pr_by_number(number: pr_number, repo: repo_name, fields: "baseRefOid,baseRefName,headRefOid,headRefName")
        unless pr
          UI.user_error!("Could not find PR ##{pr_number} in repository #{Github.org}/#{repo_name}")
        end
        base_branch = pr[:baseRefName]
        base_sha = params[:sha_base] || pr[:baseRefOid]
        head_branch = pr[:headRefName]
        head_sha = pr[:headRefOid]
        workflow_url = "https://github.com/#{Github.org}/#{repo_name}/blob/#{base_branch}/.github/workflows/#{workflow_base}"
        UI.header("Starting benchmark dispatch for PR ##{pr_number} in #{Github.org}/#{repo_name}...")

        begin
          # Create initial status check
          Github.post_status(
            repo: repo_name,
            sha: head_sha,
            context: benchmark_context,
            state: "pending",
            description: "Dispatching runs...",
            target_url: benchmark_url
          )

          runs = [
            *dispatch_runs(
              workflow: workflow_base,
              branch: base_branch,
              sha: base_sha,
              repo: repo_name,
              limit: run_total,
              dry_run: params[:dry_run]
            ),
            *dispatch_runs(
              workflow: workflow_head,
              branch: head_branch,
              sha: head_sha,
              repo: repo_name,
              limit: run_total,
              dry_run: params[:dry_run]
            ),
          ]

          # Wait for all runs to complete with timeout
          start_time = Time.now
          while runs.select { |run| ["queued", "in_progress"].include?(run[:status]) }.count > 0
            pending_runs = runs.select { |run| ["queued", "in_progress"].include?(run[:status]) }
            Github.post_status(
              repo: repo_name,
              sha: head_sha,
              context: benchmark_context,
              state: "pending",
              description: "Waiting for #{pending_runs.count} pending run(s) to complete...",
              target_url: benchmark_url
            )
            PrettyTable.print(data: runs, title: "Run Status")
            UI.message("Waiting for #{pending_runs.count} pending run(s) to complete...")
            sleep interval_seconds

            runs = runs.map do |run|
              run = if ["queued", "in_progress"].include?(run[:status])
                {
                  **run,
                  **Github::Workflows.get_run(run_id: run[:databaseId], repo: repo_name),
                }
              end
              UI.message("#{run[:headBranch].yellow}: #{run[:status].upcase.cyan} #{run[:url]}")
              run
            end

            elapsed_minutes = ((Time.now - start_time).to_i / 60).floor
            if elapsed_minutes > timeout_minutes
              # Cancel pending runs and mark them as timed out
              UI.important("Timeout reached, canceling pending runs...")
              runs = runs.map do |run|
                if ["queued", "in_progress"].include?(run[:status])
                  UI.message("Canceling run #{run[:databaseId]}...")
                  begin
                    Github::Workflows.cancel(run_id: run[:databaseId], repo: repo_name)
                  rescue => e
                    UI.important("Failed to cancel run #{run[:databaseId]}: #{e.message}")
                  end
                  {
                    **run,
                    status: "workflow_timeout",
                    updatedAt: Time.now,
                  }
                else
                  run
                end
              end
              break
            end
          end

          UI.header("Generating benchmark report")
          runs = runs.map do |run|
            {
              **run,
              duration: run[:status] == 'completed' ? (Time.parse(run[:updatedAt]) - Time.parse(run[:createdAt])).to_i : 0
            }
          end
          runs_base = runs.filter { |run| run[:headBranch] == base_branch }
          runs_head = runs.filter { |run| run[:headBranch] == head_branch }

          stats_base = compute_stats(runs: runs_base)
          stats_head = compute_stats(runs: runs_head)
          stats_diff = compare_stats(base_stats: stats_base, head_stats: stats_head, threshold: threshold)
          # Generate markdown report
          benchmark_summary = <<~MARKDOWN
          <h2>CI Benchmark / #{workflow_head}</h2>
          <table>
            <thead>
              <tr>
                <th><a href="#{base_sha}">Base Branch</a></th>
                <th>Runtime</th>
                <th><a href="#{head_sha}">PR Branch</a></th>
                <th>Runtime</th>
                <th>Diff</th>
              </tr>
            </thead>
            <tbody>
              #{(0..[runs_base.count, runs_head.count].min-1).map { |i| "<tr><td align='center'>#{format_run_id(runs_base[i], i)}</td><td>#{format_run_result(runs_base[i])}</td><td align='center'>#{format_run_id(runs_head[i], i)}</td><td>#{format_run_result(runs_head[i])}</td><td>#{compare_runs(base_run: runs_base[i], head_run: runs_head[i])}</td></tr>" }.join("\n")}
            </tbody>
            <tfoot>
              <tr>
                <th>Summary</th>
                <td><strong>#{stats_base[:summary]}</strong></td>
                <th>Summary</th>
                <td><strong>#{stats_head[:summary]}</strong></td>
                <td><strong>#{stats_diff[:emoji]} #{stats_diff[:diff]}</strong></td>
              </tr>
              <tr>
                <th colspan="5">#{stats_diff[:emoji]} #{stats_diff[:summary]} #{stats_diff[:emoji]}</th>
              </tr>
            </tfoot>
          </table>

          <em>Generated by <a href='https://github.com/#{Github.org}/#{Github.repo}/blob/main/docs/benchmark.md'>CI Benchmark</a> command: `/benchmark workflow:#{workflow_head} runs:#{run_total} timeout:#{timeout_minutes} minutes`</em>

          MARKDOWN

          Github.post_comment(
            pr_number: pr_number,
            body: benchmark_summary,
            repo: repo_name,
            edit_last: params[:edit_last]
          )

          Github.post_status(
            repo: repo_name,
            sha: head_sha,
            context: benchmark_context,
            state: "success",
            description: stats_diff[:summary],
            target_url: benchmark_url
          )
        rescue => e
          UI.message("Benchmark failed due to a system error!")
          Github.post_status(
            repo: repo_name,
            sha: head_sha,
            context: benchmark_context,
            state: "failure",
            description: "Benchmark system error: #{e.message}",
            target_url: benchmark_url
          )
          raise e
        end
      end

      def self.dispatch_runs(workflow:, branch:, sha:, repo:, limit:, dry_run:)
        if dry_run
          UI.message("Checking existing runs on branch #{branch}...")
          existing_runs = Github::Workflows.get_runs(workflow: workflow, branch: branch, sha: sha, repo: repo, limit: limit)
          UI.message("Found #{existing_runs.count} run(s) on branch #{branch}!")
          existing_runs
        else
          limit.times.map do |i|
            UI.message("Dispatching run #{i + 1}/#{limit} [#{branch}]...")
            Github::Workflows.dispatch(
              workflow: workflow,
              branch: branch,
              sha: sha,
              repo: repo,
              inputs: {
                concurrency_key: "benchmark-#{repo}-#{workflow}-#{branch}-#{sha}-#{i + 1}"
              }
            )
          end
        end
      end

      def self.format_duration(seconds)
        hours = (seconds / 3600).to_i
        minutes = ((seconds % 3600) / 60).to_i
        seconds = (seconds % 60).to_i
        [{value: hours, suffix: 'h'}, {value: minutes, suffix: 'm'}, {value: seconds, suffix: 's'}].reject { |s| s[:value].zero? }.map { |s| "#{s[:value]}#{s[:suffix]}" }.slice(0, 2).join('')
      end

      def self.format_duration_diff(base_duration:, head_duration:)
        change = head_duration - base_duration
        abs_change = change.abs
        pct_change = base_duration > 0 ? (abs_change * 100.0 / base_duration ).round(2) : 0
        if change == 0
          return "0s"
        end
        sign = change < 0 ? "-" : "+"
        return "#{sign}#{format_duration(abs_change)} (#{sign}#{pct_change}%)"
      end

      def self.compare_runs(base_run:, head_run:)
        if (base_run[:status] != "completed") || (head_run[:status] != "completed")
          return ":question:"
        end
        diff = format_duration_diff(base_duration: base_run[:duration], head_duration: head_run[:duration])
        if base_run[:duration] > head_run[:duration]
          return ":arrow_down_small: #{diff}"
        elsif base_run[:duration] < head_run[:duration]
          return ":arrow_up_small: #{diff}"
        else
          return ":shrug: #{diff}"
        end
      end

      def self.compute_stats(runs:)
        successful_runs = runs.filter { |r| r[:status] == 'completed' }
        count = successful_runs.count
        if count > 0
          durations = successful_runs.map { |r| (Time.parse(r[:updatedAt]) - Time.parse(r[:createdAt])).to_i }
          mean = durations.sum / count
          variance = durations.map { |d| (d - mean) ** 2 }.sum / count
          std_dev = Math.sqrt(variance)
          error_margin = 1.96 * std_dev / Math.sqrt(count) # 95% confidence interval
          summary = ":clock3: #{format_duration(mean)} #{error_margin != 0 ? "±#{format_duration(error_margin)}" : ''}"
        else
          mean = 0
          std_dev = 0
          error_margin = 0
          summary = ":x: No successful runs"
        end
        return {
          count: count,
          mean: mean,
          stdDev: std_dev,
          errorMargin: error_margin,
          summary: summary
        }
      end

      def self.compare_stats(base_stats:, head_stats:, threshold:)
        pct_change = base_stats[:mean] > 0 ? ((head_stats[:mean] - base_stats[:mean]) * 100.0 / base_stats[:mean]).round(2) : 0
        duration_diff = format_duration_diff(base_duration: base_stats[:mean], head_duration: head_stats[:mean])
        if base_stats[:count] == 0 || head_stats[:count] == 0
          {
            emoji: ":x:",
            summary: "Unable to complete the benchmark due to workflow failures",
            diff: ""
          }
        elsif pct_change.abs <= threshold
          {
            emoji: ":shrug:",
            diff: duration_diff,
            summary: "No significant change in runtime",
          }
        elsif head_stats[:mean] < base_stats[:mean]
          {
            emoji: ":green_heart:",
            diff: duration_diff,
            summary: "Significant runtime improvement",
          }
        else
          {
            emoji: ":bangbang:",
            diff: duration_diff,
            summary: "Significant runtime regression",
          }
        end
      end



      def self.format_run_result(run)
        if run[:status] == 'completed'
          return ":clock3: #{format_duration(run[:duration])}"
        else
          return ":x: #{run[:status] || 'dispatch_error'}"
        end
      end

      def self.format_run_id(run, index)
        "<a href='#{run[:url] || ''}'>Run #{index + 1}</a>"
      end

      def self.description
        "Generate and post benchmark report comparing workflow performance between base and head commits"
      end

      def self.return_value
        "Posts benchmark comparison comment to the PR"
      end

      def self.available_options
        [
          FastlaneCore::ConfigItem.new(
            key: :repo,
            description: "Repository name without the org prefix",
            type: String,
            verify_block: proc do |value|
              UI.user_error!(':repo option cannot be empty') if value.to_s.empty?
            end
          ),
          FastlaneCore::ConfigItem.new(
            key: :pr,
            description: "Pull request number",
            type: String,
            verify_block: proc do |value|
              UI.user_error!(':pr option cannot be empty') if value.to_s.empty?
            end
          ),
          FastlaneCore::ConfigItem.new(
            key: :comment,
            description: "Comment ID that triggered the benchmark",
            type: String,
            optional: true
          ),
          FastlaneCore::ConfigItem.new(
            key: :workflow,
            description: "Workflow name (e.g., 'ci.yml')",
            type: String,
            verify_block: proc do |value|
              UI.user_error!(':workflow option cannot be empty') if value.to_s.empty?
            end
          ),
          FastlaneCore::ConfigItem.new(
            key: :workflow_base,
            description: "Workflow name for the base runs (e.g., 'ci.yml')",
            optional: true,
            type: String,
          ),
          FastlaneCore::ConfigItem.new(
            key: :sha_base,
            description: "SHA of the base commit",
            optional: true,
            type: String,
          ),
          FastlaneCore::ConfigItem.new(
            key: :runs,
            description: "Number of benchmark runs performed",
            type: Integer,
            default_value: 5
          ),
          FastlaneCore::ConfigItem.new(
            key: :timeout,
            description: "Timeout in minutes for waiting for the workflow run to start",
            type: Integer,
            default_value: 60
          ),
          FastlaneCore::ConfigItem.new(
            key: :interval,
            description: "Interval in seconds for waiting for the workflow run to start",
            type: Integer,
            default_value: 30
          ),
          FastlaneCore::ConfigItem.new(
            key: :benchmark_run_id,
            env_name: "GITHUB_RUN_ID",
            description: "Run ID that triggered the benchmark",
            type: String,
            optional: true
          ),
          FastlaneCore::ConfigItem.new(
            key: :benchmark_repo,
            env_name: "GITHUB_REPOSITORY",
            description: "Repository that triggered the benchmark",
            type: String,
            optional: true
          ),
          FastlaneCore::ConfigItem.new(
            key: :threshold,
            description: "Threshold in percentage for the benchmark comparison",
            type: Integer,
            default_value: 5
          ),
          FastlaneCore::ConfigItem.new(
            key: :edit_last,
            description: "Edit the last comment if it exists",
            type: Boolean,
            default_value: false
          ),
          FastlaneCore::ConfigItem.new(
            key: :dry_run,
            description: "Avoid dispatching new runs, and look up existing runs for analysis",
            type: Boolean,
            default_value: false
          )
        ]
      end

      def self.is_supported?(platform)
        true
      end
    end
  end
end