import Foundation
import Darwin

private let defaultFunction = "{{DEFAULT_FUNCTION}}"
private let defaultIterations: UInt32 = 20
private let defaultWarmup: UInt32 = 3

struct BenchParams {
    let function: String
    let iterations: UInt32
    let warmup: UInt32

    private struct EncodedBenchSpec: Decodable {
        let function: String
        let iterations: UInt32
        let warmup: UInt32
    }

    static func fromBundle() -> BenchParams? {
        guard let url = Bundle.main.url(forResource: "bench_spec", withExtension: "json") else {
            print("[BenchRunner] No bench_spec.json found in bundle, will use process info or defaults")
            return nil
        }
        do {
            let data = try Data(contentsOf: url)
            let decoded = try JSONDecoder().decode(EncodedBenchSpec.self, from: data)
            print("[BenchRunner] Loaded config from bench_spec.json: function=\(decoded.function), iterations=\(decoded.iterations), warmup=\(decoded.warmup)")
            return BenchParams(function: decoded.function, iterations: decoded.iterations, warmup: decoded.warmup)
        } catch {
            print("[BenchRunner] ERROR: Failed to parse bench_spec.json: \(error)")
            print("[BenchRunner] Will fall back to process info or defaults")
            return nil
        }
    }

    static func fromProcessInfo() -> BenchParams {
        let info = ProcessInfo.processInfo
        var function = defaultFunction
        var iterations = defaultIterations
        var warmup = defaultWarmup
        var sources: [String] = []

        // Check environment variables
        if let envFunction = info.environment["BENCH_FUNCTION"], !envFunction.isEmpty {
            function = envFunction
            sources.append("function from BENCH_FUNCTION env")
        }

        if let envIterations = info.environment["BENCH_ITERATIONS"], !envIterations.isEmpty {
            if let parsed = UInt32(envIterations) {
                iterations = parsed
                sources.append("iterations from BENCH_ITERATIONS env")
            } else {
                print("[BenchRunner] WARNING: Failed to parse BENCH_ITERATIONS='\(envIterations)' as integer, using default: \(defaultIterations)")
            }
        }

        if let envWarmup = info.environment["BENCH_WARMUP"], !envWarmup.isEmpty {
            if let parsed = UInt32(envWarmup) {
                warmup = parsed
                sources.append("warmup from BENCH_WARMUP env")
            } else {
                print("[BenchRunner] WARNING: Failed to parse BENCH_WARMUP='\(envWarmup)' as integer, using default: \(defaultWarmup)")
            }
        }

        // Check launch arguments (override environment variables)
        for arg in info.arguments {
            if arg.hasPrefix("--bench-function=") {
                if let value = arg.split(separator: "=", maxSplits: 1).last {
                    function = String(value)
                    sources.append("function from --bench-function arg")
                }
            } else if arg.hasPrefix("--bench-iterations=") {
                if let value = arg.split(separator: "=", maxSplits: 1).last {
                    if let parsed = UInt32(value) {
                        iterations = parsed
                        sources.append("iterations from --bench-iterations arg")
                    } else {
                        print("[BenchRunner] WARNING: Failed to parse --bench-iterations='\(value)' as integer, using current value: \(iterations)")
                    }
                }
            } else if arg.hasPrefix("--bench-warmup=") {
                if let value = arg.split(separator: "=", maxSplits: 1).last {
                    if let parsed = UInt32(value) {
                        warmup = parsed
                        sources.append("warmup from --bench-warmup arg")
                    } else {
                        print("[BenchRunner] WARNING: Failed to parse --bench-warmup='\(value)' as integer, using current value: \(warmup)")
                    }
                }
            }
        }

        // Log resolution source
        if sources.isEmpty {
            print("[BenchRunner] Using hardcoded defaults: function=\(function), iterations=\(iterations), warmup=\(warmup)")
        } else {
            print("[BenchRunner] Resolved params from process info: function=\(function), iterations=\(iterations), warmup=\(warmup) (sources: \(sources.joined(separator: ", ")))")
        }

        return BenchParams(function: function, iterations: iterations, warmup: warmup)
    }

    static func resolved() -> BenchParams {
        if let bundled = fromBundle() {
            return bundled
        }
        return fromProcessInfo()
    }
}

