var shop = use('internal/shop')
var pkg = use('package')
var fd = use('fd')
var time = use('time')
var json = use('json')
var blob = use('blob')
var os = use('os')
var testlib = use('internal/testlib')
var math = use('math/radians')

if (!args) args = []

var target_pkg = null // null = current package
var target_bench = null // null = all benchmarks, otherwise specific bench file
var all_pkgs = false

// Benchmark configuration
def WARMUP_BATCHES = 3
def SAMPLES = 11 // Number of timing samples to collect
def TARGET_SAMPLE_NS = 20000000 // 20ms per sample (fast mode)
def MIN_SAMPLE_NS = 2000000 // 2ms minimum sample duration
def MIN_BATCH_SIZE = 1
def MAX_BATCH_SIZE = 100000000 // 100M iterations max per batch

// Statistical functions
function median(arr) {
  if (arr.length == 0) return 0
  var sorted = sort(arr)
  var mid = floor(arr.length / 2)
  if (arr.length % 2 == 0) {
    return (sorted[mid - 1] + sorted[mid]) / 2
  }
  return sorted[mid]
}

function mean(arr) {
  if (arr.length == 0) return 0
  var sum = 0
  for (var i = 0; i < arr.length; i++) {
    sum += arr[i]
  }
  return sum / arr.length
}

function stddev(arr, mean_val) {
  if (arr.length < 2) return 0
  var sum_sq_diff = 0
  for (var i = 0; i < arr.length; i++) {
    var diff = arr[i] - mean_val
    sum_sq_diff += diff * diff
  }
  return math.sqrt(sum_sq_diff / (arr.length - 1))
}

function percentile(arr, p) {
  if (arr.length == 0) return 0
  var sorted = sort(arr)
  var idx = floor(arr.length * p / 100)
  if (idx >= arr.length) idx = arr.length - 1
  return sorted[idx]
}

function min_val(arr) {
  if (arr.length == 0) return 0
  var m = arr[0]
  for (var i = 1; i < arr.length; i++) {
    if (arr[i] < m) m = arr[i]
  }
  return m
}

function max_val(arr) {
  if (arr.length == 0) return 0
  var m = arr[0]
  for (var i = 1; i < arr.length; i++) {
    if (arr[i] > m) m = arr[i]
  }
  return m
}

// Parse arguments similar to test.ce
function parse_args() {
  if (args.length == 0) {
    if (!testlib.is_valid_package('.')) {
      log.console('No cell.toml found in current directory')
      return false
    }
    target_pkg = null
    return true
  }

  if (args[0] == 'all') {
    if (!testlib.is_valid_package('.')) {
      log.console('No cell.toml found in current directory')
      return false
    }
    target_pkg = null
    return true
  }

  if (args[0] == 'package') {
    if (args.length < 2) {
      log.console('Usage: cell bench package <name> [bench]')
      log.console('       cell bench package all')
      return false
    }

    if (args[1] == 'all') {
      all_pkgs = true
      log.console('Benchmarking all packages...')
      return true
    }

    var name = args[1]
    var lock = shop.load_lock()
    if (lock[name]) {
      target_pkg = name
    } else if (starts_with(name, '/') && testlib.is_valid_package(name)) {
      target_pkg = name
    } else {
      if (testlib.is_valid_package('.')) {
        var resolved = pkg.alias_to_package(null, name)
        if (resolved) {
          target_pkg = resolved
        } else {
          log.console(`Package not found: ${name}`)
          return false
        }
      } else {
        log.console(`Package not found: ${name}`)
        return false
      }
    }

    if (args.length >= 3) {
      target_bench = args[2]
    }

    log.console(`Benchmarking package: ${target_pkg}`)
    return true
  }

  // cell bench benches/suite or cell bench <path>
  var bench_path = args[0]

  // Normalize path - add benches/ prefix if not present
  if (!starts_with(bench_path, 'benches/') && !starts_with(bench_path, '/')) {
    if (!fd.is_file(bench_path + '.cm') && !fd.is_file(bench_path)) {
      if (fd.is_file('benches/' + bench_path + '.cm') || fd.is_file('benches/' + bench_path)) {
        bench_path = 'benches/' + bench_path
      }
    }
  }

  target_bench = bench_path
  target_pkg = null

  if (!testlib.is_valid_package('.')) {
    log.console('No cell.toml found in current directory')
    return false
  }

  return true
}

if (!parse_args()) {
  $stop()
  return
}

