import Foundation

struct CompileResult {
    let success: Bool
    let output: String
    let error: String
}

func runCommand(_ command: String, timeout: TimeInterval = 30) -> CompileResult {
    let task = Process()
    task.executableURL = URL(fileURLWithPath: "/bin/zsh")
    task.arguments = ["-c", command]
    
    let outputPipe = Pipe()
    let errorPipe = Pipe()
    task.standardOutput = outputPipe
    task.standardError = errorPipe
    
    do {
        try task.run()
        task.waitUntilExit()
        
        let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile()
        let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile()
        
        let output = String(data: outputData, encoding: .utf8) ?? ""
        let error = String(data: errorData, encoding: .utf8) ?? ""
        
        return CompileResult(success: task.terminationStatus == 0, output: output.trimmingCharacters(in: .whitespacesAndNewlines), error: error.trimmingCharacters(in: .whitespacesAndNewlines))
    } catch {
        return CompileResult(success: false, output: "", error: error.localizedDescription)
    }
}

func compileAndRunSwift(code: String, testInput: String = "") -> CompileResult {
    let tempFile = "/tmp/test_swift_code.swift"
    do {
        try code.write(toFile: tempFile, atomically: true, encoding: .utf8)
    } catch {
        return CompileResult(success: false, output: "", error: "Failed to write file: \(error)")
    }
    
    let compileResult = runCommand("swiftc \(tempFile) -o /tmp/test_swift_exec", timeout: 60)
    if !compileResult.success {
        return compileResult
    }
    
    return runCommand("/tmp/test_swift_exec \(testInput)", timeout: 10)
}

func compileAndRunPython(code: String, testInput: String = "") -> CompileResult {
    let tempFile = "/tmp/test_python_code.py"
    do {
        try code.write(toFile: tempFile, atomically: true, encoding: .utf8)
    } catch {
        return CompileResult(success: false, output: "", error: "Failed to write file: \(error)")
    }
    
    return runCommand("python3 \(tempFile) \(testInput)", timeout: 30)
}

func compileAndRunCpp(code: String, testInput: String = "") -> CompileResult {
    let tempFile = "/tmp/test_cpp_code.cpp"
    do {
        try code.write(toFile: tempFile, atomically: true, encoding: .utf8)
    } catch {
        return CompileResult(success: false, output: "", error: "Failed to write file: \(error)")
    }
    
    let compileResult = runCommand("clang++ \(tempFile) -o /tmp/test_cpp_exec", timeout: 60)
    if !compileResult.success {
        return compileResult
    }
    
    return runCommand("/tmp/test_cpp_exec \(testInput)", timeout: 10)
}

func compileAndRunJavaScript(code: String, testInput: String = "") -> CompileResult {
    let tempFile = "/tmp/test_js_code.js"
    do {
        try code.write(toFile: tempFile, atomically: true, encoding: .utf8)
    } catch {
        return CompileResult(success: false, output: "", error: "Failed to write file: \(error)")
    }
    
    return runCommand("node \(tempFile) \(testInput)", timeout: 30)
}

func compileAndRunRust(code: String, testInput: String = "") -> CompileResult {
    let tempFile = "/tmp/test_rust_code.rs"
    do {
        try code.write(toFile: tempFile, atomically: true, encoding: .utf8)
    } catch {
        return CompileResult(success: false, output: "", error: "Failed to write file: \(error)")
    }
    
    let compileResult = runCommand("source $HOME/.cargo/env && rustc \(tempFile) -o /tmp/test_rust_exec", timeout: 60)
    if !compileResult.success {
        return compileResult
    }
    
    return runCommand("source $HOME/.cargo/env && /tmp/test_rust_exec \(testInput)", timeout: 10)
}

func verifyOutput(actual: String, expected: String) -> Bool {
    let actualTrimmed = actual.trimmingCharacters(in: .whitespacesAndNewlines)
    let expectedTrimmed = expected.trimmingCharacters(in: .whitespacesAndNewlines)
    return actualTrimmed == expectedTrimmed || actualTrimmed.contains(expectedTrimmed)
}

func sampleTopK(logits: [Float], k: Int = 40, temperature: Float = 0.9) -> Int {
    let scaledLogits = logits.map { $0 / temperature }
    
    let sorted = scaledLogits.enumerated().sorted { $0.element > $1.element }
    let topK = sorted.prefix(k)
    
    let maxLogit = topK.first?.element ?? 0
    let probs = topK.map { Float(exp(Double($0.element - maxLogit))) }
    let sum = probs.reduce(0, +)
    let normalizedProbs = probs.map { $0 / sum }
    
    let random = Float.random(in: 0..<1)
    var cumulative: Float = 0.0
    for (index, prob) in normalizedProbs.enumerated() {
        cumulative += prob
        if random <= cumulative {
            return topK[index].offset
        }
    }
    
    return topK.last?.offset ?? 0
}