/// Result of running a benchmark, containing both display text and JSON report
struct BenchmarkResult {
    let displayText: String
    let jsonReport: String
}

enum {{PROJECT_NAME_PASCAL}}FFI {
    static func runCurrentBenchmark() -> BenchmarkResult {
        let params = BenchParams.resolved()
        return run(params: params)
    }

    static func run(params: BenchParams) -> BenchmarkResult {
        let spec = BenchSpec(
            name: params.function,
            iterations: params.iterations,
            warmup: params.warmup
        )

        do {
            let processMemorySampler = ProcessMemorySampler()
            processMemorySampler.start()
            let report: BenchReport
            do {
                report = try runBenchmark(spec: spec)
            } catch {
                _ = processMemorySampler.stop()
                throw error
            }
            let runProcessPeakMemoryKb = processMemorySampler.stop()
            let displayText = formatBenchReport(report)
            let jsonReport = generateJSONReport(report, runProcessPeakMemoryKb: runProcessPeakMemoryKb)
            return BenchmarkResult(displayText: displayText, jsonReport: jsonReport)
        } catch let error as BenchError {
            print("[BenchRunner] ERROR: Benchmark failed: \(error)")
            let errorText = formatBenchError(error)
            let errorJSON = generateErrorJSON(error)
            return BenchmarkResult(displayText: errorText, jsonReport: errorJSON)
        } catch {
            print("[BenchRunner] ERROR: Unexpected error during benchmark execution: \(error)")
            let errorText = "Unexpected error: \(error.localizedDescription)"
            let errorJSON = "{\"error\": \"Unexpected error: \(error.localizedDescription)\"}"
            return BenchmarkResult(displayText: errorText, jsonReport: errorJSON)
        }
    }

    /// Generates a JSON report that matches the Android BENCH_JSON-style payload
    /// consumed by mobench so cross-platform parsers can stay aligned.
    private static func generateJSONReport(_ report: BenchReport, runProcessPeakMemoryKb: UInt64?) -> String {
        var json: [String: Any] = [:]

        // Spec section
        let specDict: [String: Any] = [
            "name": report.spec.name,
            "iterations": report.spec.iterations,
            "warmup": report.spec.warmup
        ]
        json["spec"] = specDict

        // Function name at top level (for compatibility with existing parsers)
        json["function"] = report.spec.name

        // Samples as array of duration_ns values
        let samplesNs = report.samples.map { $0.durationNs }
        json["samples_ns"] = samplesNs

        // Also include samples in object format for compatibility
        let samplesArray = report.samples.map { sample -> [String: Any] in
            var sampleJson: [String: Any] = ["duration_ns": sample.durationNs]
            if let cpuTimeMs = optionalCpuTimeMs(sample) {
                sampleJson["cpu_time_ms"] = cpuTimeMs
            }
            if let peakMemoryKb = optionalPeakMemoryKb(sample) {
                sampleJson["peak_memory_kb"] = peakMemoryKb
            }
            if let processPeakMemoryKb = optionalProcessPeakMemoryKb(sample) {
                sampleJson["process_peak_memory_kb"] = processPeakMemoryKb
            }
            return sampleJson
        }
        json["samples"] = samplesArray

        let phases = optionalPhases(report).compactMap { phase -> [String: Any]? in
            guard let name = optionalStringProperty(phase, named: "name"),
                  let durationNs = optionalUInt64Property(phase, named: "durationNs") else {
                return nil
            }
            return [
                "name": name,
                "duration_ns": durationNs
            ] as [String: Any]
        }
        json["phases"] = phases

        // Statistics
        if !report.samples.isEmpty {
            let durations = report.samples.map { $0.durationNs }
            let minNs = durations.min() ?? 0
            let maxNs = durations.max() ?? 0
            let sumNs = durations.reduce(0, +)
            let avgNs = Double(sumNs) / Double(durations.count)

            // Compute median
            let sorted = durations.sorted()
            let medianNs: UInt64
            if sorted.count % 2 == 0 {
                medianNs = (sorted[sorted.count / 2 - 1] + sorted[sorted.count / 2]) / 2
            } else {
                medianNs = sorted[sorted.count / 2]
            }

            let stats: [String: Any] = [
                "min_ns": minNs,
                "max_ns": maxNs,
                "avg_ns": avgNs,
                "mean_ns": UInt64(avgNs),
                "median_ns": medianNs
            ]
            json["stats"] = stats

            // Also add top-level stats for compatibility with different parsers
            json["mean_ns"] = UInt64(avgNs)
            json["min_ns"] = minNs
            json["max_ns"] = maxNs
        }

        let cpuSamplesMs = report.samples.compactMap { optionalCpuTimeMs($0) }
        let peakSamplesKb = report.samples.compactMap { optionalPeakMemoryKb($0) }
        let processPeakSamplesKb = report.samples.compactMap { optionalProcessPeakMemoryKb($0) }

        func median(_ values: [UInt64]) -> UInt64 {
            let sorted = values.sorted()
            if sorted.count % 2 == 0 {
                return (sorted[sorted.count / 2 - 1] + sorted[sorted.count / 2]) / 2
            } else {
                return sorted[sorted.count / 2]
            }
        }

        // Resource metrics derived from the measured loop itself.
        var resources: [String: Any] = [
            "platform": "ios",
            "memory_process": "benchmark_app",
            "timestamp_ms": Int64(Date().timeIntervalSince1970 * 1000)
        ]
        if !cpuSamplesMs.isEmpty {
            let cpuTotalMs = cpuSamplesMs.reduce(0, +)
            resources["cpu_total_ms"] = cpuTotalMs
            resources["cpu_median_ms"] = median(cpuSamplesMs)
            resources["elapsed_cpu_ms"] = cpuTotalMs
        }
        if let peakMemoryKb = peakSamplesKb.max() {
            resources["peak_memory_kb"] = peakMemoryKb
            resources["peak_memory_growth_kb"] = peakMemoryKb
        }
        if let processPeakMemoryKb = processPeakSamplesKb.max() ?? runProcessPeakMemoryKb {
            resources["process_peak_memory_kb"] = processPeakMemoryKb
        }
        json["resources"] = resources

        // Serialize to JSON string
        do {
            let data = try JSONSerialization.data(withJSONObject: json, options: [.sortedKeys])
            return String(data: data, encoding: .utf8) ?? "{}"
        } catch {
            print("[BenchRunner] ERROR: Failed to serialize JSON report: \(error)")
            return "{}"
        }
    }