// Collect benchmark files from a package
function collect_benches(package_name, specific_bench) {
  var prefix = testlib.get_pkg_dir(package_name)
  var benches_dir = prefix + '/benches'

  if (!fd.is_dir(benches_dir)) return []

  var files = pkg.list_files(package_name)
  var bench_files = []
  for (var i = 0; i < files.length; i++) {
    var f = files[i]
    if (starts_with(f, "benches/") && ends_with(f, ".cm")) {
      if (specific_bench) {
        var bench_name = text(f, 0, -3)
        var match_name = specific_bench
        if (!starts_with(match_name, 'benches/')) match_name = 'benches/' + match_name
        var match_base = ends_with(match_name, '.cm') ? text(match_name, 0, -3) : match_name
        if (bench_name != match_base) continue
      }
      bench_files.push(f)
    }
  }
  return bench_files
}

// Calibrate batch size for a benchmark
function calibrate_batch_size(bench_fn, is_batch) {
  if (!is_batch) return 1

  var n = MIN_BATCH_SIZE
  var dt = 0

  // Find a batch size that takes at least MIN_SAMPLE_NS
  while (n < MAX_BATCH_SIZE) {
    // Ensure n is a valid number before calling
    if (!is_number(n) || n < 1) {
      n = 1
      break
    }

    var start = os.now()
    bench_fn(n)
    dt = os.now() - start

    if (dt >= MIN_SAMPLE_NS) break

    // Double the batch size
    var new_n = n * 2
    // Check if multiplication produced a valid number
    if (!is_number(new_n) || new_n > MAX_BATCH_SIZE) {
      n = MAX_BATCH_SIZE
      break
    }
    n = new_n
  }

  // Adjust to target sample duration
  if (dt > 0 && dt < TARGET_SAMPLE_NS && is_number(n) && is_number(dt)) {
    var calc = n * TARGET_SAMPLE_NS / dt
    if (is_number(calc) && calc > 0) {
      var target_n = floor(calc)
      // Check if floor returned a valid number
      if (is_number(target_n) && target_n > 0) {
        if (target_n > MAX_BATCH_SIZE) target_n = MAX_BATCH_SIZE
        if (target_n < MIN_BATCH_SIZE) target_n = MIN_BATCH_SIZE
        n = target_n
      }
    }
  }

  // Safety check - ensure we always return a valid batch size
  if (!is_number(n) || n < 1) {
    n = 1
  }

  return n
}

// Run a single benchmark function
function run_single_bench(bench_fn, bench_name) {
  var timings_per_op = []

  // Detect benchmark format:
  // 1. Object with { setup, run, teardown } - structured format
  // 2. Function that accepts (n) - batch format
  // 3. Function that accepts () - legacy format
  var is_structured = is_object(bench_fn) && bench_fn.run
  var is_batch = false
  var batch_size = 1
  var setup_fn = null
  var run_fn = null
  var teardown_fn = null

  if (is_structured) {
    setup_fn = bench_fn.setup || function() { return null }
    run_fn = bench_fn.run
    teardown_fn = bench_fn.teardown || function(state) {}

    // Check if run function accepts batch size
    try {
      var test_state = setup_fn()
      run_fn(1, test_state)
      is_batch = true
      if (teardown_fn) teardown_fn(test_state)
    } catch (e) {
      is_batch = false
    }

    // Create wrapper for calibration
    var calibrate_fn = function(n) {
      var state = setup_fn()
      run_fn(n, state)
      if (teardown_fn) teardown_fn(state)
    }
    batch_size = calibrate_batch_size(calibrate_fn, is_batch)

    // Safety check for structured benchmarks
    if (!is_number(batch_size) || batch_size < 1) {
      batch_size = 1
    }
  } else {
    // Simple function format
    try {
      bench_fn(1)
      is_batch = true
    } catch (e) {
      is_batch = false
    }
    batch_size = calibrate_batch_size(bench_fn, is_batch)
  }

  // Safety check - ensure batch_size is valid
  if (!batch_size || batch_size < 1) {
    batch_size = 1
  }

  // Warmup phase
  for (var i = 0; i < WARMUP_BATCHES; i++) {
    // Ensure batch_size is valid before warmup
    if (!is_number(batch_size) || batch_size < 1) {
      var type_str = is_null(batch_size) ? 'null' : is_number(batch_size) ? 'number' : is_text(batch_size) ? 'text' : is_object(batch_size) ? 'object' : is_array(batch_size) ? 'array' : is_function(batch_size) ? 'function' : is_logical(batch_size) ? 'logical' : 'unknown'
      log.console(`WARNING: batch_size became ${type_str} = ${batch_size}, resetting to 1`)
      batch_size = 1
    }

    if (is_structured) {
      var state = setup_fn()
      if (is_batch) {
        run_fn(batch_size, state)
      } else {
        run_fn(state)
      }
      if (teardown_fn) teardown_fn(state)
    } else {
      if (is_batch) {
        bench_fn(batch_size)
      } else {
        bench_fn()
      }
    }
  }

  // Measurement phase - collect SAMPLES timing samples
  for (var i = 0; i < SAMPLES; i++) {
    // Double-check batch_size is valid (should never happen, but defensive)
    if (!is_number(batch_size) || batch_size < 1) {
      batch_size = 1
    }

    if (is_structured) {
      var state = setup_fn()
      var start = os.now()
      if (is_batch) {
        run_fn(batch_size, state)
      } else {
        run_fn(state)
      }
      var duration = os.now() - start
      if (teardown_fn) teardown_fn(state)

      var ns_per_op = is_batch ? duration / batch_size : duration
      timings_per_op.push(ns_per_op)
    } else {
      var start = os.now()
      if (is_batch) {
        bench_fn(batch_size)
      } else {
        bench_fn()
      }
      var duration = os.now() - start

      var ns_per_op = is_batch ? duration / batch_size : duration
      timings_per_op.push(ns_per_op)
    }
  }

  // Calculate statistics
  var mean_ns = mean(timings_per_op)
  var median_ns = median(timings_per_op)
  var min_ns = min_val(timings_per_op)
  var max_ns = max_val(timings_per_op)
  var stddev_ns = stddev(timings_per_op, mean_ns)
  var p95_ns = percentile(timings_per_op, 95)
  var p99_ns = percentile(timings_per_op, 99)

  // Calculate ops/s from median
  var ops_per_sec = 0
  if (median_ns > 0) {
    ops_per_sec = floor(1000000000 / median_ns)
  }

  return {
    name: bench_name,
    batch_size: batch_size,
    samples: SAMPLES,
    mean_ns: round(mean_ns),
    median_ns: round(median_ns),
    min_ns: round(min_ns),
    max_ns: round(max_ns),
    stddev_ns: round(stddev_ns),
    p95_ns: round(p95_ns),
    p99_ns: round(p99_ns),
    ops_per_sec: ops_per_sec
  }
}

