import XCTest
@testable import MarkBase

class CodeGenerationTest: XCTestCase {
    
    func testSwiftCodeGeneration() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Swift Code Generation Test (Improved Sampling)")
        print("═════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        
        let prompt = swiftPrompts[0]
        print("Prompt: \(prompt.prompt)")
        
        let tokens = tokenizer.encode(text: prompt.prompt)
        print("Prompt tokens: \(tokens.count)")
        
        var allTokens = tokens
        var generatedTokens: [Int] = []
        
        for i in 0..<80 {
            let pos = i < tokens.count ? i : tokens.count + (i - tokens.count)
            let currentToken = allTokens[i]
            
            let logits = try model.forwardOptimized(tokenId: currentToken, position: pos)
            let nextToken = sampleTopK(logits: logits, k: 50, temperature: 0.8)
            
            allTokens.append(nextToken)
            generatedTokens.append(nextToken)
            
            print("Token \(i): pos=\(pos), token=\(nextToken)")
            
            if tokenizer.eosTokenIds.contains(nextToken) || i >= tokens.count + 60 {
                break
            }
        }
        
        let generatedCode = tokenizer.decode(tokens: generatedTokens)
        print("\nGenerated Code:\n\(generatedCode)")
        
        let result = compileAndRunSwift(code: generatedCode, testInput: prompt.testInput)
        
        if result.success {
            print("✅ PASS - Code compiled and ran successfully")
            print("Output: \(result.output)")
        } else {
            print("⚠️ FAIL - Compilation/runtime error")
            print("Error: \(result.error)")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
    
    func testPythonCodeGeneration() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Python Code Generation Test (Improved Sampling)")
        print("═════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        
        let prompt = pythonPrompts[0]
        print("Prompt: \(prompt.prompt)")
        
        let tokens = tokenizer.encode(text: prompt.prompt)
        print("Prompt tokens: \(tokens.count)")
        
        var allTokens = tokens
        var generatedTokens: [Int] = []
        
        for i in 0..<80 {
            let pos = i < tokens.count ? i : tokens.count + (i - tokens.count)
            let currentToken = allTokens[i]
            
            let logits = try model.forwardOptimized(tokenId: currentToken, position: pos)
            let nextToken = sampleTopK(logits: logits, k: 50, temperature: 0.8)
            
            allTokens.append(nextToken)
            generatedTokens.append(nextToken)
            
            if tokenizer.eosTokenIds.contains(nextToken) || i >= tokens.count + 60 {
                break
            }
        }
        
        let generatedCode = tokenizer.decode(tokens: generatedTokens)
        print("\nGenerated Code:\n\(generatedCode)")
        
        let result = compileAndRunPython(code: generatedCode, testInput: prompt.testInput)
        
        if result.success {
            print("✅ PASS - Code ran successfully")
            print("Output: \(result.output)")
        } else {
            print("⚠️ FAIL - Runtime error")
            print("Error: \(result.error)")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
}