    private static func optionalProcessPeakMemoryKb(_ sample: Any) -> UInt64? {
        return optionalUInt64Property(sample, named: "processPeakMemoryKb")
    }

    private static func optionalCpuTimeMs(_ sample: Any) -> UInt64? {
        return optionalUInt64Property(sample, named: "cpuTimeMs")
    }

    private static func optionalPeakMemoryKb(_ sample: Any) -> UInt64? {
        return optionalUInt64Property(sample, named: "peakMemoryKb")
    }

    private static func optionalPhases(_ report: Any) -> [Any] {
        for child in Mirror(reflecting: report).children where child.label == "phases" {
            return coerceToArray(child.value)
        }
        return []
    }

    private static func coerceToArray(_ value: Any) -> [Any] {
        let mirror = Mirror(reflecting: value)
        if mirror.displayStyle == .optional {
            guard let child = mirror.children.first else {
                return []
            }
            return coerceToArray(child.value)
        }
        if let array = value as? [Any] {
            return array
        }
        if mirror.displayStyle == .collection {
            return mirror.children.map { $0.value }
        }
        return []
    }

    private static func optionalUInt64Property(_ instance: Any, named propertyName: String) -> UInt64? {
        for child in Mirror(reflecting: instance).children where child.label == propertyName {
            return coerceToUInt64(child.value)
        }
        return nil
    }

    private static func optionalStringProperty(_ instance: Any, named propertyName: String) -> String? {
        for child in Mirror(reflecting: instance).children where child.label == propertyName {
            return String(describing: child.value)
        }
        return nil
    }