// Format nanoseconds for display
function format_ns(ns) {
  if (ns < 1000) return `${ns}ns`
  if (ns < 1000000) return `${round(ns / 1000 * 100) / 100}µs`
  if (ns < 1000000000) return `${round(ns / 1000000 * 100) / 100}ms`
  return `${round(ns / 1000000000 * 100) / 100}s`
}

// Format ops/sec for display
function format_ops(ops) {
  if (ops < 1000) return `${ops} ops/s`
  if (ops < 1000000) return `${round(ops / 1000 * 100) / 100}K ops/s`
  if (ops < 1000000000) return `${round(ops / 1000000 * 100) / 100}M ops/s`
  return `${round(ops / 1000000000 * 100) / 100}G ops/s`
}

// Run benchmarks for a package
function run_benchmarks(package_name, specific_bench) {
  var bench_files = collect_benches(package_name, specific_bench)

  var pkg_result = {
    package: package_name || "local",
    files: [],
    total: 0
  }

  if (bench_files.length == 0) return pkg_result

  if (package_name) log.console(`Running benchmarks for ${package_name}`)
  else log.console(`Running benchmarks for local package`)

  for (var i = 0; i < bench_files.length; i++) {
    var f = bench_files[i]
    var mod_path = text(f, 0, -3)

    var file_result = {
      name: f,
      benchmarks: []
    }

    try {
      var bench_mod
      var use_pkg = package_name ? package_name : fd.realpath('.')
      bench_mod = shop.use(mod_path, use_pkg)

      var benches = []
      if (is_function(bench_mod)) {
        benches.push({name: 'main', fn: bench_mod})
      } else if (is_object(bench_mod)) {
        for (var k in bench_mod) {
          if (is_function(bench_mod[k])) {
            benches.push({name: k, fn: bench_mod[k]})
          }
        }
      }

      if (benches.length > 0) {
        log.console(`  ${f}`)
        for (var j = 0; j < benches.length; j++) {
          var b = benches[j]
          try {
            var result = run_single_bench(b.fn, b.name)
            result.package = pkg_result.package
            file_result.benchmarks.push(result)
            pkg_result.total++

            log.console(`    ${result.name}`)
            log.console(`      ${format_ns(result.median_ns)}/op  ${format_ops(result.ops_per_sec)}`)
            log.console(`      min: ${format_ns(result.min_ns)}  max: ${format_ns(result.max_ns)}  stddev: ${format_ns(result.stddev_ns)}`)
            if (result.batch_size > 1) {
              log.console(`      batch: ${result.batch_size}  samples: ${result.samples}`)
            }
          } catch (e) {
            log.console(`    ERROR ${b.name}: ${e}`)
            log.error(e)
            var error_result = {
              package: pkg_result.package,
              name: b.name,
              error: e.toString()
            }
            file_result.benchmarks.push(error_result)
            pkg_result.total++
          }
        }
      }
    } catch (e) {
      log.console(`  Error loading ${f}: ${e}`)
      var error_result = {
        package: pkg_result.package,
        name: "load_module",
        error: `Error loading module: ${e}`
      }
      file_result.benchmarks.push(error_result)
      pkg_result.total++
    }

    if (file_result.benchmarks.length > 0) {
      pkg_result.files.push(file_result)
    }
  }

  return pkg_result
}

