import XCTest
@testable import MarkBase

class CodeGenerationTest: XCTestCase {
    
    func testSwiftCodeGeneration() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Swift Code Generation Test")
        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)
        var generatedTokens: [Int] = []
        
        for i in 0..<50 {
            let pos = tokens.count + i - 1
            let logits = try model.forwardOptimized(tokenId: tokens.last ?? 0, position: pos)
            let nextToken = logits.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
            generatedTokens.append(nextToken)
            
            if tokenizer.eosTokenIds.contains(nextToken) { break }
        }
        
        let generatedCode = tokenizer.decode(tokens: generatedTokens)
        print("Generated: \(generatedCode)")
        
        let result = compileAndRunSwift(code: generatedCode, testInput: prompt.testInput)
        
        if result.success {
            print("✅ PASS")
        } else {
            print("⚠ FAIL: \(result.error)")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
    
    func testPythonCodeGeneration() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Python Code Generation Test")
        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)
        var generatedTokens: [Int] = []
        
        for i in 0..<50 {
            let pos = tokens.count + i - 1
            let logits = try model.forwardOptimized(tokenId: tokens.last ?? 0, position: pos)
            let nextToken = logits.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0
            generatedTokens.append(nextToken)
            
            if tokenizer.eosTokenIds.contains(nextToken) { break }
        }
        
        let generatedCode = tokenizer.decode(tokens: generatedTokens)
        print("Generated: \(generatedCode)")
        
        let result = compileAndRunPython(code: generatedCode, testInput: prompt.testInput)
        
        if result.success {
            print("✅ PASS")
        } else {
            print("⚠ FAIL: \(result.error)")
        }
        
        model.kvCaches.forEach { cache in cache.reset() }
    }
}