    private static func coerceToUInt64(_ value: Any) -> UInt64? {
        let mirror = Mirror(reflecting: value)
        if mirror.displayStyle == .optional {
            guard let child = mirror.children.first else {
                return nil
            }
            return coerceToUInt64(child.value)
        }

        if let value = value as? UInt64 {
            return value
        }
        if let value = value as? UInt32 {
            return UInt64(value)
        }
        if let value = value as? UInt {
            return UInt64(value)
        }
        if let value = value as? Int64, value >= 0 {
            return UInt64(value)
        }
        if let value = value as? Int, value >= 0 {
            return UInt64(value)
        }
        return nil
    }

    /// Generates a JSON error report
    private static func generateErrorJSON(_ error: BenchError) -> String {
        let errorDict: [String: Any] = [
            "error": true,
            "message": error.localizedDescription
        ]

        do {
            let data = try JSONSerialization.data(withJSONObject: errorDict, options: [.sortedKeys])
            return String(data: data, encoding: .utf8) ?? "{\"error\": true}"
        } catch {
            return "{\"error\": true, \"message\": \"Failed to serialize error\"}"
        }
    }

    /// Formats a duration in nanoseconds to a human-readable string.
    /// Uses milliseconds (ms) by default, switches to seconds (s) if >= 1000ms.
    private static func formatDuration(_ ns: UInt64) -> String {
        let ms = Double(ns) / 1_000_000.0
        if ms >= 1000.0 {
            let secs = ms / 1000.0
            return String(format: "%.3fs", secs)
        } else {
            return String(format: "%.3fms", ms)
        }
    }

    private static func formatBenchReport(_ report: BenchReport) -> String {
        var output = "=== Benchmark Results ===\n\n"
        output += "Function: \(report.spec.name)\n"
        output += "Iterations: \(report.spec.iterations)\n"
        output += "Warmup: \(report.spec.warmup)\n\n"

        output += "Samples (\(report.samples.count)):\n"
        for (index, sample) in report.samples.enumerated() {
            output += "  \(index + 1). \(formatDuration(sample.durationNs))\n"
        }

        if !report.samples.isEmpty {
            let durations = report.samples.map { $0.durationNs }
            let min = durations.min() ?? 0
            let max = durations.max() ?? 0
            let avg = durations.reduce(0, +) / UInt64(durations.count)

            output += "\nStatistics:\n"
            output += "  Min: \(formatDuration(min))\n"
            output += "  Max: \(formatDuration(max))\n"
            output += "  Avg: \(formatDuration(avg))\n"
        }

        return output
    }

    private static func formatBenchError(_ error: BenchError) -> String {
        // Generic error formatting - works with any BenchError variant
        // This avoids hardcoding specific error cases that may not exist in all schemas
        return "Benchmark error: \(error.localizedDescription)"
    }
}

private final class ProcessMemorySampler {
    private let sampleInterval: TimeInterval
    private let queue = DispatchQueue(label: "mobench-process-memory-sampler")
    private var timer: DispatchSourceTimer?
    private var peakKb: UInt64 = 0

    init(sampleInterval: TimeInterval = 0.01) {
        self.sampleInterval = sampleInterval
    }

    func start() {
        queue.sync {
            guard timer == nil else {
                return
            }

            recordCurrentLocked()
            let source = DispatchSource.makeTimerSource(queue: queue)
            source.schedule(deadline: .now(), repeating: sampleInterval)
            source.setEventHandler { [weak self] in
                self?.recordCurrentLocked()
            }
            source.resume()
            timer = source
        }
    }

    func stop() -> UInt64? {
        queue.sync {
            timer?.cancel()
            timer = nil
            recordCurrentLocked()
            return peakKb > 0 ? peakKb : nil
        }
    }

    private func recordCurrentLocked() {
        if let currentKb = currentProcessResidentMemoryKb(), currentKb > peakKb {
            peakKb = currentKb
        }
    }
}

private func currentProcessResidentMemoryKb() -> UInt64? {
    var info = mach_task_basic_info()
    var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size) / 4
    let result = withUnsafeMutablePointer(to: &info) { infoPointer in
        infoPointer.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { reboundPointer in
            task_info(
                mach_task_self_,
                task_flavor_t(MACH_TASK_BASIC_INFO),
                reboundPointer,
                &count
            )
        }
    }

    guard result == KERN_SUCCESS else {
        return nil
    }

    return UInt64(info.resident_size) / 1024
}