// Run all benchmarks
var all_results = []

if (all_pkgs) {
  if (testlib.is_valid_package('.')) {
    all_results.push(run_benchmarks(null, null))
  }

  var packages = shop.list_packages()
  for (var i = 0; i < packages.length; i++) {
    all_results.push(run_benchmarks(packages[i], null))
  }
} else {
  all_results.push(run_benchmarks(target_pkg, target_bench))
}

// Calculate totals
var total_benches = 0
for (var i = 0; i < all_results.length; i++) {
  total_benches += all_results[i].total
}

log.console(`----------------------------------------`)
log.console(`Benchmarks: ${total_benches} total`)

// Generate reports
function generate_reports() {
  var timestamp = text(floor(time.number()))
  var report_dir = shop.get_reports_dir() + '/bench_' + timestamp
  testlib.ensure_dir(report_dir)

  var txt_report = `BENCHMARK REPORT
Date: ${time.text(time.number())}
Total benchmarks: ${total_benches}

=== SUMMARY ===
`
  for (var i = 0; i < all_results.length; i++) {
    var pkg_res = all_results[i]
    if (pkg_res.total == 0) continue
    txt_report += `Package: ${pkg_res.package}\n`
    for (var j = 0; j < pkg_res.files.length; j++) {
      var f = pkg_res.files[j]
      txt_report += `  ${f.name}\n`
      for (var k = 0; k < f.benchmarks.length; k++) {
        var b = f.benchmarks[k]
        if (b.error) {
          txt_report += `    ERROR ${b.name}: ${b.error}\n`
        } else {
          txt_report += `    ${b.name}: ${format_ns(b.median_ns)}/op (${format_ops(b.ops_per_sec)})\n`
        }
      }
    }
  }

  txt_report += `\n=== DETAILED RESULTS ===\n`
  for (var i = 0; i < all_results.length; i++) {
    var pkg_res = all_results[i]
    if (pkg_res.total == 0) continue

    for (var j = 0; j < pkg_res.files.length; j++) {
      var f = pkg_res.files[j]
      for (var k = 0; k < f.benchmarks.length; k++) {
        var b = f.benchmarks[k]
        if (b.error) continue

        txt_report += `\n${pkg_res.package}::${b.name}\n`
        txt_report += `  batch_size: ${b.batch_size}  samples: ${b.samples}\n`
        txt_report += `  median: ${format_ns(b.median_ns)}/op\n`
        txt_report += `  mean: ${format_ns(b.mean_ns)}/op\n`
        txt_report += `  min: ${format_ns(b.min_ns)}\n`
        txt_report += `  max: ${format_ns(b.max_ns)}\n`
        txt_report += `  stddev: ${format_ns(b.stddev_ns)}\n`
        txt_report += `  p95: ${format_ns(b.p95_ns)}\n`
        txt_report += `  p99: ${format_ns(b.p99_ns)}\n`
        txt_report += `  ops/s: ${format_ops(b.ops_per_sec)}\n`
      }
    }
  }

  testlib.ensure_dir(report_dir)
  fd.slurpwrite(`${report_dir}/bench.txt`, stone(blob(txt_report)))
  log.console(`Report written to ${report_dir}/bench.txt`)

  // Generate JSON per package
  for (var i = 0; i < all_results.length; i++) {
    var pkg_res = all_results[i]
    if (pkg_res.total == 0) continue

    var pkg_benches = []
    for (var j = 0; j < pkg_res.files.length; j++) {
      var f = pkg_res.files[j]
      for (var k = 0; k < f.benchmarks.length; k++) {
        pkg_benches.push(f.benchmarks[k])
      }
    }

    var json_path = `${report_dir}/${replace(pkg_res.package, /\//, '_')}.json`
    fd.slurpwrite(json_path, stone(blob(json.encode(pkg_benches))))
  }
}

generate_reports()
$stop()